sfa/util/prefixTree.py

   1 from __future__ import print_function
   2
   3 class prefixNode:
   4
   5     def __init__(self, prefix):
   6         self.prefix = prefix
   7         self.children = []
   8
   9
  10 class prefixTree:
  11
  12     def __init__(self):
  13         self.root = prefixNode("")
  14
  15     def insert(self, prefix, node = None):
  16         """
  17         insert a prefix into the tree
  18         """
  19         if not node:
  20             node = self.root
  21
  22         parts = prefix.split(".")
  23         length = len(parts)
  24
  25         if length > 1:
  26             for i in range(1, length + 1):
  27                 name = ".".join(parts[:i])
  28                 if not self.exists(name) and not name == prefix:
  29                     self.insert(name)
  30
  31         if prefix.startswith(node.prefix):
  32             if prefix == node.prefix:
  33                 pass
  34             elif not node.children:
  35                 node.children.append(prefixNode(prefix))
  36             else:
  37                 inserted  = False
  38                 for child in node.children:
  39                     if prefix.startswith(child.prefix):
  40                         self.insert(prefix, child)
  41                         inserted = True
  42                 if not inserted:
  43                     node.children.append(prefixNode(prefix))
  44
  45     def load(self, prefix_list):
  46         """
  47         load a list of prefixes into the tree
  48         """
  49         for prefix in prefix_list:
  50             self.insert(prefix)
  51
  52     def exists(self, prefix, node = None):
  53         """
  54         returns true if the specified prefix exists anywhere in the tree,
  55         false if it doesnt.
  56         """
  57         if not node:
  58             node = self.root
  59
  60         if not prefix.startswith(node.prefix):
  61             return False
  62         elif node.prefix == prefix:
  63             return True
  64         elif not node.children:
  65             return False
  66         else:
  67             for child in node.children:
  68                 if prefix.startswith(child.prefix):
  69                     return self.exists(prefix, child)
  70
  71     def best_match(self, prefix, node = None):
  72         """
  73         searches the tree and returns the prefix that best matches the
  74         specified prefix
  75         """
  76         if not node:
  77             node = self.root
  78
  79         if prefix.startswith(node.prefix):
  80             if not node.children:
  81                 return node.prefix
  82             for child in node.children:
  83                 if prefix.startswith(child.prefix):
  84                     return self.best_match(prefix, child)
  85             return node.prefix
  86
  87     def dump(self, node = None):
  88         """
  89         print the tree
  90         """
  91         if not node:
  92             node = self.root
  93             print(node.prefix)
  94
  95         for child in node.children:
  96             print(child.prefix, end=' ')
  97
  98         for child in node.children:
  99             self.dump(child)