wip

Author: NicolasIRAGNE

src/gitingest/entrypoint.py

@@ -15,7 +15,7 @@
 from gitingest.clone import clone_repo
 from gitingest.config import MAX_FILE_SIZE
 from gitingest.ingestion import ingest_query
-from gitingest.output_formatter import generate_digest
+from gitingest.output_formatter import DefaultFormatter
 from gitingest.query_parser import parse_local_dir_path, parse_remote_repo
 from gitingest.utils.auth import resolve_token
 from gitingest.utils.compat_func import removesuffix
@@ -51,7 +51,7 @@ async def ingest_async(
     This function analyzes a source (URL or local path), clones the corresponding repository (if applicable),
     and processes its files according to the specified query parameters. It returns a single digest string.
 
-    The output is generated lazily using a ContextV1 object and the generate_digest() function.
+    The output is generated lazily using a ContextV1 object and the DefaultFormatter class.
 
     Parameters
     ----------
@@ -141,7 +141,8 @@ async def ingest_async(
         if output:
             logger.debug("Writing output to file", extra={"output_path": output})
         context = ingest_query(query)
-        digest = generate_digest(context)
+        formatter = DefaultFormatter()
+        digest = formatter.format(context, context.query)
         await _write_output(digest, content=None, target=output)
         logger.info("Ingestion completed successfully")
         return digest
@@ -165,7 +166,7 @@ def ingest(
     This function analyzes a source (URL or local path), clones the corresponding repository (if applicable),
     and processes its files according to the specified query parameters. It returns a single digest string.
 
-    The output is generated lazily using a ContextV1 object and the generate_digest() function.
+    The output is generated lazily using a ContextV1 object and the DefaultFormatter class.
 
     Parameters
     ----------

src/gitingest/format/DefaultFormatter/ContextV1.j2

@@ -1,11 +1,11 @@
-# Generated using https://gitingest.com/{{ context.query.user_name }}/{{ context.query.repo_name }}{{ context.query.subpath }}
+# Generated using https://gitingest.com/{{ source.query.user_name }}/{{ source.query.repo_name }}{{ source.query.subpath }}
 
 Sources used:
-{%- for source in context %}
-- {{ source.name }}: {{ source.__class__.__name__ }}
+{%- for src in source %}
+- {{ src.name }}: {{ src.__class__.__name__ }}
 {% endfor %}
 
-{%- for source in context.sources %}
-{{ formatter.format(source, context.query) }}
+{%- for src in source.sources %}
+{{ formatter.format(src, source.query) }}
 {%- endfor %}
-# End of https://gitingest.com/{{ context.query.user_name }}/{{ context.query.repo_name }}{{ context.query.subpath }}
+# End of https://gitingest.com/{{ source.query.user_name }}/{{ source.query.repo_name }}{{ source.query.subpath }}

src/gitingest/format/DefaultFormatter/FileSystemDirectory.j2

@@ -1,7 +1,7 @@
-{%- if node.depth == 0 %}{{ node.name }}:
-{{ node.tree }}
+{%- if source.depth == 0 %}{{ source.name }}:
+{{ source.tree }}
 
 {% endif -%}
-{%- for child in node.children -%}
+{%- for child in source.children -%}
 {{ formatter.format(child, query) }}
 {%- endfor -%}

src/gitingest/format/DefaultFormatter/FileSystemFile.j2

@@ -1,4 +1,4 @@
 {{ SEPARATOR }}
-{{ node.name }}
+{{ source.name }}
 {{ SEPARATOR }}
-{{ node.content }}
+{{ source.content }}

src/gitingest/format/DefaultFormatter/FileSystemSymlink.j2

@@ -1,3 +1,3 @@
 {{ SEPARATOR }}
-{{ node.name }}{% if node.target %} -> {{ node.target }}{% endif %}
+{{ source.name }}{% if source.target %} -> {{ source.target }}{% endif %}
 {{ SEPARATOR }}

src/gitingest/format/DefaultFormatter/GitRepository.j2

@@ -1,7 +1,7 @@
-{%- if node.depth == 0 %}🔗 Git Repository: {{ node.name }}
-{{ node.tree }}
+{%- if source.depth == 0 %}🔗 Git Repository: {{ source.name }}
+{{ source.tree }}
 
 {% endif -%}
-{%- for child in node.children -%}
+{%- for child in source.children -%}
 {{ formatter.format(child, query) }}
 {%- endfor -%}

src/gitingest/format/SummaryFormatter/ContextV1.j2

@@ -1,3 +1,5 @@
-Repository: {{ context.query.user_name }}/{{ context.query.repo_name }}
-Commit: {{ context.query.commit }}
-Files analyzed: {{ context.file_count }}
+Repository: {{ source.query.user_name }}/{{ source.query.repo_name }}
+Commit: {{ source.query.commit }}
+Files analyzed: {{ source.file_count }}
+
+Estimated tokens: {{ source.token_count }}

src/gitingest/format/SummaryFormatter/FileSystemDirectory.j2

@@ -1,2 +1,2 @@
 Directory structure:
-{{ node.tree }}
+{{ source.tree }}

src/gitingest/ingestion.py

@@ -40,7 +40,7 @@ def ingest_query(query: IngestionQuery) -> ContextV1:
     -------
     ContextV1
         A ContextV1 object representing the ingested file system nodes.
-        Use generate_digest(context) to get the summary, directory structure,
+        Use str(DefaultFormatter(context)) to get the summary, directory structure,
         and file contents.
 
     Raises
@@ -96,7 +96,7 @@ def ingest_query(query: IngestionQuery) -> ContextV1:
                 "file_size": file_node.size,
             },
         )
-        return ContextV1([file_node], query)
+        return ContextV1(sources=[file_node], query=query)
 
     # Check if this is a git repository and create appropriate node type
     if _is_git_repository(path):
@@ -127,7 +127,7 @@ def ingest_query(query: IngestionQuery) -> ContextV1:
         },
     )
 
-    return ContextV1([root_node], query)
+    return ContextV1(sources=[root_node], query=query)
 
 
 def _process_node(node: FileSystemNode, query: IngestionQuery, stats: FileSystemStats) -> None:  # noqa: C901

src/gitingest/output_formatter.py

@@ -11,7 +11,7 @@
 import tiktoken
 from jinja2 import Environment, FileSystemLoader, Template, TemplateNotFound
 
-from gitingest.schemas import ContextV1, FileSystemNode, Source
+from gitingest.schemas import FileSystemNode, Source
 from gitingest.schemas.filesystem import SEPARATOR, FileSystemNodeType
 from gitingest.utils.compat_func import readlink
 from gitingest.utils.logging_config import get_logger
@@ -136,15 +136,15 @@ def _format_token_count(text: str) -> str | None:
     return str(total_tokens)
 
 
-def generate_digest(context: ContextV1) -> str:
-    """Generate a digest string from a ContextV1 object.
+def generate_digest(context: Source) -> str:
+    """Generate a digest string from a Source object.
 
-    This is a convenience function that uses the DefaultFormatter to format a ContextV1.
+    This is a convenience function that uses the DefaultFormatter to format a Source.
 
     Parameters
     ----------
-    context : ContextV1
-        The ContextV1 object containing sources and query information.
+    context : Source
+        The Source object containing sources and query information.
 
     Returns
     -------
@@ -156,57 +156,103 @@ def generate_digest(context: ContextV1) -> str:
     return formatter.format(context, context.query)
 
 
-class DefaultFormatter:
-    """Default formatter for rendering filesystem nodes using Jinja2 templates."""
+class Formatter:
+    """Base formatter class."""
 
-    def __init__(self) -> None:
+    def __init__(self, template_subdir: str) -> None:
         self.separator = SEPARATOR
-        template_dir = Path(__file__).parent / "format" / "DefaultFormatter"
+        template_dir = Path(__file__).parent / "format" / template_subdir
         self.env = Environment(loader=FileSystemLoader(template_dir), autoescape=True)
 
     def _get_template_for_node(self, node: Source) -> Template:
         """Get template based on node class name."""
         template_name = f"{node.__class__.__name__}.j2"
         return self.env.get_template(template_name)
 
+
+class DefaultFormatter(Formatter):
+    """Default formatter for rendering filesystem nodes using Jinja2 templates."""
+
+    def __init__(self) -> None:
+        super().__init__("DefaultFormatter")
+
+    def format(self, source: Source, query: IngestionQuery) -> str:
+        """Format a source with the given query."""
+        if query is None:
+            # Handle case where query is None (shouldn't happen in normal usage)
+            raise ValueError("ContextV1 must have a valid query object")
+        
+        # Calculate and set token count for ContextV1
+        if hasattr(source, '_token_count'):
+            token_count = self._calculate_token_count(source)
+            source._token_count = token_count
+            # Also set token count in the extra dict
+            source.extra["token_count"] = token_count
+        
+        try:
+            return self._format_node(source, query)
+        except Exception as e:
+            # Log the error for debugging
+            import logging
+            logging.error(f"Error in DefaultFormatter: {e}")
+            raise
+
+    def _calculate_token_count(self, source: Source) -> str:
+        """Calculate token count for the entire source."""
+        # Gather all content from the source
+        content = self._gather_all_content(source)
+        return _format_token_count(content) or "Unknown"
+
+    def _gather_all_content(self, node: Source) -> str:
+        """Recursively gather all content from the source tree."""
+        content_parts = []
+        
+        # Add content from the current node
+        if hasattr(node, 'content'):
+            content_parts.append(node.content)
+        
+        # Add content from all sources if it's a ContextV1
+        if hasattr(node, 'sources'):
+            for source in node.sources:
+                content_parts.append(self._gather_all_content(source))
+        
+        # Add content from children if it's a directory
+        if hasattr(node, 'children'):
+            for child in node.children:
+                content_parts.append(self._gather_all_content(child))
+        
+        return "\n".join(filter(None, content_parts))
+
     @singledispatchmethod
-    def format(self, node: Source, query: IngestionQuery) -> str:
+    def _format_node(self, node: Source, query: IngestionQuery) -> str:
         """Dynamically format any node type based on available templates."""
         try:
             template = self._get_template_for_node(node)
             # Provide common template variables
             context_vars = {
-                "node": node,
+                "source": node,
                 "query": query,
                 "formatter": self,
                 "SEPARATOR": SEPARATOR,
             }
-            # Special handling for ContextV1 objects
-            if isinstance(node, ContextV1):
-                context_vars["context"] = node
-                # Use ContextV1 for backward compatibility
-                template = self.env.get_template("ContextV1.j2")
 
             return template.render(**context_vars)
         except TemplateNotFound:
             # Fallback: return content if available, otherwise empty string
             return f"{getattr(node, 'content', '')}"
 
 
-class DebugFormatter:
+class DebugFormatter(Formatter):
     """Debug formatter that shows detailed information about filesystem nodes."""
 
     def __init__(self) -> None:
-        self.separator = SEPARATOR
-        template_dir = Path(__file__).parent / "format" / "DebugFormatter"
-        self.env = Environment(loader=FileSystemLoader(template_dir), autoescape=True)
+        super().__init__("DebugFormatter")
 
     def _get_template_for_node(self, node: Source) -> Template:
         """Get template based on node class name."""
         template_name = f"{node.__class__.__name__}.j2"
         return self.env.get_template(template_name)
 
-    @singledispatchmethod
     def format(self, node: Source, query: IngestionQuery) -> str:
         """Dynamically format any node type with debug information."""
         try:
@@ -254,37 +300,31 @@ def _raise_no_dataclass_fields() -> None:
             return f"DEBUG: {node.__class__.__name__}"
 
 
-class SummaryFormatter:
+class SummaryFormatter(Formatter):
     """Dedicated formatter for generating summaries of filesystem nodes."""
 
     def __init__(self) -> None:
-        template_dir = Path(__file__).parent / "format" / "SummaryFormatter"
-        self.env = Environment(loader=FileSystemLoader(template_dir), autoescape=True)
+        super().__init__("SummaryFormatter")
 
-    def _get_template_for_node(self, node: Source) -> Template:
-        """Get template based on node class name."""
-        template_name = f"{node.__class__.__name__}.j2"
-        return self.env.get_template(template_name)
+    def format(self, source: Source, query: IngestionQuery) -> str:
+        """Generate the summary output."""
+        if query is None:
+            # Handle case where query is None (shouldn't happen in normal usage)
+            raise ValueError("ContextV1 must have a valid query object")
+        return self.summary(source, query)
 
     @singledispatchmethod
     def summary(self, node: Source, query: IngestionQuery) -> str:
         """Dynamically generate summary for any node type based on available templates."""
         try:
             # Provide common template variables
             context_vars = {
-                "node": node,
+                "source": node,
                 "query": query,
                 "formatter": self,
             }
 
-            # Special handling for ContextV1 objects
-            if isinstance(node, ContextV1):
-                context_vars["context"] = node
-                # Use ContextV1 for backward compatibility
-                template = self.env.get_template("ContextV1.j2")
-            else:
-                template = self._get_template_for_node(node)
-
+            template = self._get_template_for_node(node)
             return template.render(**context_vars)
         except TemplateNotFound:
             # Fallback: return name if available