Completed
Push — master ( f15b1c...db0cb7 )
by Chris
01:03
created

get_tree()   B

Complexity

Conditions 4

Size

Total Lines 24

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 4
c 2
b 0
f 0
dl 0
loc 24
rs 8.6845
1
#!/usr/bin/env python
2
3
"""
4
A utility for getting d3 friendly hierarchical data structures
5
from the list of files and directories on a given path.
6
7
Re-purposed from:
8
github.com/christabor/MoAL/blob/master/MOAL/get_file_tree.py
9
"""
10
11
import os
12
from pprint import pprint
13
import errno
14
import json
15
16
import click
17
18
19
def path_hierarchy(path):
20
    """Create a json representation of a file tree.
21
22
    Format is suitable for d3.js application.
23
24
    Taken from:
25
    http://unix.stackexchange.com/questions/164602/
26
        how-to-output-the-directory-structure-to-json-format
27
    """
28
    name = os.path.basename(path)
29
    hierarchy = {
30
        'type': 'folder',
31
        'name': name,
32
        'path': path,
33
    }
34
    try:
35
        hierarchy['children'] = [
36
            path_hierarchy(os.path.join(path, contents))
37
            for contents in os.listdir(path)
38
        ]
39
    except OSError as e:
40
        if e.errno != errno.ENOTDIR:
41
            raise
42
        hierarchy['type'] = 'file'
43
    return hierarchy
44
45
46
@click.command()
47
@click.option('--ppr/--no-ppr',
48
              default=False,
49
              help='Pretty-print results.')
50
@click.option('--indent', '-i',
51
              default=4,
52
              help='How far to indent if using json.')
53
@click.option('--jsonfile', '-j',
54
              default=None,
55
              help='Output specified file as json.')
56
@click.option('--path', '-p',
57
              default='.',
58
              help='The starting path')
59
def get_tree(path, jsonfile, ppr, indent):
60
    """CLI wrapper for recursive function."""
61
    res = path_hierarchy(path)
62
    if jsonfile is not None:
63
        with open(jsonfile, 'wb+') as jsonfile:
64
            jsonfile.write(json.dumps(res, indent=indent))
65
        return
66
    if ppr:
67
        pprint(res, indent=indent)
68
    else:
69
        print(res)
70
71
72
if __name__ == '__main__':
73
    get_tree()
74