Add du for checking disk usage.

ywangd · ywangd · commit ccaa7497c228 · 2016-01-22T00:17:34.000+11:00
Refactor sizeof_fmt function and move it to libcore.
diff --git a/bin/du.py b/bin/du.py
@@ -0,0 +1,78 @@
+import os
+import sys
+from argparse import ArgumentParser
+from fnmatch import fnmatch
+
+
+def is_excluded(path, pattern):
+    if pattern:
+        while path != '':
+            prefix, tail = os.path.split(path)
+            if fnmatch(tail, pattern):
+                return True
+            else:
+                path = prefix
+        return False
+    else:
+        return False
+
+
+def main(args):
+    ap = ArgumentParser(
+        description='Summarize disk usage of the set of FILEs, recursively for directories.'
+    )
+    ap.add_argument('-s', '--summarize', action='store_true',
+                    help='display only a total for each argument')
+    ap.add_argument('--exclude', dest='exclude_pattern',
+                    metavar='PATTERN',
+                    help='exclude files that match PATTERN')
+    ap.add_argument('FILEs', nargs='*', default=['.'],
+                    help='files to summarize (default to current working directory')
+
+    ns = ap.parse_args(args)
+
+    exclude_pattern = ns.exclude_pattern if ns.exclude_pattern else None
+
+    sizeof_fmt = globals()['_stash'].libcore.sizeof_fmt
+
+    for path in ns.FILEs:
+
+        path_base = os.path.dirname(path)
+
+        # Use relative path because of the following facts:
+        # du A/B --exclude="B"  -> no output
+        # du A/B --exclude="A"  -> normal output
+        if is_excluded(os.path.relpath(path, path_base), exclude_pattern):
+            continue
+
+        if os.path.isdir(path):
+            dirs_dict = {}
+            # We need to walk the tree from the bottom up so that a directory can have easy
+            # access to the size of its subdirectories.
+            for root, dirs, files in os.walk(path, topdown=False):
+                # This is to make sure the directory is not exclude from its ancestor
+                if is_excluded(os.path.relpath(root, path_base), exclude_pattern):
+                    continue
+
+                # Loop through every non directory file in this directory and sum their sizes
+                size = sum(os.path.getsize(os.path.join(root, name))
+                           for name in files if not is_excluded(name, exclude_pattern))
+
+                # Look at all of the subdirectories and add up their sizes from the `dirs_dict`
+                subdir_size = sum(dirs_dict[os.path.join(root, d)]
+                                  for d in dirs if not is_excluded(d, exclude_pattern))
+
+                # store the size of this directory (plus subdirectories) in a dict so we
+                # can access it later
+                my_size = dirs_dict[root] = size + subdir_size
+
+                if ns.summarize and root != path:
+                    continue
+
+                print '%-8s %s' % (sizeof_fmt(my_size), root)
+        else:
+            print '%-8s %s' % (sizeof_fmt(os.path.getsize(path)), path)
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])
diff --git a/bin/ls.py b/bin/ls.py
@@ -8,36 +8,27 @@
 from argparse import ArgumentParser
 
 
-def sizeof_fmt(num):
-    for x in ['B', 'KB', 'MB', 'GB']:
-        if num < 1024.0:
-            if (x == 'bytes'):
-                return "%s %s" % (num, x)
-            else:
-                return "%3.1f %s" % (num, x)
-        num /= 1024.0
-    return "%3.1f%s" % (num, 'TB')
-
-
 def main(args):
 
     ap = ArgumentParser()
     ap.add_argument('-1', '--one-line', action='store_true', help='List one file per line')
     ap.add_argument('-a', '--all', action='store_true', help='do not ignore entries starting with .')
     ap.add_argument('-l', '--long', action='store_true', help='use a long listing format')
     ap.add_argument('files', nargs='*', help='files to be listed')
-    args = ap.parse_args(args)
+    ns = ap.parse_args(args)
+
+    sizeof_fmt = globals()['_stash'].libcore.sizeof_fmt
 
-    joiner = '\n' if args.one_line or args.long else ' '
+    joiner = '\n' if ns.one_line or ns.long else ' '
 
-    if args.all:
+    if ns.all:
         def _filter(filename):
             return True
     else:
         def _filter(filename):
             return False if filename.startswith('.') else True
 
-    if args.long:
+    if ns.long:
         def _fmt(filename, dirname=''):
             _stat = os.stat(os.path.join(dirname, filename))
 
@@ -71,15 +62,15 @@ def _fmt(filename, dirname=''):
             else:
                 return filename
 
-    if len(args.files) == 0:
+    if len(ns.files) == 0:
         out = joiner.join(_fmt(f) for f in os.listdir('.') if _filter(f))
         print out
 
     else:
         out_dir = []
         out_file = []
         out_miss = []
-        for f in args.files:
+        for f in ns.files:
             if not os.path.exists(f):
                 out_miss.append('ls: %s: No such file or directory' % f)
             elif os.path.isdir(f):
diff --git a/lib/libcore.py b/lib/libcore.py
@@ -59,3 +59,10 @@ def input_stream(files=()):
     finally:
         fileinput.close()
 
+
+def sizeof_fmt(num):
+    for unit in ['B', 'K', 'M', 'G']:
+        if num < 1024:
+            return "%3.1f%s" % (num, unit)
+        num /= 1024.0
+    return "%3.1f%s" % (num, 'T')
