Disk Usage Analyser

analyser.py

import os
import sys
import colors


class Node():
    _OPT = {
        'depth': False,
        'colors': True,
        'abspath': False,
        'files': False,
        'dirs': True,
        'dot': False
    }

    def __init__(self, path, name=''):
        self.children = []
        self.path = path
        self.name = name
        self.size = 0

    def output(self, _depth=0, **options):
        for key in self._OPT:
            if key not in options:
                options[key] = self._OPT[key]

        if options['abspath']:
            name = self.abspath
        elif self.name:
            name = self.name
        else:
            name = self.path

        string = '{} {} {}'.format(
            _depth * '--',
            name,
            self.size
        )

        if (self.isdir and options['dirs']
                or self.isfile and options['files']):
            if options['colors']:
                if self.isdir:
                    print(colors.colorStr(
                        text=string,
                        fg=_depth,
                        style=colors.BOLD
                    ))
                else:
                    print(colors.colorStr(
                        text=string,
                        fg=_depth
                    ))
            else:
                print(string)

        if not options['depth'] or _depth < options['depth']:
            for child in self.children:
                hidden = child.name.startswith('.')
                if not hidden or options['dot'] and hidden:
                    child.output(_depth + 1, **options)

    def explore(self):
        path = self.abspath
        for name in os.listdir(path):
            child = Node(path, name)
            if child.isdir:
                child.size = child.explore()
            else:
                child.size = os.path.getsize(child.abspath)
            self.children.append(child)
            self.size += child.size
        return self.size

    @property
    def isdir(self):
        return os.path.isdir(self.abspath)

    @property
    def isfile(self):
        return os.path.isfile(self.abspath)

    @property
    def abspath(self):
        return os.path.join(self.path, self.name)

    @property
    def depth(self):
        depths = [child.depth for child in self.children]
        return (1 + max(depths)) if depths else 0


if __name__ == '__main__':
    try:
        directory = os.path.realpath(sys.argv[1])
    except IndexError:
        directory = os.getcwd()

    root = Node(directory)
    root.explore()
    root.output(files=True)