tree

# coding: utf-8
#
# Copyright (C) 2012, Niklas Rosenstein
#
# Render *.tree file to HD:
#
#   python tree -o out.tree
#   python tree -o out.tree -s /
#   python tree -o out.tree -s /usr/bin/
#
# Read *.tree file and format to human readable text
#
#   python tree -o out-formatted.txt -i out.tree
#
#

import os
import sys
import struct
import argparse
import cStringIO

struct_i = struct.Struct('<i')
pack_i = struct_i.pack
unpack_i = lambda x: struct_i.unpack(x)[0]

struct_b = struct.Struct('<?')
pack_b = struct_b.pack
unpack_b = lambda x: struct_b.unpack(x)[0]

class ReadError(object):

    def format(self):
        return "<ReadError>"

class UnkownParentNodeError(ReadError):

    def __init__(self, id, name, isdir, parent):
        self.id = id
        self.name = name
        self.isdir = isdir
        self.parent = parent

    def format(self):
        return ("Node (id:%d,name:%s,isdir:%s) points to a parent (id:%d) that "
                "could not be found") % (self.id, self.name, bool(self.isdir),
                                         self.parent)

class Node(object):

    def __init__(self, name):
        super(Node, self).__init__()
        self.name = os.path.basename(name)
        self.files = []
        self.folders = []

    def add_folder(self, node):
        if not isinstance(node, Node):
            raise TypeError
        self.folders.append(node)

    def add_file(self, filename):
        filename = os.path.basename(filename)
        self.files.append(filename)

    def format(self, dest, i=0, t='    '):
        dest.write(i * t + self.name + '/\n')
        i += 1
        for filename in self.files:
            dest.write(i * t + filename + '/\n')
        for child in self.folders:
            child.format(dest, i, t)

class Base(object):

    def __init__(self):
        super(Base, self).__init__()
        self.fp = None

    def init(self, fp):
        self.fp = fp

    def free(self):
        self.fp = None

    def close(self):
        self.fp.close()
        self.free()

class Reader(Base):

    def __init__(self):
        super(Reader, self).__init__()

    def init(self, fp):
        super(Reader, self).init(fp)
        self.dirs = {}
        self.files = {}

    def free(self):
        super(Reader, self).free()
        self.dirs = None
        self.files = None

    def read(self, fp):
        self.init(fp)

        errors = []
        root = Node('<root>')

        for name, isdir, id, parent_id in self.iter_nodes():
            if isdir:
                self.dirs[id] = Node(name), parent_id
            else:
                self.files[id] = (name, parent_id)

        for id, (filename, parent_id) in self.files.iteritems():
            parent = None
            if parent_id <= 0:
                parent = root
            elif parent_id in self.dirs:
                parent = self.dirs[parent_id][0]
            else:
                errors.append(UnkownParentNodeError(id, filename, True, parent_id))

            if parent:
                parent.add_file(filename)

        for id, (node, parent_id) in self.dirs.iteritems():
            parent = None
            if parent_id <= 0:
                parent = root
            elif parent_id in self.dirs:
                parent = self.dirs[parent_id][0]
            else:
                errors.append(UnkownParentNodeError(id, node.name, False, parent_id))

            if parent:
                parent.add_folder(node)

        self.free()
        return root, errors

    def read_node(self):
        length = self.fp.read(4)
        if len(length) != 4:
            return None

        length = unpack_i(length)
        name = self.fp.read(length)
        is_dir = unpack_b(self.fp.read(1))
        id = unpack_i(self.fp.read(4))
        parent_id = unpack_i(self.fp.read(4))

        return name, is_dir, id, parent_id

    def iter_nodes(self):
        while True:
            data = self.read_node()
            if not data:
                break
            else:
                yield data

class Writer(Base):

    def write(self, fp, node):
        self.init(fp)
        self.write_hierarchy(node)
        self.free()

    def write_hierarchy(self, node, id=1, parent_id=0):
        for child in node.folders:
            self.write_node(child.name, True, id, parent_id)
            curr_id = id
            id += 1

            for filename in child.files:
                self.write_node(filename, False, id, curr_id)
                id += 1

            self.write_hierarchy(child, id, curr_id)

    def write_node(self, name, isdir, id, parent_id):
        self.fp.write(pack_i(len(name)))
        self.fp.write(name)
        self.fp.write(pack_b(isdir))
        self.fp.write(pack_i(id))
        self.fp.write(pack_i(parent_id))

def collect_nodes(dirname, dest_node=None):
    if not dest_node:
        dest_node = Node('<root>')

    for name in os.listdir(dirname):
        name = os.path.join(dirname, name)
        if os.path.isdir(name):
            node = Node(name)
            collect_nodes(name, node)
            dest_node.add_folder(node)
        else:
            dest_node.add_file(name)

    return dest_node

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', nargs='?', default=None)
    parser.add_argument('-s', '--source', nargs='?', default=os.getcwd())
    parser.add_argument('-o', '--output', nargs='?', default='<out>')
    args = parser.parse_args()

    if args.output == '<out>':
        out = sys.stdout
    else:
        out = open(args.output, 'wb')

    if args.input:
        fp = open(args.input, 'rb')
        reader = Reader()
        root, errors = reader.read(fp)
        root.format(out)

        if not errors:
            print "No errors in file", args.input
        else:
            for error in errors:
                print error.format()

    else:
        if not os.path.isdir(args.source):
            parser.error("path is not a directory")



        node = collect_nodes(args.source)
        writer = Writer()
        writer.write(out, node)

if __name__ == '__main__':
    main()