|
14 | 14 | from gitingest.schemas import Source, FileSystemFile, FileSystemDirectory, FileSystemSymlink |
15 | 15 | from gitingest.schemas.filesystem import SEPARATOR |
16 | 16 | from gitingest.utils.logging_config import get_logger |
| 17 | +from jinja2 import Environment, BaseLoader |
17 | 18 |
|
18 | 19 | if TYPE_CHECKING: |
19 | 20 | from gitingest.schemas import IngestionQuery |
|
27 | 28 | ] |
28 | 29 |
|
29 | 30 |
|
30 | | -class Formatter: |
31 | | - @singledispatchmethod |
32 | | - def format(self, node: Source, query): |
33 | | - return f"{getattr(node, 'content', '')}" |
34 | | - |
35 | | - @format.register |
36 | | - def _(self, node: FileSystemFile, query): |
37 | | - return ( |
38 | | - f"{SEPARATOR}\n" |
39 | | - f"{node.name}\n" |
40 | | - f"{SEPARATOR}\n\n" |
41 | | - f"{node.content}" |
42 | | - ) |
43 | | - |
44 | | - @format.register |
45 | | - def _(self, node: FileSystemDirectory, query): |
46 | | - formatted = [] |
47 | | - for child in node.children: |
48 | | - formatted.append(self.format(child, query)) |
49 | | - return "\n".join(formatted) |
50 | | - |
51 | | - @format.register |
52 | | - def _(self, node: FileSystemSymlink, query): |
53 | | - target = getattr(node, 'target', None) |
54 | | - target_str = f" -> {target}" if target else "" |
55 | | - return ( |
56 | | - f"{SEPARATOR}\n" |
57 | | - f"{node.name}{target_str}\n" |
58 | | - f"{SEPARATOR}\n" |
59 | | - ) |
60 | | - |
61 | | -class DefaultFormatter(Formatter): |
62 | | - pass |
63 | | - |
64 | 31 | # Backward compatibility |
65 | 32 |
|
66 | 33 | def _create_summary_prefix(query: IngestionQuery, *, single_file: bool = False) -> str: |
@@ -209,3 +176,71 @@ def _format_token_count(text: str) -> str | None: |
209 | 176 | return f"{total_tokens / threshold:.1f}{suffix}" |
210 | 177 |
|
211 | 178 | return str(total_tokens) |
| 179 | + |
| 180 | +# Rename JinjaFormatter to DefaultFormatter throughout the file |
| 181 | +class DefaultFormatter: |
| 182 | + def __init__(self): |
| 183 | + self.env = Environment(loader=BaseLoader()) |
| 184 | + |
| 185 | + @singledispatchmethod |
| 186 | + def format(self, node: Source, query): |
| 187 | + return f"{getattr(node, 'content', '')}" |
| 188 | + |
| 189 | + @singledispatchmethod |
| 190 | + def summary(self, node: Source, query): |
| 191 | + # Default summary: just the name |
| 192 | + return f"{getattr(node, 'name', '')}" |
| 193 | + |
| 194 | + @format.register |
| 195 | + def _(self, node: FileSystemFile, query): |
| 196 | + template = \ |
| 197 | +""" |
| 198 | +{{ SEPARATOR }} |
| 199 | +{{ node.name }} |
| 200 | +{{ SEPARATOR }} |
| 201 | +
|
| 202 | +{{ node.content }} |
| 203 | +""" |
| 204 | + file_template = self.env.from_string(template) |
| 205 | + return file_template.render(SEPARATOR=SEPARATOR, node=node, query=query, formatter=self) |
| 206 | + |
| 207 | + @format.register |
| 208 | + def _(self, node: FileSystemDirectory, query): |
| 209 | + template = \ |
| 210 | +""" |
| 211 | +{% for child in node.children %} |
| 212 | +{{ formatter.format(child, query) }} |
| 213 | +{% endfor %} |
| 214 | +""" |
| 215 | + dir_template = self.env.from_string(template) |
| 216 | + return dir_template.render(node=node, query=query, formatter=self) |
| 217 | + |
| 218 | + @summary.register |
| 219 | + def _(self, node: FileSystemDirectory, query): |
| 220 | + template = """ |
| 221 | +{%- macro render_tree(node, prefix='', is_last=True) -%} |
| 222 | + {{ prefix }}{{ '└── ' if is_last else '├── ' }}{{ node.name }}{% if node.type == 'directory' %}/{% endif %} |
| 223 | + {%- if node.type == 'directory' and node.children %} |
| 224 | + {%- for i, child in enumerate(node.children) %} |
| 225 | + {{ render_tree(child, prefix + (' ' if is_last else '│ '), i == (node.children | length - 1)) }} |
| 226 | + {%- endfor %} |
| 227 | + {%- endif %} |
| 228 | +{%- endmacro %} |
| 229 | +
|
| 230 | +Directory structure: |
| 231 | +{{ render_tree(node) }} |
| 232 | +""" |
| 233 | + summary_template = self.env.from_string(template) |
| 234 | + return summary_template.render(node=node, query=query, formatter=self) |
| 235 | + |
| 236 | + |
| 237 | + @format.register |
| 238 | + def _(self, node: FileSystemSymlink, query): |
| 239 | + template = \ |
| 240 | +""" |
| 241 | +{{ SEPARATOR }} |
| 242 | +{{ node.name }}{% if node.target %} -> {{ node.target }}{% endif %} |
| 243 | +{{ SEPARATOR }} |
| 244 | +""" |
| 245 | + symlink_template = self.env.from_string(template) |
| 246 | + return symlink_template.render(SEPARATOR=SEPARATOR, node=node, query=query, formatter=self) |
0 commit comments