-
Notifications
You must be signed in to change notification settings - Fork 11
/
update_tags_init.py
114 lines (90 loc) · 3.25 KB
/
update_tags_init.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import argparse
import ast
import glob
import io
import itertools
import os
import warnings
from collections import defaultdict
class Universe:
def __init__(self):
self.cls_modules = {}
self.cls_bases = defaultdict(set)
def get_all_bases(self, name):
bases = set()
open = list(self.cls_bases[name])
while open:
name = open.pop(0)
if name not in bases:
yield name
bases.add(name)
open.extend(list(self.cls_bases[name]))
def get_classes_deriving_from(self, base_names):
for name in list(self.cls_bases.keys()):
for base in self.get_all_bases(name):
if base in base_names:
yield name
break
class ClassDefWalker(ast.NodeVisitor):
def __init__(self, module, universe):
self.module = module
self.universe = universe
def _stringify_id_node(self, node):
if isinstance(node, ast.Name):
return node.id
elif isinstance(node, ast.Attribute):
return f'{self._stringify_id_node(node.value)}.{self._stringify_id_node(node.attr)}'
elif isinstance(node, str):
return node
else:
warnings.warn(f'Dont know how to stringify {node}')
return '<none>'
def visit_ClassDef(self, classdef):
for base in classdef.bases:
base_name = self._stringify_id_node(base)
self.universe.cls_bases[classdef.name].add(base_name)
self.universe.cls_modules[classdef.name] = self.module
def find_tags_in_modules():
tags_per_module = defaultdict(set)
universe = Universe()
for filepath in glob.glob('emrichen/tags/*.py'):
modname = os.path.splitext(os.path.basename(filepath))[0]
if modname.startswith('_'):
continue
with open(filepath) as infp:
tree = ast.parse(infp.read(), filepath)
cdw = ClassDefWalker(module=modname, universe=universe)
cdw.visit(tree)
for tag in universe.get_classes_deriving_from(('BaseTag',)):
if not tag.startswith('_'):
tags_per_module[universe.cls_modules[tag]].add(tag)
return tags_per_module
def format_tags_init_file(tags_per_module):
sio = io.StringIO()
print("# This file is autogenerated by update_tags_init.py; please just run it.\n", file=sio)
for modname, tags in sorted(tags_per_module.items()):
if tags:
print(f'from .{modname} import {", ".join(sorted(tags))}', file=sio)
print('\n__all__ = [', file=sio)
for tag in sorted(itertools.chain(*tags_per_module.values())):
print(f' {tag!r},', file=sio)
print(']', file=sio)
return sio.getvalue()
def main():
ap = argparse.ArgumentParser()
ap.add_argument(
'--dry-run',
'-n',
action='store_true',
help='do not update tags/__init__, just print the new one',
)
args = ap.parse_args()
tags_per_module = find_tags_in_modules()
init_file = format_tags_init_file(tags_per_module)
if not args.dry_run:
with open('emrichen/tags/__init__.py', 'w') as outfp:
outfp.write(init_file)
else:
print(init_file)
if __name__ == '__main__':
main()