| 1 | """Scan workdir for promptinclude files. No agent/tool dependencies.""" |
| 2 | |
| 3 | import fnmatch |
| 4 | import os |
| 5 | from typing import Literal, TypedDict |
| 6 | |
| 7 | from pathspec import PathSpec |
| 8 | |
| 9 | from helpers import tokens |
| 10 | |
| 11 | |
| 12 | # ------------------------------------------------------------------ |
| 13 | # Types |
| 14 | # ------------------------------------------------------------------ |
| 15 | |
| 16 | class FileEntry(TypedDict): |
| 17 | path: str |
| 18 | content: str |
| 19 | token_count: int |
| 20 | status: Literal["ok", "cropped", "skipped"] |
| 21 | |
| 22 | |
| 23 | class ScanResult(TypedDict): |
| 24 | files: list[FileEntry] |
| 25 | skipped_count: int |
| 26 | |
| 27 | |
| 28 | # ------------------------------------------------------------------ |
| 29 | # Public API |
| 30 | # ------------------------------------------------------------------ |
| 31 | |
| 32 | def scan_promptinclude_files( |
| 33 | root: str, |
| 34 | *, |
| 35 | name_pattern: str = "*.promptinclude.md", |
| 36 | max_depth: int = 10, |
| 37 | max_file_tokens: int = 2000, |
| 38 | max_file_count: int = 50, |
| 39 | max_total_tokens: int = 8000, |
| 40 | gitignore: str = "", |
| 41 | ) -> ScanResult: |
| 42 | ignore_spec = _build_ignore_spec(gitignore) |
| 43 | matched = _find_matching_files(root, name_pattern, max_depth, ignore_spec) |
| 44 | matched.sort() |
| 45 | |
| 46 | result_files: list[FileEntry] = [] |
| 47 | total_tokens_used = 0 |
| 48 | skipped_count = 0 |
| 49 | budget_exhausted = False |
| 50 | |
| 51 | for path in matched: |
| 52 | if budget_exhausted or len(result_files) >= max_file_count: |
| 53 | skipped_count += 1 |
| 54 | continue |
| 55 | |
| 56 | try: |
| 57 | with open(path, "r", encoding="utf-8", errors="replace") as f: |
| 58 | raw = f.read() |
| 59 | except (OSError, IOError): |
| 60 | skipped_count += 1 |
| 61 | continue |
| 62 | |
| 63 | if not raw.strip(): |
| 64 | continue |
| 65 | |
| 66 | file_tokens = tokens.count_tokens(raw) |
| 67 | |
| 68 | # check if adding path line alone exceeds budget |
| 69 | path_tokens = tokens.count_tokens(path) + 5 # overhead for formatting |
| 70 | if total_tokens_used + path_tokens > max_total_tokens: |
| 71 | skipped_count += 1 |
| 72 | budget_exhausted = True |
| 73 | continue |
| 74 | |
| 75 | # per-file token cap |
| 76 | capped = min(file_tokens, max_file_tokens) |
| 77 | |
| 78 | if total_tokens_used + path_tokens + capped > max_total_tokens: |
| 79 | # try to fit partial |
| 80 | remaining = max_total_tokens - total_tokens_used - path_tokens |
| 81 | if remaining > 50: |
| 82 | trimmed = tokens.trim_to_tokens(raw, remaining, direction="start") |
| 83 | trimmed_count = tokens.count_tokens(trimmed) |
| 84 | total_tokens_used += path_tokens + trimmed_count |
| 85 | result_files.append(FileEntry( |
| 86 | path=path, content=trimmed, |
| 87 | token_count=trimmed_count, status="cropped", |
| 88 | )) |
| 89 | else: |
| 90 | result_files.append(FileEntry( |
| 91 | path=path, content="", |
| 92 | token_count=0, status="skipped", |
| 93 | )) |
| 94 | budget_exhausted = True |
| 95 | continue |
| 96 | |
| 97 | if capped < file_tokens: |
| 98 | trimmed = tokens.trim_to_tokens(raw, max_file_tokens, direction="start") |
| 99 | trimmed_count = tokens.count_tokens(trimmed) |
| 100 | total_tokens_used += path_tokens + trimmed_count |
| 101 | result_files.append(FileEntry( |
| 102 | path=path, content=trimmed, |
| 103 | token_count=trimmed_count, status="cropped", |
| 104 | )) |
| 105 | else: |
| 106 | total_tokens_used += path_tokens + file_tokens |
| 107 | result_files.append(FileEntry( |
| 108 | path=path, content=raw, |
| 109 | token_count=file_tokens, status="ok", |
| 110 | )) |
| 111 | |
| 112 | # remaining unprocessed files from matched list |
| 113 | remaining_unprocessed = len(matched) - len(result_files) - skipped_count |
| 114 | if remaining_unprocessed > 0: |
| 115 | skipped_count += remaining_unprocessed |
| 116 | |
| 117 | return ScanResult(files=result_files, skipped_count=skipped_count) |
| 118 | |
| 119 | |
| 120 | # ------------------------------------------------------------------ |
| 121 | # Internal helpers |
| 122 | # ------------------------------------------------------------------ |
| 123 | |
| 124 | def _build_ignore_spec(gitignore: str) -> PathSpec | None: |
| 125 | if not gitignore or not gitignore.strip(): |
| 126 | return None |
| 127 | lines = [ |
| 128 | line.strip() |
| 129 | for line in gitignore.splitlines() |
| 130 | if line.strip() and not line.strip().startswith("#") |
| 131 | ] |
| 132 | if not lines: |
| 133 | return None |
| 134 | return PathSpec.from_lines("gitignore", lines) |
| 135 | |
| 136 | |
| 137 | def _find_matching_files( |
| 138 | root: str, |
| 139 | name_pattern: str, |
| 140 | max_depth: int, |
| 141 | ignore_spec: PathSpec | None, |
| 142 | ) -> list[str]: |
| 143 | root = os.path.abspath(root) |
| 144 | if not os.path.isdir(root): |
| 145 | return [] |
| 146 | |
| 147 | results: list[str] = [] |
| 148 | |
| 149 | for dirpath, dirnames, filenames in os.walk(root, topdown=True): |
| 150 | depth = dirpath[len(root):].count(os.sep) |
| 151 | if depth >= max_depth: |
| 152 | dirnames.clear() |
| 153 | continue |
| 154 | |
| 155 | # filter ignored dirs in-place |
| 156 | if ignore_spec: |
| 157 | filtered_dirs = [] |
| 158 | for d in dirnames: |
| 159 | rel = os.path.relpath(os.path.join(dirpath, d), root) |
| 160 | rel_posix = rel.replace(os.sep, "/") |
| 161 | if ignore_spec.match_file(rel_posix) or ignore_spec.match_file(f"{rel_posix}/"): |
| 162 | continue |
| 163 | filtered_dirs.append(d) |
| 164 | dirnames[:] = filtered_dirs |
| 165 | |
| 166 | for fname in filenames: |
| 167 | if not fnmatch.fnmatch(fname, name_pattern): |
| 168 | continue |
| 169 | full = os.path.join(dirpath, fname) |
| 170 | if ignore_spec: |
| 171 | rel = os.path.relpath(full, root) |
| 172 | rel_posix = rel.replace(os.sep, "/") |
| 173 | if ignore_spec.match_file(rel_posix): |
| 174 | continue |
| 175 | results.append(full) |
| 176 | |
| 177 | return results |