Tawsif add support for codebase exports (#117)

# Motivation
<!-- Why is this change necessary? -->
- adds exports to codebase.exports and directory.exports + directory.get_export(name) for typescript codebases
- Made a new PR so that I didn't have to deal with merge conflicts from
graphsitter renaming to codegen-sdk
- Fixed all comments from previous PR
- Added better typehints to ensure that exports are meant only for
typescript codebases and typescript specific directories

---------

Co-authored-by: codegen-bot <[email protected]>

Author: tawsifkamal

src/codegen/sdk/core/codebase.py

@@ -39,6 +39,7 @@
 from codegen.sdk.core.detached_symbols.code_block import CodeBlock
 from codegen.sdk.core.detached_symbols.parameter import Parameter
 from codegen.sdk.core.directory import Directory
+from codegen.sdk.core.export import Export
 from codegen.sdk.core.external_module import ExternalModule
 from codegen.sdk.core.file import File, SourceFile
 from codegen.sdk.core.function import Function
@@ -59,19 +60,22 @@
 from codegen.sdk.python.file import PyFile
 from codegen.sdk.python.function import PyFunction
 from codegen.sdk.python.import_resolution import PyImport
+from codegen.sdk.python.statements.import_statement import PyImportStatement
 from codegen.sdk.python.symbol import PySymbol
 from codegen.sdk.typescript.assignment import TSAssignment
 from codegen.sdk.typescript.class_definition import TSClass
 from codegen.sdk.typescript.detached_symbols.code_block import TSCodeBlock
 from codegen.sdk.typescript.detached_symbols.parameter import TSParameter
+from codegen.sdk.typescript.export import TSExport
 from codegen.sdk.typescript.file import TSFile
 from codegen.sdk.typescript.function import TSFunction
 from codegen.sdk.typescript.import_resolution import TSImport
 from codegen.sdk.typescript.interface import TSInterface
+from codegen.sdk.typescript.statements.import_statement import TSImportStatement
 from codegen.sdk.typescript.symbol import TSSymbol
 from codegen.sdk.typescript.type_alias import TSTypeAlias
 from codegen.sdk.utils import determine_project_language
-from codegen.shared.decorators.docs import apidoc, noapidoc
+from codegen.shared.decorators.docs import apidoc, noapidoc, py_noapidoc
 from codegen.shared.exceptions.control_flow import MaxAIRequestsError
 from codegen.shared.performance.stopwatch_utils import stopwatch
 from codegen.visualizations.visualization_manager import VisualizationManager
@@ -91,6 +95,11 @@
 TTypeAlias = TypeVar("TTypeAlias", bound="TypeAlias")
 TParameter = TypeVar("TParameter", bound="Parameter")
 TCodeBlock = TypeVar("TCodeBlock", bound="CodeBlock")
+TExport = TypeVar("TExport", bound="Export")
+TSGlobalVar = TypeVar("TSGlobalVar", bound="Assignment")
+PyGlobalVar = TypeVar("PyGlobalVar", bound="Assignment")
+TSDirectory = Directory[TSFile, TSSymbol, TSImportStatement, TSGlobalVar, TSClass, TSFunction, TSImport]
+PyDirectory = Directory[PyFile, PySymbol, PyImportStatement, PyGlobalVar, PyClass, PyFunction, PyImport]
 
 
 @apidoc
@@ -263,6 +272,27 @@ def imports(self) -> list[TImport]:
         """
         return self.G.get_nodes(NodeType.IMPORT)
 
+    @property
+    @py_noapidoc
+    def exports(self: "TSCodebaseType") -> list[TSExport]:
+        """Returns a list of all Export nodes in the codebase.
+
+        Retrieves all Export nodes from the codebase graph. These exports represent all export statements across all files in the codebase,
+        including exports from both internal modules and external packages. This is a TypeScript-only codebase property.
+
+        Args:
+            None
+
+        Returns:
+            list[TSExport]: A list of Export nodes representing all exports in the codebase.
+            TExport can only be a  TSExport for TypeScript codebases.
+
+        """
+        if self.language == ProgrammingLanguage.PYTHON:
+            raise NotImplementedError("Exports are not supported for Python codebases since Python does not have an export mechanism.")
+
+        return self.G.get_nodes(NodeType.EXPORT)
+
     @property
     def external_modules(self) -> list[ExternalModule]:
         """Returns a list of all external modules in the codebase.
@@ -1145,5 +1175,5 @@ def from_repo(cls, repo_name: str, *, tmp_dir: str | None = None, commit: str |
 # The last 2 lines of code are added to the runner. See codegen-backend/cli/generate/utils.py
 # Type Aliases
 CodebaseType = Codebase[SourceFile, Directory, Symbol, Class, Function, Import, Assignment, Interface, TypeAlias, Parameter, CodeBlock]
-PyCodebaseType = Codebase[PyFile, Directory, PySymbol, PyClass, PyFunction, PyImport, PyAssignment, Interface, TypeAlias, PyParameter, PyCodeBlock]
-TSCodebaseType = Codebase[TSFile, Directory, TSSymbol, TSClass, TSFunction, TSImport, TSAssignment, TSInterface, TSTypeAlias, TSParameter, TSCodeBlock]
+PyCodebaseType = Codebase[PyFile, PyDirectory, PySymbol, PyClass, PyFunction, PyImport, PyAssignment, Interface, TypeAlias, PyParameter, PyCodeBlock]
+TSCodebaseType = Codebase[TSFile, TSDirectory, TSSymbol, TSClass, TSFunction, TSImport, TSAssignment, TSInterface, TSTypeAlias, TSParameter, TSCodeBlock]

src/codegen/sdk/core/directory.py

@@ -3,15 +3,22 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
-from codegen.shared.decorators.docs import apidoc
+from codegen.shared.decorators.docs import apidoc, py_noapidoc
 
 if TYPE_CHECKING:
     from codegen.sdk.core.assignment import Assignment
     from codegen.sdk.core.class_definition import Class
     from codegen.sdk.core.file import File
     from codegen.sdk.core.function import Function
-    from codegen.sdk.core.import_resolution import ImportStatement
+    from codegen.sdk.core.import_resolution import Import, ImportStatement
     from codegen.sdk.core.symbol import Symbol
+    from codegen.sdk.typescript.class_definition import TSClass
+    from codegen.sdk.typescript.export import TSExport
+    from codegen.sdk.typescript.file import TSFile
+    from codegen.sdk.typescript.function import TSFunction
+    from codegen.sdk.typescript.import_resolution import TSImport
+    from codegen.sdk.typescript.statements.import_statement import TSImportStatement
+    from codegen.sdk.typescript.symbol import TSSymbol
 
 import logging
 
@@ -24,10 +31,13 @@
 TGlobalVar = TypeVar("TGlobalVar", bound="Assignment")
 TClass = TypeVar("TClass", bound="Class")
 TFunction = TypeVar("TFunction", bound="Function")
+TImport = TypeVar("TImport", bound="Import")
+
+TSGlobalVar = TypeVar("TSGlobalVar", bound="Assignment")
 
 
 @apidoc
-class Directory(Generic[TFile, TSymbol, TImportStatement, TGlobalVar, TClass, TFunction]):
+class Directory(Generic[TFile, TSymbol, TImportStatement, TGlobalVar, TClass, TFunction, TImport]):
     """Directory representation for codebase.
     GraphSitter abstraction of a file directory that can be used to look for files and symbols within a specific directory.
     """
@@ -133,6 +143,17 @@ def functions(self) -> list[TFunction]:
         """Get a recursive list of all functions in the directory and its subdirectories."""
         return list(chain.from_iterable(f.functions for f in self.files))
 
+    @property
+    @py_noapidoc
+    def exports(self: "Directory[TSFile, TSSymbol, TSImportStatement, TSGlobalVar, TSClass, TSFunction, TSImport]") -> "list[TSExport]":
+        """Get a recursive list of all exports in the directory and its subdirectories."""
+        return list(chain.from_iterable(f.exports for f in self.files))
+
+    @property
+    def imports(self) -> list[TImport]:
+        """Get a recursive list of all imports in the directory and its subdirectories."""
+        return list(chain.from_iterable(f.imports for f in self.files))
+
     def get_symbol(self, name: str) -> TSymbol | None:
         """Get a symbol by name in the directory and its subdirectories."""
         return next((s for s in self.symbols if s.name == name), None)
@@ -176,6 +197,15 @@ def get_file(self, filename: str, ignore_case: bool = False) -> TFile | None:
             return next((f for name, f in self.items.items() if name.lower() == filename.lower() and isinstance(f, File)), None)
         return self.items.get(filename, None)
 
+    @py_noapidoc
+    def get_export(self: "Directory[TSFile, TSSymbol, TSImportStatement, TSGlobalVar, TSClass, TSFunction, TSImport]", name: str) -> "TSExport | None":
+        """Get an export by name in the directory and its subdirectories (supports only typescript)."""
+        return next((s for s in self.exports if s.name == name), None)
+
+    def get_import(self, name: str) -> TImport | None:
+        """Get an import by name in the directory and its subdirectories."""
+        return next((s for s in self.imports if s.name == name), None)
+
     def add_subdirectory(self, subdirectory: Self) -> None:
         """Add a subdirectory to the directory."""
         rel_path = os.path.relpath(subdirectory.dirpath, self.dirpath)

src/codegen/shared/decorators/docs.py

@@ -70,6 +70,20 @@ def noapidoc(obj: T) -> T:
     return obj
 
 
+py_no_apidoc_objects: list[DocumentedObject] = []
+py_no_apidoc_signatures: set[str] = set()
+
+
+def py_noapidoc(obj: T) -> T:
+    """Decorator for things that are hidden from the Python API documentation for AI-agent prompts."""
+    obj._py_apidoc = False
+    obj._api_doc_lang = "python"
+    if doc_obj := get_documented_object(obj):
+        bisect.insort(py_no_apidoc_objects, doc_obj)
+        py_no_apidoc_signatures.add(doc_obj.signature())
+    return obj
+
+
 def get_documented_object(obj) -> DocumentedObject | None:
     module = inspect.getmodule(obj)
     module_name = module.__name__ if module else ""

tests/unit/codegen/sdk/python/codebase/test_codebase_raise_error.py

@@ -0,0 +1,18 @@
+import pytest
+
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.sdk.enums import ProgrammingLanguage
+
+
+def test_python_exports_not_supported(tmpdir):
+    """Test that exports are not supported in Python codebases."""
+    # language=python
+    content = """
+def hello():
+    pass
+    """
+    # Create a Python codebase with a simple Python file
+    with get_codebase_session(tmpdir=tmpdir, files={"test.py": content}, programming_language=ProgrammingLanguage.PYTHON) as codebase:
+        # Verify that accessing exports raises NotImplementedError
+        with pytest.raises(NotImplementedError):
+            _ = codebase.exports

tests/unit/codegen/sdk/typescript/codebase/test_codebase_exports.py

@@ -0,0 +1,48 @@
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.sdk.enums import ProgrammingLanguage
+
+
+def test_codebase_exports(tmpdir) -> None:
+    # language=typescript
+    content = """
+    export const a = 1;
+    export let b = 2;
+    export var c = 3;
+    export function foo() {}
+    export class Bar {}
+    export interface IFoo {}
+    export type MyType = string;
+    export { foo as default };
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 8
+        export_names = {exp.name for exp in codebase.exports}
+        assert export_names == {"a", "b", "c", "foo", "Bar", "IFoo", "MyType", "default"}
+
+
+def test_codebase_reexports(tmpdir) -> None:
+    # language=typescript
+    content1 = """
+    export const x = 1;
+    export const y = 2;
+    """
+    content2 = """
+    export { x } from './file1';
+    export { y as z } from './file1';
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file1.ts": content1, "file2.ts": content2}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 4
+        export_names = {exp.name for exp in codebase.exports}
+        assert export_names == {"x", "y", "z"}
+
+
+def test_codebase_default_exports(tmpdir) -> None:
+    # language=typescript
+    content = """
+    const value = 42;
+    export default value;
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 1
+        export = codebase.exports[0]
+        assert export.name == "value"

tests/unit/codegen/sdk/typescript/export/test_directory_exports.py

@@ -0,0 +1,48 @@
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.sdk.enums import ProgrammingLanguage
+
+
+def test_codebase_exports(tmpdir) -> None:
+    # language=typescript
+    content = """
+    export const a = 1;
+    export let b = 2;
+    export var c = 3;
+    export function foo() {}
+    export class Bar {}
+    export interface IFoo {}
+    export type MyType = string;
+    export { foo as default };
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 8
+        export_names = {exp.name for exp in codebase.exports}
+        assert export_names == {"a", "b", "c", "foo", "Bar", "IFoo", "MyType", "default"}
+
+
+def test_codebase_reexports(tmpdir) -> None:
+    # language=typescript
+    content1 = """
+    export const x = 1;
+    export const y = 2;
+    """
+    content2 = """
+    export { x } from './file1';
+    export { y as z } from './file1';
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file1.ts": content1, "file2.ts": content2}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 4
+        export_names = {exp.name for exp in codebase.exports}
+        assert export_names == {"x", "y", "z"}
+
+
+def test_codebase_default_exports(tmpdir) -> None:
+    # language=typescript
+    content = """
+    const value = 42;
+    export default value;
+    """
+    with get_codebase_session(tmpdir=tmpdir, files={"file.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+        assert len(codebase.exports) == 1
+        export = codebase.exports[0]
+        assert export.name == "value"

tests/unit/codegen/sdk/typescript/import_resolution/test_directory_import.py

@@ -0,0 +1,57 @@
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.sdk.enums import ProgrammingLanguage
+
+
+def test_directory_imports(tmpdir) -> None:
+    # language=typescript
+    content1 = """
+    import { a, b } from '../shared';
+    import type { IFoo } from './types';
+    """
+    content2 = """
+    import { c } from '../shared';
+    import defaultExport from './module';
+    """
+    with get_codebase_session(
+        tmpdir=tmpdir, files={"dir1/file1.ts": content1, "dir1/file2.ts": content2, "dir2/file3.ts": "import { d } from '../shared';"}, programming_language=ProgrammingLanguage.TYPESCRIPT
+    ) as codebase:
+        dir1 = codebase.get_directory("dir1")
+        dir2 = codebase.get_directory("dir2")
+
+        # Test dir1 imports
+        assert len(dir1.imports) == 5
+        dir1_import_names = {imp.name for imp in dir1.imports}
+        assert dir1_import_names == {"a", "b", "IFoo", "c", "defaultExport"}
+
+        # Test dir2 imports
+        assert len(dir2.imports) == 1
+        assert dir2.imports[0].name == "d"
+
+        # Test get_import method
+        assert dir1.get_import("a") is not None
+        assert dir1.get_import("d") is None
+        assert dir2.get_import("d") is not None
+
+
+def test_directory_nested_imports(tmpdir) -> None:
+    # language=typescript
+    content1 = """
+    import { a } from './module1';
+    """
+    content2 = """
+    import { b } from '../module2';
+    """
+    content3 = """
+    import { c } from '../../module3';
+    """
+    with get_codebase_session(
+        tmpdir=tmpdir, files={"dir1/file1.ts": content1, "dir1/subdir/file2.ts": content2, "dir1/subdir/deepdir/file3.ts": content3}, programming_language=ProgrammingLanguage.TYPESCRIPT
+    ) as codebase:
+        dir1 = codebase.get_directory("dir1")
+        subdir = codebase.get_directory("dir1/subdir")
+        deepdir = codebase.get_directory("dir1/subdir/deepdir")
+
+        # Test imports at each directory level
+        assert len(dir1.imports) == 3  # Should include all nested imports
+        assert len(subdir.imports) == 2  # Should include its own and deeper imports
+        assert len(deepdir.imports) == 1  # Should only include its own imports