Completed
Push — master ( c2a5ab...6a5067 )
by Chris
01:05
created

get_tree()   A

Complexity

Conditions 4

Size

Total Lines 21

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 4
dl 0
loc 21
rs 9.0534
c 1
b 0
f 0
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',
48
              default=False,
49
              help='Pretty-print results.')
50
@click.option('--jsonfile', '-j',
51
              default=None,
52
              help='Output specified file as json.')
53
@click.option('--path', '-p',
54
              default='.',
55
              help='The starting path')
56
def get_tree(path, jsonfile, ppr):
57
    """CLI wrapper for recursive function."""
58
    res = path_hierarchy(path)
59
    if jsonfile is not None:
60
        with open(jsonfile, 'wb+') as jsonfile:
61
            jsonfile.write(json.dumps(res, indent=4))
62
        return
63
    if ppr:
64
        pprint(res)
65
    else:
66
        print(res)
67
68
69
if __name__ == '__main__':
70
    get_tree()
71