Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- # 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()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement