refactor: Use @graph_sitter

.codegen/codemods/test_language/test_language.py

@@ -1,9 +1,9 @@
-import codegen
-from codegen.sdk.core.codebase import Codebase
-from codegen.shared.enums.programming_language import ProgrammingLanguage
+import graph_sitter
+from graph_sitter.core.codebase import Codebase
+from graph_sitter.shared.enums.programming_language import ProgrammingLanguage
 
 
-@codegen.function("test-language", subdirectories=["src/codegen/cli"], language=ProgrammingLanguage.PYTHON)
+@graph_sitter.function("test-language", subdirectories=["src/codegen/cli"], language=ProgrammingLanguage.PYTHON)
 def run(codebase: Codebase):
     file = codebase.get_file("src/codegen/cli/errors.py")
     print(f"File: {file.path}")

.github/workflows/test.yml

@@ -23,6 +23,10 @@ jobs:
   unit-tests:
     needs: access-check
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        group: [1, 2, 3, 4, 5, 6, 7, 8]
+
     steps:
       - uses: actions/checkout@v4
         with:
@@ -38,9 +42,10 @@ jobs:
           uv run pytest \
             -n auto \
             --cov src \
-            --cov-report=json \
+            --splits 8 --group ${{ matrix.group }} \
+            # --cov-report=json \
             --timeout 15 \
-            -o junit_suite_name="${{github.job}}" \
+            -o junit_suite_name="${{github.job}}-${{ matrix.group }}" \
             tests/unit
 
       - uses: ./.github/actions/report

docs/building-with-graph-sitter/function-decorator.mdx

@@ -1,6 +1,6 @@
 ---
 title: Function Decorator
-sidebarTitle: "@codegen.function"
+sidebarTitle: "@graph_sitter.function"
 icon: "at"
 iconType: "solid"
 ---
@@ -11,17 +11,17 @@ The `function` decorator is used to define codegen functions within your applica
 
 ## Usage
 
-To use the `function` decorator, simply annotate your function with `@codegen.function` and provide a name as an argument.
+To use the `function` decorator, simply annotate your function with `@graph_sitter.function` and provide a name as an argument.
 
 ### Example
 
 ```python
-@codegen.function('my-function')
+@graph_sitter.function('my-function')
 def run(codebase):
     pass
 ```
 
-In this example, the function `run` is decorated with `@codegen.function` and given the name `'my-function'`. This name will be used when the function is ran.
+In this example, the function `run` is decorated with `@graph_sitter.function` and given the name `'my-function'`. This name will be used when the function is ran.
 
 ## Parameters
 

docs/building-with-graph-sitter/reusable-codemods.mdx

@@ -5,7 +5,7 @@ icon: "arrows-rotate"
 iconType: "solid"
 ---
 
-Graph-sitter enables you to create reusable code transformations using Python functions decorated with `@codegen.function`. These codemods can be shared, versioned, and run by your team.
+Graph-sitter enables you to create reusable code transformations using Python functions decorated with `@graph_sitter.function`. These codemods can be shared, versioned, and run by your team.
 
 ## Creating Codemods
 
@@ -18,10 +18,10 @@ gs create rename-function .
 This creates a new codemod in your `.codegen/codemods` directory:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here
@@ -64,15 +64,15 @@ The execution flow:
 
 A codemod consists of three main parts:
 
-1. The `@codegen.function` decorator that names your codemod
+1. The `@graph_sitter.function` decorator that names your codemod
 2. A `run` function that takes a `Codebase` parameter
 3. Your transformation logic using the Codebase API
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("update-imports")
+@graph_sitter.function("update-imports")
 def run(codebase: Codebase):
     """Update import statements to use new package names."""
     for file in codebase.files:
@@ -93,7 +93,7 @@ class RenameArgs(BaseModel):
     old_name: str
     new_name: str
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase, arguments: RenameArgs):
     """Rename a function across the codebase."""
     old_func = codebase.get_function(arguments.old_name)

docs/cli/create.mdx

@@ -41,10 +41,10 @@ When you run `gs create rename-function .`, it creates:
 The generated codemod will have this structure:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here

docs/tutorials/training-data.mdx

@@ -31,7 +31,7 @@ First, we will do a "graph expansion" for each function - grab the function's so
 First, let's import the types we need from Codegen:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from codegen.sdk.core.external_module import ExternalModule
 from codegen.sdk.core.import_resolution import Import

examples/STRUCTURE.md

@@ -25,7 +25,7 @@ Your `run.py` should follow this structure, demonstrated well in the `generate_t
 1. **Imports at the top**
 
    ```python
-   import codegen
+   import graph_sitter
    from graph_sitter import Codebase
    from codegen.sdk.core import Function
    # ... other imports
@@ -42,7 +42,7 @@ Your `run.py` should follow this structure, demonstrated well in the `generate_t
 1. **Main Graph-sitter function with decorator**
 
    ```python
-   @codegen.function("your-function-name")
+   @graph_sitter.function("your-function-name")
    def run(codebase: Codebase):
        """Clear docstring explaining what the function does.
 
@@ -98,7 +98,7 @@ Choose between these approaches based on:
 
 1. **Function Decorator**
 
-   - Always use `@codegen.function()` with a descriptive name
+   - Always use `@graph_sitter.function()` with a descriptive name
    - Name should match the example's purpose
 
 1. **Utility Functions**
@@ -136,7 +136,7 @@ def get_function_context(function) -> dict:
 
 
 # Main transformation with decorator
-@codegen.function("generate-training-data")
+@graph_sitter.function("generate-training-data")
 def run(codebase: Codebase):
     """Generate training data using a node2vec-like approach...
 

examples/examples/codegen-mcp-server/server.py

@@ -167,7 +167,7 @@ def codegen_system_prompt():
 ```python
 from graph_sitter import Codebase
 
-@codegen.function('{name}')
+@graph_sitter.function('{name}')
 def codemod(codebase: Codebase):
     for function in codebase.functions:
         if not function.usages:
@@ -355,9 +355,9 @@ async def run_codemod(
                 return {"error": "Invalid JSON in arguments parameter"}
 
         # Create a session for the codemod
-        from graph_sitter.cli.auth.session import CodegenSession
+        from graph_sitter.cli.auth.session import CliSession
 
-        session = CodegenSession(state.parsed_codebase.repo_path)
+        session = CliSession(state.parsed_codebase.repo_path)
         session.codebase = state.parsed_codebase
 
         # Capture output

examples/examples/cyclomatic_complexity/run.py

@@ -1,12 +1,12 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.statements.for_loop_statement import ForLoopStatement
 from graph_sitter.core.statements.if_block_statement import IfBlockStatement
 from graph_sitter.core.statements.try_catch_statement import TryCatchStatement
 from graph_sitter.core.statements.while_statement import WhileStatement
 
 
-@codegen.function("cyclomatic-complexity")
+@graph_sitter.function("cyclomatic-complexity")
 def run(codebase: Codebase):
     def calculate_cyclomatic_complexity(code_block):
         # Initialize cyclomatic complexity count

examples/examples/delete_dead_code/run.py

@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("delete-dead-code")
+@graph_sitter.function("delete-dead-code")
 def run(codebase: Codebase):
     removed_functions_count = 0
     removed_variables_count = 0

examples/examples/dict_to_schema/run.py

@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("dict-to-pydantic-schema")
+@graph_sitter.function("dict-to-pydantic-schema")
 def run(codebase: Codebase):
     """Convert dictionary literals to Pydantic models in a Python codebase.
 

examples/examples/document_functions/run.py

@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.external_module import ExternalModule
 from graph_sitter.core.import_resolution import Import
@@ -54,7 +54,7 @@ def get_extended_context(symbol: Symbol, degree: int) -> tuple[set[Symbol], set[
     return dependencies, usages
 
 
-@codegen.function("document-functions")
+@graph_sitter.function("document-functions")
 def run(codebase: Codebase):
     # Define the maximum degree of dependencies and usages to consider for context
     N_DEGREE = 2

examples/examples/fragment_to_shorthand/run.py

@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("fragment_to_shorthand")
+@graph_sitter.function("fragment_to_shorthand")
 def run(codebase: Codebase):
     print("🔍 Starting Fragment syntax conversion...")
 

examples/examples/freezegun_to_timemachine_migration/run.py

@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("freezegun-to-timemachine")
+@graph_sitter.function("freezegun-to-timemachine")
 def run(codebase: Codebase):
     """Convert FreezeGun usage to TimeMachine in test files.
 

examples/examples/generate_training_data/run.py

@@ -1,6 +1,6 @@
 import json
 
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.external_module import ExternalModule
 from graph_sitter.core.import_resolution import Import
@@ -42,7 +42,7 @@ def get_function_context(function) -> dict:
     return context
 
 
-@codegen.function("generate-training-data")
+@graph_sitter.function("generate-training-data")
 def run(codebase: Codebase):
     """Generate training data using a node2vec-like approach for code embeddings.
 

examples/examples/modules_dependencies/run.py

@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 
 
-@codegen.function("visualize-modules-dependencies")
+@graph_sitter.function("visualize-modules-dependencies")
 def run(codebase: Codebase):
     # Create a directed graph
     G = nx.DiGraph()

examples/examples/openapi_decorators/run.py

@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
@@ -169,7 +169,7 @@ def analyze_method_params(method) -> dict:
     return schema
 
 
-@codegen.function("add-openapi-decorators")
+@graph_sitter.function("add-openapi-decorators")
 def run(codebase: Codebase):
     """Add OpenAPI decorators (@response and @expect) to API endpoints."""
     analytics = {}

examples/examples/promises_to_async_await/run.py

@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.statements.statement import StatementType
 
 
-@codegen.function("promises-to-async-await")
+@graph_sitter.function("promises-to-async-await")
 def run(codebase: Codebase):
     """Convert a repeated use of a promise chain to async await in the official twilio js client library twilio/twilio-node repository.
 

examples/examples/python2_to_python3/run.py

@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 # Initialize codebase
@@ -115,7 +115,7 @@ def update_iterators(file):
                     stmt.edit(new_stmt)
 
 
-@codegen.function("python2-to-python3")
+@graph_sitter.function("python2-to-python3")
 def run():
     """Main function to run the Python 2 to 3 conversion"""
     print("🚀 Starting Python 2 to 3 conversion...\n")

examples/examples/reexport_management/run.py

@@ -1,11 +1,11 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.typescript.file import TSImport
 
 processed_imports = set()
 
 
-@codegen.function("reexport_management")
+@graph_sitter.function("reexport_management")
 def run(codebase: Codebase):
     print("🚀 Starting reexport analysis...")
     for file in codebase.files:

examples/examples/remove_default_exports/run.py

@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.typescript.file import TSFile
 
 
-@codegen.function("remove-default-exports")
+@graph_sitter.function("remove-default-exports")
 def run(codebase: Codebase):
     """Convert default exports to named exports in TypeScript files.
 

examples/examples/sqlalchemy_1.6_to_2.0/run.py

@@ -1,10 +1,10 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 from graph_sitter.core.expressions.chained_attribute import ChainedAttribute
 
 
-@codegen.function("sqlalchemy-1.4-to-2.0")
+@graph_sitter.function("sqlalchemy-1.4-to-2.0")
 def run(codebase: Codebase):
     """
     Convert SQLAlchemy 1.4 codebases to 2.0.

examples/examples/sqlalchemy_soft_delete/run.py

@@ -2,7 +2,7 @@
 import subprocess
 from pathlib import Path
 
-import codegen
+import graph_sitter
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 
 
@@ -61,7 +61,7 @@ def clone_repo(repo_url: str, repo_path: Path) -> None:
     subprocess.run(["git", "clone", repo_url, str(repo_path)], check=True)
 
 
-@codegen.function("sqlalchemy-soft-delete")
+@graph_sitter.function("sqlalchemy-soft-delete")
 def process_soft_deletes(codebase):
     """Process soft delete conditions for join methods in the codebase."""
     soft_delete_models = {