Skills

Common GraphSitter code snippets

add_copyright_header_skill

Add a copyright header to all files in the codebase

add_decorator_to_foo_functions_skill

This skill adds a specified decorator to all functions within a codebase that start with a particular prefix, ensuring that the decorator is properly imported if not already present.

# get the decorator_function symbol
decorator_symbol = codebase.get_symbol("decorator_function")
# for each file in the codebase
for file in codebase.files:
    # for each function in the file
    for function in file.functions:
        # if the function name starts with 'foo'
        if function.name.startswith("foo"):
            # if the decorator is not imported or declared in the file
            if not file.has_import("decorator_function") and decorator_symbol.file != file:
                # add an import for the decorator function
                file.add_symbol_import(decorator_symbol)
            # add the decorator to the function
            function.add_decorator(f"@{decorator_symbol.name}")

add_decorator_to_function

Adds a specified decorator to all functions and methods in a codebase, ensuring that it is not already present, and handles the necessary imports for the decorator.

# get the decorator symbol
decorator_symbol = codebase.get_symbol("my_decorator")

# iterate through each file
for file in codebase.files:
    # if the file does not have the decorator symbol and the decorator symbol is not in the same file
    if not file.has_import(decorator_symbol.name) and decorator_symbol.file != file:
        # import the decorator symbol
        file.add_symbol_import(decorator_symbol)

    # iterate through each function in the file
    for function in file.functions:
        # if the function is the decorator symbol, skip
        if function == decorator_symbol:
            continue
        # add the decorator to the function, don't add the decorator if it already exists
        function.add_decorator(f"@{decorator_symbol.name}", skip_if_exists=True)
    # iterate through each class in the file
    for cls in file.classes:
        # iterate through each method in the class
        for method in cls.methods:
            # add the decorator to the method, don't add the decorator if it already exists
            method.add_decorator(f"@{decorator_symbol.name}", skip_if_exists=True)

add_docstrings_skill

This skill adds docstrings to all functions and methods in a codebase, ensuring that each function has a descriptive comment explaining its purpose and functionality.

for file in codebase.files:
    for function in file.functions:
        if function.docstring is None:
            docstring = '"""This function does something."""'
            function.set_docstring(docstring)

add_return_type_hint_skill

Adds an integer return type hint to all functions whose names start with ‘foo’.

# for each function in the codebase
for function in codebase.functions:
    # if the function name starts with 'foo'
    if function.name.startswith("foo"):
        # add an int return type hint to the function
        function.return_type.edit("int")

add_type_hints_skill

This skill adds type hints to function parameters and return values in a codebase, enhancing type safety and improving code readability.

add_wrapper_function_skill

Adds a trivial wrapper function around each function and class method, creating a new function that simply calls the original function.

for file in codebase.files:
    for function in file.functions:
        wrapper_name = f"new_{function.name}"
        wrapper_code = f"def {wrapper_name}():\n    return {function.name}()"
        file.add_symbol_from_source(wrapper_code)

    for cls in file.classes:
        for method in cls.methods:
            wrapper_name = f"new_{method.name}"
            wrapper_code = f"def {wrapper_name}(self):\n        return self.{method.name}()"
            cls.add_source(wrapper_code)

append_parameter_skill

Appends a parameter to the signature of a specified function in both Python and TypeScript codebases.

append_to_global_list

Skill to append 2 to global list variable ‘a’.

append_type_to_union_type_skill

Appends an additional ‘str’ type to a piped union type in Python

asyncify_function_skill

Given a synchronous function ‘func_to_convert’, convert its signature to be asynchronous. The function’s call sites as well as the call sites’ functions are recursively converted as well.

asyncify_type_alias_elements

Given a type alias ‘MyMapper’ containing synchronous methods, convert ‘getFromMethod’ method to be asynchronous. The inherited implementations of the type alias as well as all their call sites are also update.

from collections.abc import MutableMapping

FUNC_NAME_TO_CONVERT = "convert"

mapper_symbol: TypeAlias = codebase.get_symbol("MyMapper")
mapper_dict: Dict = mapper_symbol.value
# Update the base type alias definition
mapper_dict[FUNC_NAME_TO_CONVERT].asyncify()

# Collect all implementations of the mapper type
mapper_impl: list[tuple[Symbol, MutableMapping]] = []

for symbol in codebase.symbols:
    if isinstance(symbol, Assignment):
        # Check if the global variable initializes a mapper implementation
        if symbol.type and symbol.type.name == "MyMapper" and isinstance(symbol.value, Dict):
            mapper_impl.append((symbol, symbol.value))
    elif isinstance(symbol, Function):
        if symbol.return_type and symbol.return_type.name == "MyMapper":
            # Search for the assignment statement that implements the mapper type
            for statement in symbol.code_block.assignment_statements:
                if (val := statement.right) and isinstance(val, Dict) and set(val.keys()) == set(mapper_dict.keys()):
                    mapper_impl.append((symbol, val))
                    break
    elif isinstance(symbol, Class):
        if mapper_symbol in symbol.superclasses:
            mapper_impl.append((symbol, {**{x.name: x for x in symbol.methods}}))

# Update the mapper type alias implementations
usages_to_check = []
for symbol, val in mapper_impl:
    func_to_convert = val[FUNC_NAME_TO_CONVERT]
    if not func_to_convert.is_async:
        func_to_convert.asyncify()
    # Collect usages of the type alias implementations
    usages_to_check.extend(symbol.symbol_usages)

files_to_check = set(u.file for u in usages_to_check)
funcs_to_asyncify = []
for file in files_to_check:
    for f_call in file.function_calls:
        if FUNC_NAME_TO_CONVERT in f_call.name:
            if not f_call.is_awaited:
                f_call.edit(f"await ({f_call.source})")
            if parent_func := f_call.parent_function:
                funcs_to_asyncify.append(parent_func)

# Asyncify all functions that are called by the async functions
processed = set()
while funcs_to_asyncify:
    f = funcs_to_asyncify.pop()
    if f in processed:
        continue

    processed.add(f)
    if not f.is_async:
        f.asyncify()

        for call_site in f.call_sites:
            if call_site.parent and isinstance(call_site.parent, Function):
                funcs_to_asyncify.append(call_site.parent)

call_graph_filter

This skill shows a visualization of the call graph from a given function or symbol. It iterates through the usages of the starting function and its subsequent calls, creating a directed graph of function calls. The skill filters out test files and includes only methods with specific names (post, get, patch, delete). By default, the call graph uses red for the starting node, yellow for class methods, and can be customized based on user requests. The graph is limited to a specified depth to manage complexity. In its current form, it ignores recursive calls and external modules but can be modified trivially to include them

call_graph_from_node

This skill creates a directed call graph for a given function. Starting from the specified function, it recursively iterates through its function calls and the functions called by them, building a graph of the call paths to a maximum depth. The root of the directed graph is the starting function, each node represents a function call, and edge from node A to node B indicates that function A calls function B. In its current form, it ignores recursive calls and external modules but can be modified trivially to include them. Furthermore, this skill can easily be adapted to support creating a call graph for a class method. In order to do this one simply needs to replace

function_to_trace = codebase.get_function("function_to_trace")

with

function_to_trace = codebase.get_class("class_of_method_to_trace").get_method("method_to_trace")

call_paths_between_nodes

This skill generates and visualizes a call graph between two specified functions. It starts from a given function and iteratively traverses through its function calls, building a directed graph of the call paths. The skill then identifies all simple paths between the start and end functions, creating a subgraph that includes only the nodes in these paths.

By default, the call graph uses blue for the starting node and red for the ending node, but these colors can be customized based on user preferences. The visualization provides a clear representation of how functions are interconnected, helping developers understand the flow of execution and dependencies between different parts of the codebase.

In its current form, it ignores recursive calls and external modules but can be modified trivially to include them

clear_global_list

Skill to clear global list variable ‘a’.

convert_statement_to_argument

Converts http status code assertion statements into an expect_status argument for test functions that make a call to a http method.

convert_to_built_in_type_skill

Replaces type annotations using typing module with builtin types.

Examples: typing.List -> list typing.Dict -> dict typing.Set -> set typing.Tuple -> tuple

import_replacements = {"List": "list", "Dict": "dict", "Set": "set", "Tuple": "tuple"}

# Iterate over all imports in the codebase
for imported in codebase.imports:
    # Check if the import is from the typing module and is a builtin type
    if imported.module == "typing" and imported.name in import_replacements:
        # Remove the type import
        imported.remove()
        # Iterate over all symbols that use this imported module
        for symbol in imported.symbol_usages:
            # Find all exact matches (Editables) in the symbol with this imported module
            for usage in symbol.find(imported.name, exact=True):
                # Replace the usage with the builtin type
                usage.edit(import_replacements[imported.name])

dead_code

This skill shows a visualization of the dead code in the codebase. It iterates through all functions in the codebase, identifying those that have no usages and are not in test files or decorated. These functions are considered ‘dead code’ and are added to a directed graph. The skill then explores the dependencies of these dead code functions, adding them to the graph as well. This process helps to identify not only directly unused code but also code that might only be used by other dead code (second-order dead code). The resulting visualization provides a clear picture of potentially removable code, helping developers to clean up and optimize their codebase.

delete_rolled_out_feature_flag_skill

Locates a fully rolled out feature flag that has changed from a default value of False to True, and deletes all uses of the flag assuming the flag value is True. This skill simplifies conditional statements and removes unnecessary imports related to the feature flag.

delete_unused_logger_skill

Removes all global variables that are defined as logger instances if they are unused. This skill works for both Python and TypeScript codebases, with slight variations in the logger initialization pattern.

delete_unused_symbols_skill

Deletes all unused symbols in the codebase except for those starting with bar (case insensitive) and deletes all files without any remaining symbols.

dict_to_schema

Converts a dictionary into a Schema object. Converts the key value pairs into arguments for the constructor

# iterate over all global vars
for v in codebase.global_vars:
    # if the variable is a dictionary
    if isinstance(v.value, Dict):
        # convert it to a Schema object
        v.set_value(f"Schema({", ".join(f"{k}={v}" for k, v in v.value.items())})")

eslint_comment_skill

Remove eslint disable comments for an eslint rule. This is useful if a codebase is making an eslint rule either not required or reducing it to warning level. If the rule is no longer required/warning level, the disable comments are also no longer required and can be cleaned up.

ESLINT_RULE = "@typescript-eslint/no-explicit-any"

# Iterate over all files in the codebase
for file in codebase.files:
    # Iterate over all comment statements in the file
    for comment in file.code_block.comments:
        if "eslint-disable" in comment.source:
            pattern = r"(.*eslint-disable(?:-next-line?)?)(?:\s+([@a-z-\/,\s]+))?(.*)"
            match = re.search(pattern, comment.source)
            if not match:
                continue

            rules = [r.strip() for r in match.group(2).split(",")]
            if ESLINT_RULE in rules:
                # Case: the only rule being disabled is the one being de-activated. Delete whole comment
                if len(rules) == 1:
                    print(f"Deleting comment: {comment.source}")
                    comment.remove()
                # Case: multiples rules are being disabled. Remove just ESLINT_RULE from the comment
                else:
                    print(f"Removing {ESLINT_RULE} from comment: {comment.source}")
                    rules.remove(ESLINT_RULE)
                    new_comment_source = f"{match.group(1)} {', '.join(rules)}{match.group(3)}"
                    comment.edit(new_comment_source)

export_skills

Convert default exports to named exports in TypeScript

for file in codebase.files:
    for export in file.exports:
        export.make_non_default()

file_app_import_graph

This skill visualizes the import relationships for a specific file in the codebase. It creates a directed graph where nodes represent the target file and its imported modules. Edges in the graph indicate the import relationships, pointing from the file to its imports. The skill focuses on a particular file (‘path/to/file.py’) and analyzes its import statements to construct the graph. This visualization helps developers understand the dependencies and structure of imports within the specified file, which can be useful for code organization and dependency management.

foreign_key_graph

This skill helps analyze a data schema by creating a graph representation of SQLAlchemy models and their foreign key relationships.

It processes a collection of SQLAlchemy models with foreign keys referencing each other. All of these models are classes that inherit from BaseModel, similar to the one in this file. Foreign keys are typically defined in the following format: agent_run_id = Column(BigInteger, ForeignKey(“AgentRun.id”, ondelete=“CASCADE”), nullable=False)

The skill iterates through all classes in the codebase, identifying those that are subclasses of BaseModel. For each relevant class, it examines the attributes to find ForeignKey definitions. It then builds a mapping of these relationships.

Using this mapping, the skill constructs a directed graph where:

• Nodes represent the models (with the ‘Model’ suffix stripped from their names)
• Edges represent the foreign key relationships between models

This graph visualization allows for easy analysis of the data schema, showing how different models are interconnected through their foreign key relationships. The resulting graph can be used to understand data dependencies, optimize queries, or refactor the database schema.

foreign_key_mapping = {}

# Iterate through all classes in the codebase
for cls in codebase.classes:
    # Check if the class is a subclass of BaseModel and defined in the correct file
    if cls.is_subclass_of("BaseModel") and "from app.models.base import BaseModel" in cls.file.content:
        # Initialize an empty list for the current class
        foreign_key_mapping[cls.name] = []

        # Iterate through the attributes of the class
        for attr in cls.attributes:
            # Check if the attribute's source contains a ForeignKey definition
            if "ForeignKey" in attr.source:
                # Extract the table name from the ForeignKey string
                start_index = attr.source.find('("') + 2
                end_index = attr.source.find(".id", start_index)
                if end_index != -1:
                    target_table = attr.source[start_index:end_index]
                    # Append the target table to the mapping, avoiding duplicates
                    if target_table not in foreign_key_mapping[cls.name]:
                        foreign_key_mapping[cls.name].append(target_table)

# Now foreign_key_mapping contains the desired relationships
# print(foreign_key_mapping)

# Create a directed graph
G = nx.DiGraph()

# Iterate through the foreign_key_mapping to add nodes and edges
for model, targets in foreign_key_mapping.items():
    # Add the model node (strip 'Model' suffix)
    model_name = model.replace("Model", "")
    G.add_node(model_name)

    # Add edges to the target tables
    for target in targets:
        G.add_node(target)  # Ensure the target is also a node
        G.add_edge(model_name, target)

# Now G contains the directed graph of models and their foreign key relationships
# You can visualize or analyze the graph as needed
codebase.visualize(G)

##############################################################################################################
# IN DEGREE
##############################################################################################################

# Calculate in-degrees for each node
in_degrees = G.in_degree()

# Create a list of nodes with their in-degree counts
in_degree_list = [(node, degree) for node, degree in in_degrees]

# Sort the list by in-degree in descending order
sorted_in_degrees = sorted(in_degree_list, key=lambda x: x[1], reverse=True)

# Print the nodes with their in-degrees
for node, degree in sorted_in_degrees:
    print(f"Node: {node}, In-Degree: {degree}")
    if degree == 0:
        G.nodes[node]["color"] = "red"

##############################################################################################################
# FIND MODELS MAPPING TO TASK
##############################################################################################################

# Collect models that map to the Task model
models_mapping_to_task = []
for model, targets in foreign_key_mapping.items():
    if "Task" in targets:
        models_mapping_to_task.append(model)

# Print the models that map to Task
print("Models mapping to 'Task':")
for model in models_mapping_to_task:
    print(f"> {model}")

generate_docstrings

This skill generates docstrings for all class methods. This is another example of Codebase AI, where is it being used to generate textual content that will then be used as the docstring for the class methods.

This is also an example usecase of the context feature, where additional context is provided to the AI to help it generate the docstrings.

mark_internal_functions_skill

This skill identifies functions that are exclusively used within the application directory and marks them as internal by appending an @internal tag to their docstrings.

# for each file in the codebase
for file in codebase.files:
    # skip files that are not in the app directory
    if "app" not in file.filepath.split("/"):
        continue
    # for each function in the file
    for function in file.functions:
        is_internal = True
        # for each usage of the function
        for usage in function.usages:
            # resolve the usage symbol
            usage_symbol = usage.usage_symbol
            # if the usage symbol is not in the app directory
            if "app" not in usage_symbol.filepath.split("/"):
                # the function is not internal
                is_internal = False
                break
        # if the function is internal
        if is_internal:
            # if the function does not have a docstring add one
            if function.docstring is None:
                updated_docstring = "\n@internal\n"
            else:
                # add the @internal tag to the bottom of the docstring
                current_docstring = function.docstring.text or ""
                updated_docstring = current_docstring.strip() + "\n\n@internal\n"
            # update the function docstring
            function.set_docstring(updated_docstring)

move_dataclasses_skills

Moves all classes decorated with @dataclasses into a dedicated directory

# Iterate over all files in the codebase
for file in codebase.files:
    # Check if the file is not a dataclasses file
    if "dataclasses" not in file.filepath and "dataclasses" not in file.name:
        for cls in file.classes:
            # Check if the class is a dataclass
            if "@dataclass" in cls.source:
                # Get a new filename
                # Note: extension is included in the file name
                new_filename = "dataclasses.py"

                # Ensure the file exists
                if not codebase.has_file(new_filename):
                    dst_file = codebase.create_file(new_filename, "")
                else:
                    dst_file = codebase.get_file(new_filename)

                # Move the symbol and it's dependencies, adding a "back edge" import to the original file
                cls.move_to_file(dst_file, include_dependencies=True, strategy="add_back_edge")

move_enums_to_separate_file_skill

Moves any enumerations found within a file into a separate file named enums.py. If the original file contains only enumerations, it renames that file to enums.py. If the enums.py file does not exist, it creates one.

# for each file in the codebase
for file in codebase.files:
    # skip the file if it is already named enums.py
    if file.name == "enums.py":
        continue
    # get all enum classes in the file
    enum_classes = [cls for cls in file.classes if cls.is_subclass_of("Enum")]

    if enum_classes:
        # construct the path for the enums file. Note: extension is added to the filepath
        parent_dir = Path(file.filepath).parent
        new_filepath = str(parent_dir / "enums.py")

        # if the file only contains enums rename it
        if len(file.symbols) == len(enum_classes):
            file.update_filepath(new_filepath)
        else:
            # get the enums file if it exists, otherwise create it
            dst_file = codebase.get_file(new_filepath) if codebase.has_file(new_filepath) else codebase.create_file(new_filepath, "from enum import Enum\n\n")
            # for each enum class in the file
            for enum_class in enum_classes:
                # move the enum class to the enums file
                enum_class.move_to_file(dst_file)

move_foo_functions_skill

Moves all functions starting with ‘foo’ to a file named foo

# get the foo.py file if it exists, otherwise create it Note: extension is included in the file name
foo_file = codebase.get_file("foo.py") if codebase.has_file("foo.py") else codebase.create_file("foo.py")

# for each function in the codebase
for function in codebase.functions:
    # if the function name starts with 'foo'
    if function.name.startswith("foo"):
        # move the function to the foo.py file
        function.move_to_file(foo_file)

move_non_default_exported_jsx_components_skill

Moves all JSX components that are not exported by default into a new file located in the same directory as the original file.

# for each file in the codebase
for file in codebase.files:
    # skip files that do not have default exports
    if not file.default_exports:
        continue
    # list to store non-default exported components
    non_default_exported_components = []
    # get the names of the default exports
    default_exports = [export.name for export in file.default_exports]
    # for each function in the file
    for function in file.functions:
        # if the function is a JSX component and is not a default export
        if function.is_jsx and function.name not in default_exports:
            # add the function to the list of non-default exported components
            non_default_exported_components.append(function)
    # if there are non-default exported components
    if non_default_exported_components:
        for component in non_default_exported_components:
            # create a new file in the same directory as the original file
            component_dir = Path(file.filepath).parent
            # create a new file path for the component
            new_file_path = component_dir / f"{component.name}.tsx"
            # if the file does not exist create it
            new_file = codebase.create_file(str(new_file_path))
            # add an import for React
            new_file.add_import_from_import_string('import React from "react";')
            # move the component to the new file
            component.move_to_file(new_file)

reduce_if_statement_condition_skill

Simplifies the if/else control flow by reducing conditions that are set to a specific value to True. This skill works for both Python and TypeScript codebases, with slight variations in the condition naming.

refactor_class

This skill refactors the given class to be shorter and more readable. It uses codebase to first find the class, then uses codebase AI to refactor the class. This is a trivial case of using Codebase AI to edit and generate new code that will then be applied to the codebase.

remove_from_global_list

Skill to remove 2 from global list variable ‘a’.

remove_unused_imports_skill

Removes all unused import statements from the codebase, ensuring that only necessary imports are retained in each file.

for file in codebase.files:
    for imp in file.imports:
        if len(imp.usages) == 0:
            imp.remove()

rename_class_skill

Renames a specified class in the codebase from an old name to a new name.

old_name = "OldName"
new_name = "NewName"
for file in codebase.files:
    for cls in file.classes:
        if cls.name == old_name:
            cls.rename(new_name)

rename_foo_to_bar_skill

Renames all functions in the codebase that start with ‘foo’ to start with ‘bar’, ensuring consistent naming conventions throughout the code.

rename_global_var_skill

Renames all global variables named ‘x’ to ‘y’ across the codebase.

for file in codebase.files:
    for v in file.global_vars:
        if v.name == "x":
            v.set_name("y")

rename_methods

This skill renames all class methods to something better. This is an example of Codebase AI generating content that is not directly written to the codebase as-in, but is chained and applied to the codebase in a later step.

In this case, the agent is asked to create a new name for each class method, then the new name is used to rename the method in the codebase.

repo_dir_tree

This skill displays the directory or repository tree structure of a codebase. It analyzes the file paths within the codebase and constructs a hierarchical representation of the directory structure. The skill creates a visual graph where each node represents a directory or file, and edges represent the parent-child relationships between directories. This visualization helps developers understand the overall organization of the codebase, making it easier to navigate and manage large projects. Additionally, it can be useful for identifying potential structural issues or inconsistencies in the project layout.

search_type_alias_inheritance_skill

Gets all implementation instances of type alias ‘MyMapper’ in TypeScript codebase.

from collections.abc import MutableMapping

mapper_symbol: TypeAlias = codebase.get_symbol("MyMapper")
mapper_dict: Dict = mapper_symbol.value

# Collect all implementations of the mapper type
mapper_impl: list[tuple[Symbol, MutableMapping]] = []

for symbol in codebase.symbols:
    if isinstance(symbol, Assignment):
        # Check if the global variable initializes a mapper implementation
        if symbol.type and symbol.type.name == "MyMapper" and isinstance(symbol.value, Dict):
            mapper_impl.append((symbol, symbol.value))
    elif isinstance(symbol, Function):
        if symbol.return_type and symbol.return_type.name == "MyMapper":
            # Search for the assignment statement that implements the mapper type
            for statement in symbol.code_block.assignment_statements:
                if (val := statement.right) and isinstance(val, Dict) and set(val.keys()) == set(mapper_dict.keys()):
                    mapper_impl.append((symbol, val))
                    break
    elif isinstance(symbol, Class):
        if mapper_symbol in symbol.superclasses:
            mapper_impl.append((symbol, {**{x.name: x for x in symbol.methods}}))

assert len(mapper_impl) == 3
assert all([set(val.keys()) == set(mapper_dict.keys()) for _, val in mapper_impl])

set_global_var_value_skill

This skill modifies the values of all global variables in the codebase, setting them to 2 if their current assigned value is 1.

for file in codebase.files:
    for v in file.global_vars:
        if v.value == "1":
            v.set_value("2")

skip_all_tests

This skill adds a decorator to all test functions in the codebase, marking them to be skipped during test execution with a specified reason.

for file in codebase.files:
    for function in file.functions:
        if function.name.startswith("test_"):
            file.add_import_from_import_string("import pytest")
            function.add_decorator('@pytest.mark.skip(reason="This is a test")')

    for cls in file.classes:
        for method in cls.methods:
            if method.name.startswith("test_"):
                file.add_import_from_import_string("import pytest")
                method.add_decorator('@pytest.mark.skip(reason="This is a test")')

unwrap_function_body

Unwraps the body of all functions in the codebase, transforming each function’s code block into a flat structure without nested scopes.

unwrap_if_statement

Unwraps the body of all if statements in the file

unwrap_with_statement

This unwraps a with statement

# for all functions in the codebase
for function in codebase.functions:
    # for each with statement in the function
    for with_statement in function.code_block.with_statements:
        # unwrap the with statement
        with_statement.code_block.unwrap()

update_doc_string_of_decorated_methods

Updates the docstring of methods whose decorator has with_user/withUser in their name by appending ‘OPERATES ON USER DATA’.

update_optional_type_hints_skill

This skill replaces type hints in functions and methods by updating instances of Optional[type] to type | None, ensuring compatibility with modern type hinting practices.

# pattern to match Optional[type]
optional_type_pattern = re.compile(r"Optional\[(.*?)]")

# update optional parameter type hints
def update_optional_parameter_type_hints(function: PyFunction):
    # for each parameter in the function
    for parameter in function.parameters:
        # if the parameter is typed
        if parameter.is_typed:
            # get the old type
            old_type = parameter.type
            # if the old type is Optional[type]
            if "Optional[" in old_type:
                # replace Optional[type] with type | None
                new_type = optional_type_pattern.sub(r"\1 | None", old_type)
                # update the parameter type hint
                parameter.set_type_annotation(new_type)

def update_optional_return_type_hints(function: PyFunction):
    # if the function has a return type
    if function.return_type:
        # get the old return type
        old_return_type = function.return_type.source
        # if the old return type is Optional[type]
        if "Optional[" in old_return_type:
            # replace Optional[type] with type | None
            new_return_type = optional_type_pattern.sub(r"\1 | None", old_return_type)
            # update the return type hint
            function.return_type.edit(new_return_type)

# for each function in the codebase
for function in codebase.functions:
    # update optional parameter type hints
    update_optional_parameter_type_hints(function)
    # update optional return type hints
    update_optional_return_type_hints(function)

# for each class in the codebase
for cls in codebase.classes:
    # for each method in the class
    for method in cls.methods:
        # update optional parameter type hints
        update_optional_parameter_type_hints(method)
        # update optional return type hints
        update_optional_return_type_hints(method)

write_test

This skill writes a test for the given function. This is an example of Codebase AI generating brand new code that will then be added to the codebase.