|
| 1 | +import os |
| 2 | +from .defaults import DEFAULT_COMMENT_STYLES, DEFAULT_EXCLUDE_DIRS |
| 3 | + |
| 4 | + |
| 5 | +def get_full_extension(filename: str) -> str: |
| 6 | + """Return everything after the first dot, e.g., 'unit.test.js' -> '.test.js'.""" |
| 7 | + if "." not in filename: |
| 8 | + return "" |
| 9 | + return filename[filename.index("."):] # keeps the dot |
| 10 | + |
| 11 | + |
| 12 | +def is_excluded_dir(filepath: str, root: str, exclude_dirs: set) -> bool: |
| 13 | + """Check if file lives in an excluded directory (relative to root).""" |
| 14 | + rel_path = os.path.relpath(filepath, root) |
| 15 | + parts = rel_path.split(os.sep) |
| 16 | + return any(part in exclude_dirs for part in parts) |
| 17 | + |
| 18 | + |
| 19 | +def is_excluded_file(filename: str, exclude_files: set) -> bool: |
| 20 | + """Check if filename matches exactly one of the excluded files.""" |
| 21 | + return filename in exclude_files |
| 22 | + |
| 23 | + |
| 24 | +def is_excluded_ext(filename: str, exclude_exts: set) -> bool: |
| 25 | + """Check if full extension matches excluded extensions.""" |
| 26 | + ext = get_full_extension(filename) |
| 27 | + return ext in exclude_exts |
| 28 | + |
| 29 | + |
| 30 | +def annotate_file(root, filepath, comment_styles): |
| 31 | + rel_path = os.path.relpath(filepath, root) |
| 32 | + ext = get_full_extension(os.path.basename(filepath)) |
| 33 | + |
| 34 | + prefix = comment_styles.get(ext) |
| 35 | + if not prefix: |
| 36 | + return |
| 37 | + |
| 38 | + annotation = f"{prefix} {rel_path}\n" |
| 39 | + if prefix in ["<!--", "/*"]: |
| 40 | + annotation = f"{prefix} {rel_path} {'-->' if prefix == '<!--' else '*/'}\n" |
| 41 | + |
| 42 | + try: |
| 43 | + with open(filepath, "r", encoding="utf-8", errors="ignore") as f: |
| 44 | + lines = f.readlines() |
| 45 | + except Exception as e: |
| 46 | + print(f"[WARN] Skipping {filepath}: {e}") |
| 47 | + return |
| 48 | + |
| 49 | + if lines and rel_path in lines[0]: |
| 50 | + return |
| 51 | + |
| 52 | + new_content = [annotation] + lines |
| 53 | + try: |
| 54 | + with open(filepath, "w", encoding="utf-8") as f: |
| 55 | + f.writelines(new_content) |
| 56 | + print(f"[OK] Annotated {rel_path}") |
| 57 | + except Exception as e: |
| 58 | + print(f"[ERROR] Failed to write {filepath}: {e}") |
| 59 | + |
| 60 | + |
| 61 | +def annotate_project(root=".", comment_styles=None, |
| 62 | + exclude_exts=None, exclude_dirs=None, exclude_files=None): |
| 63 | + if comment_styles is None: |
| 64 | + comment_styles = DEFAULT_COMMENT_STYLES |
| 65 | + if exclude_exts is None: |
| 66 | + exclude_exts = set() |
| 67 | + if exclude_dirs is None: |
| 68 | + exclude_dirs = set() |
| 69 | + if exclude_files is None: |
| 70 | + exclude_files = set() |
| 71 | + |
| 72 | + effective_exclude_dirs = set(DEFAULT_EXCLUDE_DIRS) | set(exclude_dirs) |
| 73 | + |
| 74 | + for subdir, dirs, files in os.walk(root): |
| 75 | + # skip excluded dirs at traversal level |
| 76 | + dirs[:] = [d for d in dirs if d not in effective_exclude_dirs] |
| 77 | + |
| 78 | + for file in files: |
| 79 | + filepath = os.path.join(subdir, file) |
| 80 | + |
| 81 | + # precedence rules |
| 82 | + if is_excluded_dir(filepath, root, effective_exclude_dirs): |
| 83 | + continue |
| 84 | + if is_excluded_file(file, exclude_files): |
| 85 | + continue |
| 86 | + if is_excluded_ext(file, exclude_exts): |
| 87 | + continue |
| 88 | + |
| 89 | + annotate_file(root, filepath, comment_styles) |
0 commit comments