baobab for poor

imagine.py

 #!/usr/bin/env python

import re
import os
import cairo

from hashlib import md5
from colorsys import hls_to_rgb
from optparse import OptionParser
from os.path import abspath, exists, isdir, isfile, join as joinpath

# constants

VARIANCE = 0.95
SATURATION = 0.5
COLORVALUE = 0.99
FILENAME = "out.png"
ORIENTATION = False # vertical
FUCKINGBIGNUMBER = float(2**128)

# main

def getsize(path):
    if isfile(path): return os.path.getsize(path) / 1024
    p = os.popen('du -sc "'+path+'"')
    size = p.readlines()
    p.close()
    if size: return int(size[-1].split('\t')[0])
    return 0


def get_hash(name):
    m = md5()
    m.update(name)
    return int(m.hexdigest(), 16) / FUCKINGBIGNUMBER


def gen_color(color_range, name, variance, saturation, colorvalue):
    u, v = tuple(color_range)
    h = get_hash(name) + 1 
    l = h - (v - u) * variance * 0.5
    r = h + (v - u) * variance * 0.5
    while l < 1.1 and r < 1.1:
        l += 1
        r += 1
    while l > 2.9 and r > 2.9:
        l -= 1
        r -= 1
    col = hls_to_rgb((h * (v - u) + u)%1, saturation, colorvalue)
    return col, [l, r], h


def diver(variance, saturation, colorvalue, verbose = False):
    colorsettings = variance, saturation, colorvalue
    def dive(ctx, colrange, path, paths, pos, dim, orientation, maxs=None):

        if maxs is None:
            print "* fallback"
            maxs = getsize(path)

        color, colrange, id = gen_color(colrange, path, *colorsettings)

        if verbose: print "enter", path, "  size:", maxs, "   dim:", dim

        ctx.rectangle(pos[0], pos[1], dim[0], dim[1])
        ctx.set_source_rgb(*color)
        ctx.fill()

        if not isdir(path): return
        npos = [pos[0] + 1, pos[1] + 1]
        ndim = [dim[0] - 2, dim[1] - 2]
        if ndim[0] < 2 or ndim[1] < 2: return

        names = sorted(paths.keys())
        for name in names:
            next = joinpath(path, name)
            s, ps = paths[name]
            perc = float(s) / float(maxs)

            o = orientation
            i = o * 1

            _dim = [ndim[0], ndim[1]]
            _dim[i] = _dim[i] * perc

            if _dim[0] == _dim[1]:
                o = not o
            else:
                o = _dim[0] < _dim[1]
            dive(ctx, colrange, next, ps, npos, _dim, o, s)

            npos[i] += _dim[i]
    return dive


def get_sizes(path, min_perc):
    p = os.popen('du -a "' + path + '"')
    lines = p.readlines()
    p.close()
    if not lines: return []
    
    duout = re.compile("(?P<size>\d+)\s+(?P<path>.+)")
    p = lambda m: (int(m.group('size')), m.group('path')) if m else None
    size_path = [ p(duout.match(l)) for l in lines ]
    
    if not all(size_path):
        print "cannot parse du output."
        exit(2)

    total = size_path.pop()[0]
    min = int(total * min_perc)

    return total, filter(lambda sp: sp[0] > min, size_path)


def build_tree(min_path, size_path, verbose = False):
    root = {'': (None, {})}
    for size, path in size_path:
        if verbose: print path

        cur = root
        dirs = path.split('/')
        for dir in dirs[:-1]:
            if dir not in cur:
                cur[dir] = (None, {})
            cur = cur[dir][1]

        leaf = dirs[-1]
        if leaf not in cur:
            cur[leaf] = (size, {})
        else:
            cur[leaf] = (size, cur[leaf][1])

    cur = root
    for dir in min_path.split('/'):
        if dir in cur:
            cur = cur[dir][1]
        else:
            print "huh? CRITICAL ERROR", dir, "in", min_path
            break
    return cur


def filter_tree(tree, ignore = True):
    if not ignore: return tree
    result = {}
    for path in tree:
        if not path.startswith('.'):
            item = tree[path]
            result[path] = item[0], filter_tree(item[1], ignore)
    return result


def main():
    parser = OptionParser("usage: %prog [options] path")
    parser.add_option("-v", "--verbose",
                  action="store_true", dest="verbose", default=False,
                  help="make lots of output")
    parser.add_option("-i", "--ignore",
                  action="store_true", dest="ignore", default=False,
                  help="ignore hidden folder")
    parser.add_option("-o", "--output",
                  metavar="FILE", dest="filename", default=FILENAME,
                  help="write output to FILE [default %default]")
    parser.add_option("-c", "--color-variance", type = "float",
                  metavar="PERCENTAGE", dest="variance", default=VARIANCE,
                  help="using only PERCENTAGE of the color range of the upper" +
                    "branch while walking through the tree [default %default]")
    parser.add_option("-S", "--saturation", type = "float",
                  dest="saturation", default=SATURATION,
                  help="color saturation [default %default]")
    parser.add_option("-V", "--value", type = "float",
                  metavar="VALUE", dest="colorvalue", default=COLORVALUE,
                  help="color value [default %default]")
    parser.add_option("-s", "--size", type = "int",
                  dest="size", default=400,
                  help="image size [default %default]")
    opts, args = parser.parse_args()

    path = abspath("." if len(args) != 1 else args[0])
    
    if not exists(path):
        print "path", path, "doesn't exists."
        exit(1)

    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, opts.size, opts.size)
    ctx = cairo.Context(surface)

    print "scanning", path, "..."
    total, size_path = get_sizes(path, 1 / opts.size)
    paths = filter_tree(build_tree(path, size_path, opts.verbose), opts.ignore)

    print "render image ..."
    dive = diver(opts.variance, opts.saturation, opts.colorvalue, opts.verbose)
    dive(ctx,[1,2],path,paths,(0,0),(opts.size,opts.size),ORIENTATION,total)

    print "write" , opts.filename , "..."
    surface.write_to_png(opts.filename)

if __name__ == "__main__":
    main()

It's drawing folder structure.

Usage: imagine.py [options] path

Options:
  -h, --help            show this help message and exit
  -v, --verbose         make lots of output
  -i, --ignore          ignore hidden folder
  -o FILE, --output=FILE
                          write output to FILE [default out.png]
  -c PERCENTAGE, --color-variance=PERCENTAGE
                          using only PERCENTAGE of the color range of the
                          upperbranch while walking through the tree [default
                          0.95]
  -S SATURATION, --saturation=SATURATION
                          color saturation [default 0.5]
  -V VALUE, --value=VALUE
                          color value [default 0.99]
  -s SIZE, --size=SIZE  image size [default 400]