sfa/util/prefixTree.py

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