Codebase
Main interface for codemods to interact with codebases, including utility methods etc..
Properties
classes
List of all Classes in the codebase.
Returns a sorted list of all Class nodes in the codebase. Class nodes represent class definitions in source files. Only includes top-level classes, not inner/nested classes.
Returns: list[TClass]: A sorted list of all Class nodes in the codebase.
current_commit
Returns the current Git commit that is checked out in the repository.
Args: None
Returns: GitCommit | None: The currently checked out Git commit object, or None if no commit is checked out.
default_branch
The default branch of this repository.
Returns the name of the default branch (e.g. ‘main’ or ‘master’) for the current repository.
Returns: str: The name of the default branch.
directories
List all directories in the codebase.
Returns a list of all Directory objects present in the codebase. Each Directory object represents a directory in the codebase. This property is used to access and navigate the directory structure of the codebase.
Returns: list[TDirectory]: A list of Directory objects in the codebase.
external_modules
Returns a list of all external modules in the codebase.
An external module represents a dependency that is imported but not defined within the codebase itself (e.g. third-party packages like ‘requests’ or ‘numpy’).
Returns: list[ExternalModule]: List of external module nodes from the codebase graph.
files
A list property that returns all source files in the codebase.
Returns all Files in the codebase, sorted alphabetically. For Python codebases, returns PyFiles (python files). For Typescript codebases, returns TSFiles (typescript files). Does not include non-source files like markdown or image files.
Returns: list[TSourceFile]: A sorted list of source files in the codebase.
functions
List of all Functions in the codebase.
Returns a sorted list of all top-level Function objects in the codebase, excluding class methods.
Returns: list[TFunction]: A list of Function objects representing all functions in the codebase, sorted alphabetically.
global_vars
List of all GlobalVars in the codebase.
A GlobalVar represents a global variable assignment in the source code. These are variables defined at the module level.
Returns: list[TGlobalVar]: A list of all global variable assignments in the codebase.
imports
Returns a list of all Import nodes in the codebase.
Retrieves all Import nodes from the codebase graph. These imports represent all import statements across all files in the codebase, including imports from both internal modules and external packages.
Args: None
Returns: list[TImport]: A list of Import nodes representing all imports in the codebase. TImport can be PyImport for Python codebases or TSImport for TypeScript codebases.
interfaces
Retrieves all interfaces in the codebase.
Returns a list of all Interface symbols defined at the top-level of source files in the codebase. This property is only applicable for TypeScript codebases and will return an empty list for Python codebases.
Returns: list[TInterface]: A list of Interface objects defined in the codebase’s source files.
symbols
List of all top-level Symbols (Classes, Functions, etc.) in the codebase. Excludes Class methods.
Returns: list[TSymbol]: A list of Symbol objects of all top-level symbols in the codebase. Includes classes, functions, and global variables defined at the module level, excludes methods.
types
List of all Types in the codebase (Typescript only).
Returns a list of all type aliases defined at the top level in the codebase. This method is only applicable for TypeScript codebases.
Returns: list[TTypeAlias]: A list of all type aliases defined in the codebase.
Methods
ai
Generates a response from the AI based on the provided prompt, target, and context.
A method that sends a prompt to the AI client along with optional target and context information to generate a response. Used for tasks like code generation, refactoring suggestions, and documentation improvements.
Args: prompt (str): The text prompt to send to the AI. target (Editable | None): An optional editable object (like a function, class, etc.) that provides the main focus for the AI’s response. context (Editable | list[Editable] | dict[str, Editable | list[Editable]] | None): Additional context to help inform the AI’s response. model (str): The AI model to use for generating the response. Defaults to “gpt-4o-mini”.
Returns: str: The generated response from the AI.
Raises: MaxAIRequestsError: If the maximum number of allowed AI requests (default 150) has been exceeded.
checkout
Checks out a git branch or commit and syncs the codebase graph to the new state.
This method discards any pending changes, performs a git checkout of the specified branch or commit, and then syncs the codebase graph to reflect the new state.
Args: commit (str | GitCommit | None): Hash or GitCommit object to checkout. Cannot be used with branch. branch (str | None): Name of branch to checkout. Cannot be used with commit. create_if_missing (bool): If True, creates the branch if it doesn’t exist. Defaults to False. remote (bool): If True, attempts to pull from remote when checking out branch. Defaults to False.
Returns: CheckoutResult: The result of the checkout operation.
Raises: AssertionError: If neither commit nor branch is specified, or if both are specified.
create_directory
Creates a directory at the specified path.
Args: dir_path (str): The path where the directory should be created. exist_ok (bool): If True, don’t raise an error if the directory already exists. Defaults to False. parents (bool): If True, create any necessary parent directories. Defaults to False.
Raises: FileExistsError: If the directory already exists and exist_ok is False.
create_file
Creates a new file in the codebase with specified content.
Args: filepath (str): The path where the file should be created. content (str): The content of the file to be created. Defaults to empty string. sync (bool): Whether to sync the graph after creating the file. Defaults to True.
Returns: File: The newly created file object.
Raises: ValueError: If the provided content cannot be parsed according to the file extension.
find_by_span
Finds editable objects that overlap with the given source code span.
Searches for editable objects (like functions, classes, variables) within a file that overlap with the specified byte range span. Returns an empty list if no matching file is found.
Args: span (Span): The span object containing the filepath and byte range to search within.
Returns: list[Editable]: A list of Editable objects that overlap with the given span.
flag_ai
Determines whether to flag the symbol, file, attribute, or message using AI.
Flagging returns either a CodeFlag object or None. Use this method when you want to conditionally flag an entity based on the AI’s response. By default, there is a maximum of 150 AI requests per codemod execution.
Args: prompt (str): The prompt to send to the AI. target (Editable | None): The target editable object for the AI request. context (Editable | list[Editable] | dict[str, Editable | list[Editable]] | None): Additional context for the AI request. model (str): The model to use for the AI request. Defaults to “gpt-4o-mini”.
Returns: CodeFlag | None: A CodeFlag object if the AI determines the target should be flagged, None otherwise.
Raises: MaxAIRequestsError: If the maximum number of AI requests is exceeded.
flag_instance
Flags a symbol, file or import to enable enhanced tracking of changes and splitting into smaller PRs.
This method should be called once per flaggable entity and should be called before any edits are made to the entity. Flags enable tracking of changes and can be used for various purposes like generating pull requests or applying changes selectively.
Args: symbol (TSymbol | None): The symbol to flag. Can be None if just flagging a message. **kwargs: Arguments used to construct the flag Returns: CodeFlag: A flag object representing the flagged entity.
get_class
Returns a class that matches the given name.
Args: class_name (str): The name of the class to find. optional (bool): If True, return None when class is not found instead of raising ValueError. Defaults to False.
Returns: TClass | None: The class with the given name, or None if optional=True and class not found.
Raises: ValueError: If the class is not found and optional=False, or if multiple classes with the same name exist.
get_directory
Returns Directory by dir_path
, or full path to the directory from codebase root.
Args: dir_path (str): The path to the directory to retrieve. optional (bool): If True, return None when directory is not found. If False, raise ValueError.
Returns: TDirectory | None: The Directory object if found, None if optional=True and directory not found.
Raises: ValueError: If directory not found and optional=False.
get_file
Retrieves a file from the codebase by its filepath.
This method first attempts to find the file in the graph, then checks the filesystem if not found. Files can be either source files (e.g. .py, .ts) or binary files.
Args: filepath (str): The path to the file, relative to the codebase root. optional (bool): If True, return None if file not found. If False, raise ValueError.
Returns: TSourceFile | None: The source file if found, None if optional=True and file not found.
Raises: ValueError: If file not found and optional=False.
get_function
Retrieves a function from the codebase by its name.
This method searches through all functions in the codebase to find one matching the given name. If multiple functions with the same name exist, a ValueError is raised.
Args: function_name (str): The name of the function to retrieve. optional (bool): If True, returns None when function is not found instead of raising ValueError. Defaults to False.
Returns: TFunction | None: The matching function if found. If optional=True and no match is found, returns None.
Raises: ValueError: If function is not found and optional=False, or if multiple matching functions exist.
get_symbol
Returns a Symbol by name from the codebase.
Returns the first Symbol that matches the given name. If multiple symbols are found with the same name, raises a ValueError. If no symbol is found, returns None if optional is True, otherwise raises a ValueError.
Args: symbol_name (str): The name of the symbol to find. optional (bool): If True, returns None when symbol is not found. If False, raises ValueError. Defaults to False.
Returns: TSymbol | None: The matched Symbol if found, None if not found and optional=True.
Raises: ValueError: If multiple symbols are found with the same name, or if no symbol is found and optional=False.
get_symbols
Retrieves all symbols in the codebase that match the given symbol name.
This method is used when there may be multiple symbols with the same name, in which case get_symbol() would raise a ValueError.
Args: symbol_name (str): The name of the symbols to retrieve.
Returns: list[TSymbol]: A list of Symbol objects that match the given name, sorted alphabetically.
Note: When a unique symbol is required, use get_symbol() instead. It will raise ValueError if multiple symbols are found.
git_commit
Commits all staged changes to the codebase and git.
Args: message (str): The commit message verify (bool): Whether to verify the commit before committing. Defaults to False. sync (bool): Whether to sync the codebase after committing. Defaults to True.
Returns: GitCommit | None: The commit object if changes were committed, None otherwise.
has_directory
Returns a boolean indicating if a directory exists in the codebase.
Checks if a directory exists at the specified path within the codebase.
Args: dir_path (str): The path to the directory to check for, relative to the codebase root.
Returns: bool: True if the directory exists in the codebase, False otherwise.
has_file
Determines if a file exists in the codebase.
Args: filepath (str): The filepath to check for existence.
Returns: bool: True if the file exists in the codebase, False otherwise.
has_symbol
Returns whether a symbol exists in the codebase.
This method checks if a symbol with the given name exists in the codebase.
Args: symbol_name (str): The name of the symbol to look for.
Returns: bool: True if a symbol with the given name exists in the codebase, False otherwise.
set_max_ai_requests
Sets the maximum number of AI requests that can be made during codemod execution.
Sets the global limit for AI requests, which is capped by HARD_MAX_AI_LIMIT
.
This limit prevents excessive AI calls during codemod execution. When the limit is reached, a MaxAIRequestsError
is raised.
Args: max_ai_requests (int | None): The maximum number of AI requests allowed. If None, there is no limit.
Raises:
ValueError: If max_ai_requests
exceeds HARD_MAX_AI_LIMIT
.
should_fix
Returns True if the flag should be fixed based on the current mode and active group.
Used to filter out flags that are not in the active group and determine if the flag should be processed or ignored.
Args: flag (CodeFlag): The code flag to check.
Returns: bool: True if the flag should be fixed, False if it should be ignored. Returns False in find mode. Returns True if no active group is set. Returns True if the flag’s hash exists in the active group hashes.
visualize
Visualizes a NetworkX graph or Plotly figure.
Creates a visualization of the provided graph using GraphViz. This is useful for visualizing dependency graphs, call graphs, directory structures, or other graph-based representations of code relationships.
Args: G (Graph | go.Figure): A NetworkX graph or Plotly figure to visualize root (Editable | str | int | None): The root node to visualize around. When specified, the visualization will be centered on this node. Defaults to None.
Returns: None