Mercurial > hg > config
annotate python/tree.py @ 701:de7bf9523e21
-
| author | Jeff Hammel <k0scist@gmail.com> |
|---|---|
| date | Wed, 13 Aug 2014 13:16:56 -0700 |
| parents | eeb38dfa17d0 |
| children |
| rev | line source |
|---|---|
| 382 | 1 #!/usr/bin/env python |
| 2 # -*- coding: utf-8 -*- | |
| 3 | |
| 4 """ | |
| 5 tree in python | |
| 6 """ | |
| 7 | |
| 425 | 8 # TODO: script2package |
| 9 | |
| 670 | 10 import argparse |
| 382 | 11 import os |
| 12 import sys | |
| 13 | |
| 387 | 14 # ASCII delimeters |
| 388 | 15 ascii_delimeters = { |
| 389 | 16 'vertical_line' : '|', |
| 17 'item_marker' : '+', | |
| 18 'last_child' : '\\' | |
| 388 | 19 } |
| 387 | 20 |
| 21 # unicode delimiters | |
| 389 | 22 unicode_delimeters = { |
| 23 'vertical_line' : '│', | |
| 24 'item_marker' : '├', | |
| 25 'last_child' : '└' | |
| 26 } | |
| 382 | 27 |
| 425 | 28 |
| 382 | 29 def depth(directory): |
| 387 | 30 """returns the integer depth of a directory or path relative to '/' """ |
| 31 | |
| 382 | 32 directory = os.path.abspath(directory) |
| 33 level = 0 | |
| 34 while True: | |
| 35 directory, remainder = os.path.split(directory) | |
| 36 level += 1 | |
| 37 if not remainder: | |
| 38 break | |
| 39 return level | |
| 40 | |
| 425 | 41 |
| 42 ### stuff for tree generalization | |
| 43 | |
| 44 class Tree(object): | |
| 45 """tree structure in python""" | |
| 46 | |
| 47 def __init__(self, parent=None): | |
| 48 self.parent = parent | |
| 426 | 49 self._children = [] |
| 425 | 50 |
| 51 def children(self): | |
| 52 """returns children of the tree""" | |
| 426 | 53 return self._children # base implementation |
| 425 | 54 |
| 55 def add(self, item): | |
| 56 """add a child to the tree root""" | |
| 57 | |
| 58 def update(self, tree): | |
| 59 """add a subtree to the tree""" | |
| 60 self.add(tree) | |
| 61 tree.parent = self # XXX .add should probably do this for scary reasons | |
| 62 | |
| 63 def output(self, serializer): | |
| 64 """output the tree via the given serializer""" | |
| 65 # XXX or should this method exist at all and instead the | |
| 66 # __call__ method of serializers take a Tree object? | |
| 67 | |
| 68 class DirectoryTree(Tree): | |
| 69 """directory structure as a tree""" | |
| 70 | |
| 429 | 71 def __init__(self, directory): |
| 72 self.directory = directory | |
| 73 self._return_type = os.path.abspath | |
| 74 | |
| 75 def children(self): | |
| 76 return os.listdir(self.directory) # -> self._return_type | |
| 77 | |
| 78 ## Serializers | |
| 79 | |
| 443 | 80 def tree2html(tree): |
| 81 """ | |
| 82 . | |
| 83 ├── a8e.py | |
| 84 ├── abstract.py | |
| 85 ├── accentuate.py | |
| 86 | |
| 87 -> | |
| 88 | |
| 89 <ul> | |
| 90 <li>a8e.py</li> | |
| 91 ... | |
| 92 """ | |
| 93 | |
| 94 | |
| 429 | 95 # How to serialize a tree -> JSON? |
| 96 | |
| 425 | 97 ### |
| 98 | |
| 388 | 99 def tree(directory, |
| 389 | 100 item_marker=unicode_delimeters['item_marker'], |
| 101 vertical_line=unicode_delimeters['vertical_line'], | |
| 102 last_child=unicode_delimeters['last_child'], | |
| 670 | 103 display_files=True, |
| 388 | 104 sort_key=lambda x: x.lower()): |
| 105 """ | |
| 106 display tree directory structure for `directory` | |
| 107 """ | |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
108 |
| 385 | 109 retval = [] |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
110 indent = [] |
| 384 | 111 last = {} |
| 386 | 112 top = depth(directory) |
| 113 | |
| 382 | 114 for dirpath, dirnames, filenames in os.walk(directory, topdown=True): |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
115 |
|
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
116 abspath = os.path.abspath(dirpath) |
| 384 | 117 basename = os.path.basename(abspath) |
| 118 parent = os.path.dirname(abspath) | |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
119 level = depth(abspath) - top |
|
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
120 |
| 671 | 121 # omit files if specified |
| 122 if not display_files: | |
| 123 filenames = [] | |
| 124 | |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
125 # sort articles of interest |
|
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
126 for resource in (dirnames, filenames): |
|
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
127 resource[:] = sorted(resource, key=sort_key) |
| 382 | 128 |
| 389 | 129 files_end = item_marker |
| 130 dirpath_marker = item_marker | |
| 385 | 131 |
| 132 if level > len(indent): | |
| 389 | 133 indent.append(vertical_line) |
| 385 | 134 indent = indent[:level] |
| 135 | |
|
383
8d1ad56761b0
this still, somehow, eludes my tired brain
Jeff Hammel <jhammel@mozilla.com>
parents:
382
diff
changeset
|
136 if dirnames: |
| 389 | 137 files_end = item_marker |
| 385 | 138 last[abspath] = dirnames[-1] |
| 139 else: | |
| 389 | 140 files_end = last_child |
| 384 | 141 |
| 385 | 142 if last.get(parent) == os.path.basename(abspath): |
| 143 # last directory of parent | |
| 389 | 144 dirpath_mark = last_child |
| 385 | 145 indent[-1] = ' ' |
| 146 elif not indent: | |
| 147 dirpath_mark = '' | |
| 148 else: | |
| 389 | 149 dirpath_mark = item_marker |
| 385 | 150 |
| 388 | 151 # append the directory and piece of tree structure |
| 152 # if the top-level entry directory, print as passed | |
| 153 retval.append('%s%s%s'% (''.join(indent[:-1]), | |
| 154 dirpath_mark, | |
| 155 basename if retval else directory)) | |
| 156 # add the files | |
| 385 | 157 if filenames: |
| 158 last_file = filenames[-1] | |
| 387 | 159 retval.extend([('%s%s%s' % (''.join(indent), |
| 389 | 160 files_end if filename == last_file else item_marker, |
| 385 | 161 filename)) |
| 162 for index, filename in enumerate(filenames)]) | |
| 163 | |
| 382 | 164 return '\n'.join(retval) |
| 165 | |
| 671 | 166 |
| 382 | 167 def main(args=sys.argv[1:]): |
| 670 | 168 """CLI""" |
| 382 | 169 |
|
390
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
170 # parse command line options |
| 670 | 171 parser = argparse.ArgumentParser(description=__doc__) |
| 172 parser.add_argument('-a', '--ascii', dest='use_ascii', | |
| 173 action='store_true', default=False, | |
| 174 help="use ascii delimeters ({})".format(', '.join(ascii_delimeters.values()))) | |
| 671 | 175 parser.add_argument('-d', '--no-files', dest='display_files', |
| 176 action='store_false', default=True, | |
| 177 help='only display directories') | |
| 670 | 178 parser.add_argument('path', nargs='*', |
| 671 | 179 help="directory paths to display the tree of") |
| 670 | 180 options = parser.parse_args(args) |
| 181 | |
| 182 # get paths to operate on | |
| 183 paths = options.path | |
| 184 if not paths: | |
| 185 paths = ['.'] | |
| 382 | 186 |
| 671 | 187 # ensure each path is a directory |
| 670 | 188 not_directory = [arg for arg in paths |
| 382 | 189 if not os.path.isdir(arg)] |
| 190 if not_directory: | |
| 671 | 191 parser.error("Not a directory: {}".format(', '.join(not_directory))) |
| 382 | 192 |
|
390
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
193 delimeters = unicode_delimeters |
|
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
194 if options.use_ascii: |
|
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
195 delimeters = ascii_delimeters |
|
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
196 |
| 671 | 197 # build function arguments |
| 198 kw = delimeters.copy() | |
| 199 kw['display_files'] = options.display_files | |
| 200 | |
|
390
9d02187611ae
make delimeters CLI switchable
Jeff Hammel <jhammel@mozilla.com>
parents:
389
diff
changeset
|
201 # print the tree |
| 670 | 202 for arg in paths: |
| 671 | 203 print (tree(arg, **kw)) |
| 382 | 204 |
| 205 if __name__ == '__main__': | |
| 206 main() |