core/pysize_fs_tree.py

   1 import os, stat
   2
   3 from pysize_global_fs_cache import get_dev_ino, cache_add_dir, \
   4 cache_add_hardlink, cache_get_dir_size
   5
   6 def _size_fast(path, dir_ino):
   7     try:
   8         st = os.lstat(path)
   9         if stat.S_ISDIR(st.st_mode):
  10             dev_ino = (st.st_dev, st.st_ino)
  11             cached_size = cache_get_dir_size(dev_ino)
  12             if cached_size == None:
  13                 prev = os.open('.', os.O_RDONLY)
  14                 os.chdir(path)
  15                 size = st.st_blocks * 512
  16                 dir_ino = st.st_ino
  17                 for p in os.listdir('.'):
  18                     size += _size_fast(p, dir_ino)
  19                 os.fchdir(prev)
  20                 os.close(prev)
  21                 cache_add_dir(dev_ino, size)
  22             else:
  23                 size = cached_size
  24         elif st.st_nlink > 1:
  25             if cache_add_hardlink(st.st_dev, st.st_ino, dir_ino, path):
  26                 size = st.st_blocks * 512
  27             else:
  28                 size = 0
  29         else:
  30             size = st.st_blocks * 512
  31         return size
  32     except OSError, e:
  33         print 'size', e
  34         return 0
  35
  36 def _size(path):
  37     def parent_path():
  38         elements = path.split('/')
  39         length = len(elements)
  40         found = False
  41         for i in xrange(length - 1, -1, -1):
  42             if elements[i] not in ('', '.'):
  43                 if found:
  44                     del elements[i+1:]
  45                     return '/'.join(elements)
  46                 found = True
  47         if path[0] == '/':
  48             return '/'
  49         if os.path.samefile(path, '.'):
  50             return os.path.dirname(os.getcwd())
  51         return '.'
  52     parent = parent_path()
  53     return _size_fast(path, get_dev_ino(parent)[1])
  54
  55 class _pysize_node:
  56     def __init__(self, parent, basename):
  57         self.rectangle = (0, 0, 0, 0)
  58         self.parent = parent
  59         if basename:
  60             self.basename = basename
  61             self.size = _size(basename)
  62
  63     def _compute_height(self):
  64         return 0
  65
  66     def _compute_depth(self):
  67         depth = 0
  68         node = self
  69         while node:
  70             node = node.parent
  71             depth += 1
  72         return depth
  73
  74     def is_dir(self):
  75         return False
  76
  77     def is_real(self):
  78         return True
  79
  80     def get_children(self):
  81         return []
  82
  83     def get_fullpath(self):
  84         fullpath = self.basename
  85         parent = self.parent
  86         while parent:
  87             fullpath = os.path.join(parent.basename, fullpath)
  88             parent = parent.parent
  89         return fullpath
  90
  91     def __iter__(self):
  92         yield self
  93         for child in self.get_children():
  94             for node in child:
  95                 yield node
  96
  97     def contains_point(self, p):
  98         (x0, x1, y0, y1) = self.rectangle
  99         return x0 < p.x and p.x < x1 and y0 < p.y and p.y < y1
 100
 101 class _pysize_node_remaining(_pysize_node):
 102     def __init__(self, parent, size):
 103         _pysize_node.__init__(self, parent, None)
 104         self.basename = '..///..'
 105         self.size = size
 106
 107     def is_real(self):
 108         return False
 109
 110 class _pysize_node_dir(_pysize_node):
 111     def __init__(self, parent, basename, max_depth, min_size):
 112         _pysize_node.__init__(self, parent, basename)
 113         self.children = []
 114         if max_depth != 0:
 115             children_size = 0
 116             prev = os.open('.', os.O_RDONLY)
 117             os.chdir(basename)
 118             for child in os.listdir('.'):
 119                 node = _create_node(self, child, max_depth - 1, min_size)
 120                 if node:
 121                     self.children.append(node)
 122                     children_size += node.size
 123             os.fchdir(prev)
 124             os.close(prev)
 125             self.children.sort(key=lambda t: t.size, reverse=True)
 126             self.size = max(children_size, self.size)
 127             remaining_size = self.size - children_size
 128             if remaining_size > min_size:
 129                 rem = _pysize_node_remaining(self, remaining_size)
 130                 self.children.append(rem)
 131
 132     def _compute_height(self):
 133         if self.children:
 134             children_height = max([c._compute_height() for c in self.children])
 135             return children_height + 1
 136         return 0
 137
 138     def is_dir(self):
 139         return True
 140
 141     def get_children(self):
 142         return self.children
 143
 144 class _pysize_node_file(_pysize_node):
 145     def __init__(self, parent, basename):
 146         _pysize_node.__init__(self, parent, basename)
 147
 148 class _pysize_node_hardlink(_pysize_node_file):
 149     def __init__(self, parent, basename):
 150         _pysize_node_file.__init__(self, parent, basename)
 151
 152 def _create_node(parent, basename, max_depth, min_size):
 153     if _size(basename) < min_size:
 154         return
 155     st = os.lstat(basename)
 156     if stat.S_ISDIR(st.st_mode):
 157         node = _pysize_node_dir(parent, basename, max_depth, min_size)
 158     elif st.st_nlink > 1:
 159         node = _pysize_node_hardlink(parent, basename)
 160     else:
 161         node = _pysize_node_file(parent, basename)
 162     return node
 163
 164 # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/231503
 165 # By David Eppstein
 166 def _breadth_first(tree,children=iter):
 167     """Traverse the nodes of a tree in breadth-first order.
 168     The first argument should be the tree root; children
 169     should be a function taking as argument a tree node and
 170     returning an iterator of the node's children.
 171     """
 172     yield tree
 173     last = tree
 174     for node in _breadth_first(tree,children):
 175         for child in children(node):
 176             yield child
 177             last = child
 178         if last == node:
 179             return
 180
 181 class pysize_tree:
 182     def __init__(self, path, max_depth, min_size_ratio):
 183         self.fullpath = path
 184         min_size = int(min_size_ratio * _size(path))
 185         self.root = _create_node(None, path, max_depth, min_size)
 186         if self.root:
 187             self.height = self.root._compute_height()
 188         else:
 189             self.height = 0
 190
 191     def get_next_sibling(self, node):
 192         is_next = False
 193         for child in _breadth_first(self.root, lambda c: c.get_children()):
 194             if is_next:
 195                 if child._compute_depth() == node._compute_depth():
 196                     return child
 197                 return None
 198             is_next = child is node
 199
 200     def get_previous_sibling(self, node):
 201         prev = None
 202         for child in _breadth_first(self.root, lambda c: c.get_children()):
 203             if child is node:
 204                 if prev._compute_depth() == node._compute_depth():
 205                     return prev
 206                 return None
 207             prev = child
 208
 209     def get_first_child(self, node):
 210         children = node.get_children()
 211         if children:
 212             return children[0]
 213
 214     def get_parent(self, node):
 215         return node.parent
 216
 217 def _main():
 218     tree = pysize_tree('/usr/share/doc', 4, 0.1)
 219     print 'height', tree.height
 220
 221 if __name__ == '__main__':
 222     _main()