Directory stats in Python

dir_stat.py

import os
import sys
import datetime


class ModCounter(object):
    def __init__(self, n, start=None, on_zero_handler=None):
        self._n = n
        self._counter = start if start else 0
        self._on_zero_handler = on_zero_handler if on_zero_handler else lambda: None

    def inc(self):
        self._counter += 1
        if self._counter == self._n:
            self._counter = 0
            self._on_zero_handler()
        return self._counter

    def reset(self):
        self._counter = 0

    @property
    def current_value(self):
        return self._counter

    @property
    def is_zero(self):
        return self.current_value == 0


def d_stat(dir_path, activity_log_count=None):
    """
    print dir stats
    """
    if activity_log_count is None:
        activity_log_count = 10000

    assert isinstance(activity_log_count, (int, long))
    assert activity_log_count >= 0

    use_log = activity_log_count > 0
    if use_log:
        def log_something():
            current_time = str(datetime.datetime.now())[11:][:8]
            print current_time, 'did %d actions, and still going' % activity_log_count
        mc = ModCounter(activity_log_count, on_zero_handler=log_something)

    def inc():
        if use_log:
            mc.inc()

    total_file_size = 0
    total_dir_size = 0
    total_files_counter = 0
    for root, dirs, files in os.walk(dir_path, topdown=False):
        for name in files:
            total_file_size += os.path.getsize(os.path.join(root, name))
            total_files_counter += 1
            inc()
        for name in dirs:
            total_dir_size += os.path.getsize(os.path.join(root, name))
            inc()

    print 'done calculating. results:'
    print 'total files', str(total_files_counter)
    print 'total file size', bytes2human(total_file_size)
    print 'total dir size', bytes2human(total_dir_size)


def bytes2human(n):
    symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
    prefix = {}
    for i, s in enumerate(symbols):
        prefix[s] = 1 << (i+1)*10
    for s in reversed(symbols):
        if n >= prefix[s]:
            value = float(n) / prefix[s]
            return '%.1f%s' % (value, s)
    return "%sB" % n

if __name__ == '__main__':
    if not (2 <= len(sys.argv) <= 3):
        print 'usage: %s dir_path [activity_log_count]' % sys.argv[0]
        exit(0)
    dir_path = sys.argv[1]
    activity_log_count = long(sys.argv[2]) if len(sys.argv) == 3 else None
    d_stat(dir_path, activity_log_count)