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
 
 
 async def ingest_async(
@@ -48,6 +49,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
@@ -112,7 +115,8 @@ async def ingest_async(
     async with _clone_repo_if_remote(query, token=token):
         if not include_gitignored:
             _apply_gitignores(query)
-        summary, tree, content = ingest_query(query)
+        context = ingest_query(query)
+        summary, tree, content = context.generate_digest()
         await _write_output(tree, content=content, target=output)
         return summary, tree, content
 
@@ -136,6 +140,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
@@ -175,20 +181,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,19 +7,18 @@
 
 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
 
 if TYPE_CHECKING:
     from gitingest.schemas import IngestionQuery
 
 
-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
     ----------
@@ -28,8 +27,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
     ------
@@ -65,7 +64,7 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
             msg = f"File {file_node.name} has no content"
             raise ValueError(msg)
 
-        return format_node(file_node, query=query)
+        return Context([file_node])
 
     root_node = FileSystemNode(
         name=path.name,
@@ -78,7 +77,7 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
 
     _process_node(node=root_node, query=query, stats=stats)
 
-    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

@@ -10,6 +10,7 @@
 from gitingest.query_parser import parse_remote_repo
 from gitingest.utils.git_utils import validate_github_token
 from gitingest.utils.pattern_utils import process_patterns
+from gitingest.schemas import Context
 from server.models import IngestErrorResponse, IngestResponse, IngestSuccessResponse, PatternType
 from server.s3_utils import generate_s3_file_path, is_s3_enabled, upload_to_s3
 from server.server_config import MAX_DISPLAY_SIZE
@@ -82,7 +83,8 @@ 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