core/pysize_fs_tree.py

   1 from core import compute_size
   2 from core.pysize_fs_node import create_node
   3
   4 # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/231503
   5 # By David Eppstein
   6 def _breadth_first(tree, children=iter):
   7     """Traverse the nodes of a tree in breadth-first order.
   8     The first argument should be the tree root; children
   9     should be a function taking as argument a tree node and
  10     returning an iterator of the node's children.
  11     """
  12     yield tree
  13     last = tree
  14     for node in _breadth_first(tree, children):
  15         for child in children(node):
  16             yield child
  17             last = child
  18         if last is node:
  19             return
  20
  21 class pysize_tree:
  22     """The entry point to a tree of pysize_node.
  23     min_size < 1.0 => ratio to total size."""
  24     def __init__(self, path, max_depth=1, min_size=0.0):
  25         self.fullpath = path
  26         if path and min_size < 1.0:
  27             min_size *= compute_size.slow(path)
  28         self.root = create_node(None, path, max_depth, min_size)
  29         self.height = self.root.compute_height()
  30
  31     def get_next_sibling(self, node):
  32         """Return the next pysize_node in node's level."""
  33         is_next = False
  34         for child in _breadth_first(self.root, lambda c: c.get_children()):
  35             if is_next:
  36                 if child.compute_depth() == node.compute_depth():
  37                     return child
  38                 return
  39             is_next = child is node
  40
  41     def get_previous_sibling(self, node):
  42         """Return the previous pysize_node in node's level."""
  43         prev = None
  44         for child in _breadth_first(self.root, lambda c: c.get_children()):
  45             if child == node:
  46                 if prev.compute_depth() == node.compute_depth():
  47                     return prev
  48                 return
  49             prev = child
  50
  51     def get_first_child(self, node):
  52         """Return the first pysize_node in node's children."""
  53         children = node.get_children()
  54         if children:
  55             return children[0]
  56
  57     def get_parent(self, node):
  58         """Return the saved parent of node."""
  59         return node.parent
  60
  61 def _main():
  62     tree = pysize_tree('/usr/share/doc', 4, 0.1)
  63     print 'height', tree.height
  64
  65 if __name__ == '__main__':
  66     _main()