-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathexport-module-tree.py
More file actions
189 lines (159 loc) · 5.79 KB
/
export-module-tree.py
File metadata and controls
189 lines (159 loc) · 5.79 KB
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
"""Export the module structure with docstrings to a tree diagram file.
Scans a Python package directory, builds a tree of module names with first-line
module docstrings, and writes the result to an output file using Unicode
box-drawing characters.
"""
from __future__ import annotations
import argparse
import ast
import os
import sys
from pathlib import Path
BRANCH = "├── "
LAST_BRANCH = "└── "
PIPE = "│ "
SPACE = " "
# Relative paths (from the source root) to exclude from the tree and strict
# enforcement.
AUTO_EXCLUDED = {
"_version.py", # These are typically generated by dynamic VCS-based versioning tools.
}
def main() -> int:
parser = argparse.ArgumentParser(
description="Export module structure with docstrings to a tree diagram.",
)
parser.add_argument(
"--source-root",
required=True,
help="Path to the root package directory to scan.",
)
parser.add_argument(
"--output-file",
required=True,
help="Path to the output file to write the tree diagram to.",
)
parser.add_argument(
"--strict",
action="store_true",
default=False,
help="Fail if any module is missing a docstring.",
)
args = parser.parse_args()
source_root = Path(args.source_root)
output_file = Path(args.output_file)
if not source_root.is_dir():
print(f"ERROR: Source root {source_root} is not a directory.", file=sys.stderr)
return 1
init_path = source_root / "__init__.py"
if not init_path.is_file():
print(
f"ERROR: {source_root} is not a Python package (no __init__.py).",
file=sys.stderr,
)
return 1
# Build the tree.
missing: list[Path] = []
root_name = source_root.name
root_docstring = _get_first_docstring_line(init_path)
if root_docstring is None:
missing.append(init_path)
root_comment = f"# {root_docstring}" if root_docstring else ""
root_pair = (root_name, root_comment)
child_pairs = _build_tree(source_root, "", missing, source_root)
all_pairs = [root_pair, *child_pairs]
# Align comments: pad each line so comments start at the same column.
max_width = max(len(text) for text, _ in all_pairs)
formatted: list[str] = []
for text, comment in all_pairs:
if comment:
formatted.append(f"{text:<{max_width}} {comment}")
else:
formatted.append(text)
content = os.linesep.join(formatted) + os.linesep
try:
with open(output_file, encoding="utf-8", newline="") as f:
existing = f.read()
except FileNotFoundError:
existing = None
modified = content != existing
if modified:
output_file.parent.mkdir(parents=True, exist_ok=True)
output_file.write_text(content, encoding="utf-8", newline="")
print(f"Module tree written to {output_file}.")
else:
print("Module tree is already up to date.")
if args.strict and missing:
print(
f"ERROR: {len(missing)} module(s) missing a docstring:",
file=sys.stderr,
)
for path in missing:
print(f" - {path}", file=sys.stderr)
return 1
return 1 if modified else 0
def _get_first_docstring_line(path: Path) -> str | None:
"""Return the first line of the module docstring, or None if absent."""
try:
source = path.read_text(encoding="utf-8")
except (OSError, UnicodeDecodeError):
return None
try:
tree = ast.parse(source)
except SyntaxError:
return None
docstring = ast.get_docstring(tree)
if docstring is None:
return None
return docstring.split("\n")[0].strip()
def _build_tree(
root: Path,
prefix: str,
missing: list[Path],
source_root: Path,
) -> list[tuple[str, str]]:
"""Recursively build tree lines for a package directory.
Returns a list of (line_text, comment) pairs so the caller can align
comments vertically.
"""
lines: list[tuple[str, str]] = []
# Collect .py files (excluding __init__.py) and subpackages.
py_files: list[Path] = []
subpackages: list[Path] = []
for entry in sorted(root.iterdir()):
if entry.name in ("__pycache__", "py.typed"):
continue
if entry.relative_to(source_root).as_posix() in AUTO_EXCLUDED:
continue
if entry.is_file() and entry.suffix == ".py" and entry.name != "__init__.py":
py_files.append(entry)
elif entry.is_dir() and (entry / "__init__.py").is_file():
subpackages.append(entry)
entries: list[Path] = [*py_files, *subpackages]
if not entries:
return lines
for i, entry in enumerate(entries):
is_last = i == len(entries) - 1
connector = LAST_BRANCH if is_last else BRANCH
child_prefix = prefix + (SPACE if is_last else PIPE)
if entry.is_file():
name = entry.stem
docstring = _get_first_docstring_line(entry)
if docstring is None:
missing.append(entry)
line_text = f"{prefix}{connector}{name}"
comment = f"# {docstring}" if docstring else ""
lines.append((line_text, comment))
else:
# Subpackage directory — attribute __init__.py docstring to the dir.
name = entry.name
init_path = entry / "__init__.py"
docstring = _get_first_docstring_line(init_path)
if docstring is None:
missing.append(init_path)
line_text = f"{prefix}{connector}{name}"
comment = f"# {docstring}" if docstring else ""
lines.append((line_text, comment))
lines.extend(_build_tree(entry, child_prefix, missing, source_root))
return lines
if __name__ == "__main__":
raise SystemExit(main())