test

Author: NicolasIRAGNE

src/gitingest/entrypoint.py

@@ -27,6 +27,7 @@
     from types import TracebackType
 
     from gitingest.schemas import IngestionQuery
+    from gitingest.schemas import Context
 
 # Initialize logger for this module
 logger = get_logger(__name__)
@@ -51,6 +52,8 @@ async def ingest_async(
     and processes its files according to the specified query parameters. It returns a summary, a tree-like
     structure of the files, and the content of the files. The results can optionally be written to an output file.
 
+    The output is generated lazily using a Context object and its .generate_digest() method.
+
     Parameters
     ----------
     source : str
@@ -142,6 +145,8 @@ async def ingest_async(
 
         if output:
             logger.debug("Writing output to file", extra={"output_path": output})
+        context = ingest_query(query)
+        summary, tree, content = context.generate_digest()
         await _write_output(tree, content=content, target=output)
 
         logger.info("Ingestion completed successfully")
@@ -167,6 +172,8 @@ def ingest(
     and processes its files according to the specified query parameters. It returns a summary, a tree-like
     structure of the files, and the content of the files. The results can optionally be written to an output file.
 
+    The output is generated lazily using a Context object and its .generate_digest() method.
+
     Parameters
     ----------
     source : str
@@ -206,20 +213,20 @@ def ingest(
     ``ingest_async`` : The asynchronous version of this function.
 
     """
-    return asyncio.run(
-        ingest_async(
-            source=source,
-            max_file_size=max_file_size,
-            include_patterns=include_patterns,
-            exclude_patterns=exclude_patterns,
-            branch=branch,
-            tag=tag,
-            include_gitignored=include_gitignored,
-            include_submodules=include_submodules,
-            token=token,
-            output=output,
-        ),
-    )
+    import asyncio
+    context = asyncio.run(ingest_async(
+        source,
+        max_file_size=max_file_size,
+        include_patterns=include_patterns,
+        exclude_patterns=exclude_patterns,
+        branch=branch,
+        tag=tag,
+        include_gitignored=include_gitignored,
+        include_submodules=include_submodules,
+        token=token,
+        output=output,
+    ))
+    return context.generate_digest()
 
 
 def _override_branch_and_tag(query: IngestionQuery, branch: str | None, tag: str | None) -> None:

src/gitingest/ingestion.py

@@ -7,7 +7,7 @@
 
 from gitingest.config import MAX_DIRECTORY_DEPTH, MAX_FILES, MAX_TOTAL_SIZE_BYTES
 from gitingest.output_formatter import format_node
-from gitingest.schemas import FileSystemNode, FileSystemNodeType, FileSystemStats
+from gitingest.schemas import FileSystemNode, FileSystemNodeType, FileSystemStats, Context
 from gitingest.utils.ingestion_utils import _should_exclude, _should_include
 from gitingest.utils.logging_config import get_logger
 
@@ -18,12 +18,11 @@
 logger = get_logger(__name__)
 
 
-def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
+def ingest_query(query: IngestionQuery) -> Context:
     """Run the ingestion process for a parsed query.
 
     This is the main entry point for analyzing a codebase directory or single file. It processes the query
-    parameters, reads the file or directory content, and generates a summary, directory structure, and file content,
-    along with token estimations.
+    parameters, reads the file or directory content, and returns a Context object that can generate the final output digest on demand.
 
     Parameters
     ----------
@@ -32,8 +31,8 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
 
     Returns
     -------
-    tuple[str, str, str]
-        A tuple containing the summary, directory structure, and file contents.
+    Context
+        A Context object representing the ingested file system nodes. Call .generate_digest() to get the summary, directory structure, and file contents.
 
     Raises
     ------
@@ -91,7 +90,7 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
                 "file_size": file_node.size,
             },
         )
-        return format_node(file_node, query=query)
+		return Context([file_node])
 
     logger.info("Processing directory", extra={"directory_path": str(path)})
 
@@ -117,7 +116,7 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
         },
     )
 
-    return format_node(root_node, query=query)
+    return Context([root_node])
 
 
 def _process_node(node: FileSystemNode, query: IngestionQuery, stats: FileSystemStats) -> None:

src/gitingest/schemas/__init__.py

@@ -1,7 +1,7 @@
 """Module containing the schemas for the Gitingest package."""
 
 from gitingest.schemas.cloning import CloneConfig
-from gitingest.schemas.filesystem import FileSystemNode, FileSystemNodeType, FileSystemStats
+from gitingest.schemas.filesystem import Context, FileSystemNode, FileSystemNodeType, FileSystemStats
 from gitingest.schemas.ingestion import IngestionQuery
 
-__all__ = ["CloneConfig", "FileSystemNode", "FileSystemNodeType", "FileSystemStats", "IngestionQuery"]
+__all__ = ["CloneConfig", "FileSystemNode", "FileSystemNodeType", "FileSystemStats", "IngestionQuery", "Context"]

src/gitingest/schemas/filesystem.py

@@ -159,3 +159,61 @@ def content(self) -> str:  # pylint: disable=too-many-return-statements
                 return fp.read()
         except (OSError, UnicodeDecodeError) as exc:
             return f"Error reading file with {good_enc!r}: {exc}"
+
+
+@dataclass
+class Context:
+    """Context for holding a list of FileSystemNode objects and generating a digest on demand."""
+    nodes: list[FileSystemNode]
+
+    def generate_digest(self) -> tuple[str, str, str]:
+        """Generate a summary, directory structure, and file contents for the context's nodes.
+
+        Returns
+        -------
+        tuple[str, str, str]
+            A tuple containing the summary, directory structure, and file contents.
+        """
+        summary_lines = ["Context Digest"]
+        total_files = 0
+        for node in self.nodes:
+            if node.type == FileSystemNodeType.DIRECTORY:
+                total_files += node.file_count
+            elif node.type == FileSystemNodeType.FILE:
+                total_files += 1
+        summary_lines.append(f"Files analyzed: {total_files}")
+        summary = "\n".join(summary_lines)
+
+        # Directory structure
+        tree_lines = ["Directory structure:"]
+        for node in self.nodes:
+            tree_lines.append(self._create_tree_structure(node))
+        tree = "\n".join(tree_lines)
+
+        # File contents
+        content_lines = []
+        for node in self.nodes:
+            content_lines.append(self._gather_file_contents(node))
+        content = "\n".join(content_lines)
+
+        return summary, tree, content
+
+    def _gather_file_contents(self, node: FileSystemNode) -> str:
+        if node.type != FileSystemNodeType.DIRECTORY:
+            return node.content_string
+        return "\n".join(self._gather_file_contents(child) for child in node.children)
+
+    def _create_tree_structure(self, node: FileSystemNode, prefix: str = "", is_last: bool = True) -> str:
+        tree_str = ""
+        current_prefix = "└── " if is_last else "├── "
+        display_name = node.name
+        if node.type == FileSystemNodeType.DIRECTORY:
+            display_name += "/"
+        elif node.type == FileSystemNodeType.SYMLINK:
+            display_name += " -> " + readlink(node.path).name
+        tree_str += f"{prefix}{current_prefix}{display_name}\n"
+        if node.type == FileSystemNodeType.DIRECTORY and node.children:
+            prefix += "    " if is_last else "│   "
+            for i, child in enumerate(node.children):
+                tree_str += self._create_tree_structure(child, prefix=prefix, is_last=i == len(node.children) - 1)
+        return tree_str

src/server/query_processor.py

@@ -21,6 +21,7 @@
     upload_metadata_to_s3,
     upload_to_s3,
 )
+from gitingest.schemas import Context
 from server.server_config import MAX_DISPLAY_SIZE
 
 # Initialize logger for this module
@@ -287,7 +288,10 @@ async def process_query(
         raise RuntimeError(msg)
 
     try:
-        summary, tree, content = ingest_query(query)
+        context = ingest_query(query)
+        summary, tree, content = context.generate_digest()
+
+        # Prepare the digest content (tree + content)
         digest_content = tree + "\n" + content
         _store_digest_content(query, clone_config, digest_content, summary, tree, content)
     except Exception as exc: