pysize/core/pysize_fs_node.py

   1 import os
   2 import stat
   3
   4 from pysize.core import chdir_browsing
   5 from pysize.core import compute_size
   6
   7 def _extract_prefix_suffixes(paths):
   8     """['/prefix/suffix1', '/prefix/suffix2', '/prefix/suffix3'] =>
   9         ('/prefix', ['suffix1', 'suffix2', 'suffix3'])"""
  10     prefix_len = os.path.commonprefix(paths).rfind('/') + 1
  11     prefix = paths[0][:prefix_len - 1]
  12     suffixes = [p[prefix_len:] for p in paths]
  13     return (prefix, suffixes)
  14
  15 def _join_prefix_suffixes(prefix, suffixes):
  16     if len(suffixes) == 1:
  17         return prefix + '/' + suffixes[0]
  18     return prefix + '/{' + ','.join(suffixes) + '}'
  19
  20 def _sort_nodes(nodes):
  21     def cmp_fn(n1, n2):
  22         size_diff = n2.size - n1.size
  23         if size_diff:
  24             return size_diff
  25         return cmp(n1.get_name(), n2.get_name())
  26     nodes.sort(cmp=cmp_fn)
  27
  28 class _pysize_node(object):
  29     """The parent class of all displayed nodes, as these nodes are displayed on
  30     screen, there should be few instances."""
  31     def __init__(self, parent, basename):
  32         self.rectangle = (0, 0, 0, 0)
  33         self.parent = parent
  34         self.children = []
  35         self.basename = basename
  36         if basename:
  37             self.size = compute_size.slow(basename)
  38         else:
  39             self.size = 0
  40
  41     def compute_height(self):
  42         if self.children:
  43             children_height = max([c.compute_height() for c in self.children])
  44             return children_height + 1
  45         return 0
  46
  47     def compute_depth(self):
  48         depth = 0
  49         node = self
  50         while node:
  51             node = node.parent
  52             depth += 1
  53         return depth
  54
  55     def minimum_real_size(self):
  56         res = self.size
  57         for child in self.children:
  58             child_min_real_size = child.minimum_real_size()
  59             if child_min_real_size >= 0 and child_min_real_size < res:
  60                 res = child_min_real_size
  61         return res
  62
  63     def get_dirname(self):
  64         return os.path.dirname(self.get_fullpaths()[0])
  65
  66     def is_dir(self):
  67         return False
  68
  69     def is_real(self):
  70         """Does the node correspond to a path that actually exists?"""
  71         return True
  72
  73     def get_fullname(self):
  74         fullpaths = self.get_fullpaths()
  75         (prefix, suffixes) = _extract_prefix_suffixes(fullpaths)
  76         fullname = _join_prefix_suffixes(prefix, suffixes)
  77         return fullname
  78
  79     def get_fullpaths(self):
  80         parent = self.parent
  81         fullpath = self.basename
  82         while parent:
  83             fullpath = parent.basename + '/' + fullpath
  84             parent = parent.parent
  85         return [fullpath]
  86
  87     def get_name(self):
  88         """Return the name that should be displayed for this node."""
  89         return self.basename
  90
  91     def __iter__(self):
  92         """The iterator is a depth first traversal."""
  93         yield self
  94         for child in self.children:
  95             for node in child:
  96                 yield node
  97
  98     def contains_point(self, p):
  99         """Is the point p in the graphical representation of this node?"""
 100         (x0, x1, y0, y1) = self.rectangle
 101         return x0 < p.x and p.x < x1 and y0 < p.y and p.y < y1
 102
 103 class _pysize_node_collection(_pysize_node):
 104     """A directory"""
 105     def __init__(self, parent, prefix, children, max_depth, min_size):
 106         super(_pysize_node_collection, self).__init__(parent, None)
 107         self.basename = prefix
 108         self.size = sum([compute_size.slow(prefix + '/' + p) for p in children])
 109         if max_depth != 0:
 110             children_size = 0
 111             remaining_size = 0
 112             remaining_items = []
 113             cookie = chdir_browsing.init(prefix)
 114             try:
 115                 for child in children:
 116                     node = create_node(self, child, max_depth - 1, min_size)
 117                     if node.is_real():
 118                         self.children.append(node)
 119                         children_size += node.size
 120                     else:
 121                         node.__name = child
 122                         remaining_items.append(node)
 123                         remaining_size += node.size
 124
 125                 _sort_nodes(self.children)
 126                 if remaining_size > min_size:
 127                     _sort_nodes(remaining_items)
 128                     names = [n.__name for n in remaining_items]
 129                     rem = _pysize_node_remaining(self, names)
 130                     self.children.append(rem)
 131             finally:
 132                 chdir_browsing.finalize(cookie)
 133             self.size = max(children_size + remaining_size, self.size)
 134
 135     def is_real(self):
 136         """This node does not actually exists, it is an aggregate."""
 137         return False
 138
 139 class _pysize_node_forest(_pysize_node_collection):
 140     def __init__(self, parent, children, max_depth, min_size):
 141         (prefix, suffixes) = _extract_prefix_suffixes(children)
 142         super(_pysize_node_forest, self).__init__(parent, prefix, suffixes,
 143                                                   max_depth, min_size)
 144         self.basename = prefix
 145         self.forest_items = suffixes
 146         self.forest_name = _join_prefix_suffixes(prefix, suffixes)
 147
 148     def get_name(self):
 149         return self.forest_name
 150
 151     def get_dirname(self):
 152         return self.basename
 153
 154     def get_fullpaths(self):
 155         fullpaths = self.forest_items
 156         parent = self
 157         while parent:
 158             fullpaths = [parent.basename + '/' + fp for fp in fullpaths]
 159             parent = parent.parent
 160         return fullpaths
 161
 162 class _pysize_node_dir(_pysize_node_collection):
 163     """A directory"""
 164     def __init__(self, parent, basename, max_depth, min_size):
 165         super(_pysize_node_dir, self).__init__(parent, basename,
 166                                                os.listdir(basename), max_depth,
 167                                                min_size)
 168         self.basename = basename
 169         self.size = max(self.size, compute_size.slow(basename))
 170
 171     def is_dir(self):
 172         return True
 173
 174     def is_real(self):
 175         return True
 176
 177     def get_name(self):
 178         return self.basename + '/'
 179
 180 class _pysize_node_remaining(_pysize_node_collection):
 181     """The sum of a directory's children that are too small to be drawn."""
 182     def __init__(self, parent, elements):
 183         _pysize_node.__init__(self, parent, None)
 184         # The parent constructor would visit the files
 185         self.size = sum(map(compute_size.slow, elements))
 186         self.remaining_elements = elements
 187
 188     def minimum_real_size(self):
 189         return -1
 190
 191     def get_name(self):
 192         return '{' + ','.join(self.remaining_elements) + '}'
 193
 194     def get_fullname(self):
 195         if self.remaining_elements:
 196             return _pysize_node_collection.get_fullname(self)
 197         return '' # This is the initial node
 198
 199     def get_fullpaths(self):
 200         fullpaths = self.remaining_elements
 201         parent = self.parent
 202         while parent:
 203             fullpaths = [parent.basename + '/' + fp for fp in fullpaths]
 204             parent = parent.parent
 205         return fullpaths
 206
 207 class _pysize_node_file(_pysize_node):
 208     """A file"""
 209     def __init__(self, parent, basename):
 210         super(_pysize_node_file, self).__init__(parent, basename)
 211
 212 class _pysize_node_hardlink(_pysize_node_file):
 213     """A hardlink, the canonical one, or a link"""
 214     def __init__(self, parent, basename):
 215         super(_pysize_node_hardlink, self).__init__(parent, basename)
 216
 217 def create_node(parent, what, max_depth, min_size):
 218     """Return a pysize_node for parent/basename traversing up to max_depth
 219     levels and only taking into account elements bigger than min_size."""
 220     if not what:
 221         return _pysize_node_remaining(parent, [])
 222
 223     if isinstance(what, list):
 224         size = sum(map(compute_size.slow, what))
 225         if len(what) == 1:
 226             what = what[0]
 227     else:
 228         size = compute_size.slow(what)
 229
 230     if size < min_size:
 231         if isinstance(what, str):
 232             what = [what]
 233         node = _pysize_node_remaining(parent, what)
 234     elif isinstance(what, list):
 235         node = _pysize_node_forest(parent, what, max_depth, min_size)
 236     else:
 237         st = os.lstat(what)
 238         if stat.S_ISDIR(st.st_mode):
 239             node = _pysize_node_dir(parent, what, max_depth, min_size)
 240         elif st.st_nlink > 1:
 241             node = _pysize_node_hardlink(parent, what)
 242         else:
 243             node = _pysize_node_file(parent, what)
 244     return node