Skills

Common GraphSitter code snippets

activity_definition_constant_extractor

Finds all @activity.defn(...) decorators on functions and If they have a name arg, it pulls out this value into a constants.py file in the same directory, then imports it and replaces the arg to point to this constant.

for file in codebase.files:
    # Performance optimization - only look at files with @activity
    if "@activity" not in file.content:
        continue

    # Iterate through all functions in the file
    for function in file.functions:
        # Iterate through decorators
        for decorator in function.decorators:
            # Find @activity decorators
            if "@activity.defn" in decorator.source:
                # Check to see if they have a function call
                if len(decorator.function_calls) > 0:
                    # Grab the `name` arg from this function call
                    call = decorator.function_calls[0]
                    name_arg = call.get_arg_by_parameter_name("name")
                    if name_arg:
                        # Check if the name_arg is a string (use GraphSitter `String` type, not python `str`)
                        if isinstance(name_arg.value, String):
                            # Extract the last part of the string after the dot
                            constant_name = str(name_arg.value).split(".")[-1].upper().replace("-", "_").replace(" ", "_")

                            # Create a constants.py file in the same directory
                            constants_file_path = f"{file.directory.path}/constants.py"
                            if not codebase.has_file(constants_file_path):
                                constants_file = codebase.create_file(constants_file_path, f"{constant_name} = '{name_arg.value}'\n")
                                print(f"🔵 Creating file: {constants_file}")
                            else:
                                constants_file = codebase.get_file(constants_file_path)

                            # Add the string to this file
                            if constant_name not in constants_file.source:
                                constants_file.add_symbol_from_source(f"{constant_name} = '{name_arg.value}'")
                                print(f"  🔷 Constant: {constant_name} ({name_arg.value})")

                            # Import the constant into the current file
                            if not file.has_import(f"{constant_name}"):
                                file.add_import_from_import_string(f"from .constants import {constant_name}")

                            # Set the name_arg to the constant
                            name_arg.set_value(constant_name)

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_self_parameter_to_skill_impl_methods

This codemod finds all methods that have the @skill_impl decorator on them and adds a first self parameter to the method.

for cls in codebase.classes:
    # Check each method in the class
    for method in cls.methods:
        # Check if the method has the @skill_impl decorator
        if any(["@skill_impl" in d.source for d in method.decorators]):
            # Insert 'self' as the first parameter
            method.parameters.insert(0, "self")

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’.

codebase_path_refactorer

This codemod (code modification tool) script is designed to automate a mass refactor in a whole codebase by replacing instances of a specific string. It examines all files in the codebase, checking for any occurrences of the string “/app”. It then scans through all the symbols inside each file, excluding import symbols. It next looks for appearances of “/app” in the symbol’s source using a precise search. If it finds any occurrences, they are replaced with the string “paths.app”. In doing this, it simplifies the process of changing a directory path across a large number of files. The key features of this script are the ability to iterate over files and symbols within the files, and its ability to search and replace specific instances of a string.

for file in codebase.files:
    # Check if the file contains the string "/app"
    if "/app" in file.content:
        # Iterate through all symbols in the file
        for symbol in file.symbols:
            # Exclude import symbols
            if not isinstance(symbol, Import):
                # Search for occurrences of "/app" in the symbol's source
                matches = symbol.search("/app", include_strings=True)
                for match in matches:
                    # Replace occurrences of "/app" with "paths.app"
                    match.replace("/app", "/paths.app")

convert_jsx_arrow_functions_to_named

This converts all JSX components that are arrow functions to named components.

for function in codebase.functions:
    # Check if the function is a JSX function and an arrow function
    if function.is_jsx and function.is_arrow:
        # Convert this function to a named function
        function.arrow_to_named()

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])

count_use_intl_functions_bar_chart

This counts the number of JSX components that have useIntl() used as a hook in each child directory of src, then plots it as a bar chart, filtering for 0s

import plotly.graph_objects as go

# Initialize a dictionary to hold the counts of useIntl functions for each app
app_counts = {}

# Iterate over all top-level directories in the src/ directory
for directory in codebase.directories:
    if directory.parent and directory.parent.name == "src":  # Check if it's a top-level directory under src
        # Initialize the count for this directory
        app_counts[directory.name] = 0
        # Iterate through all functions in the directory
        for function in directory.functions:
            # Check if the function contains useIntl and is a JSX function
            # Note - performance optimization since `function.function_calls` is slower
            if "useIntl()" in function.source and function.is_jsx:
                app_counts[directory.name] += 1  # Increment the count

# Filter to Non-zero
app_counts = {k: v for k, v in app_counts.items() if v > 0}

# Prepare data for the bar chart
apps = list(app_counts.keys())
counts = list(app_counts.values())

# Create the bar chart
fig = go.Figure(data=[go.Bar(x=apps, y=counts)])

# Update the layout of the graph
fig.update_layout(
    title="Count of useIntl Functions in Each App",
    xaxis_title="App",
    yaxis_title="Count of useIntl Functions",
    xaxis_tickangle=-45,  # Rotate x-axis labels for better readability
)

# Visualize the chart
codebase.visualize(fig)

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())})")

enum_attribute_mismatch_checker

None

# Iterate through all classes in the codebase
for cls in codebase.classes:
    # Check if the class is a subclass of Enum
    if cls.is_subclass_of("Enum"):
        # Initialize a flag to check for mismatched attributes
        has_mismatched_attributes = False
        # Iterate through the attributes of the Enum class
        for attr in cls.attributes:
            # Check if the attribute name and value are not the same
            if attr.name != attr.value:
                has_mismatched_attributes = True
                break
        # If there are mismatched attributes, print the class name
        if has_mismatched_attributes:
            print(f"Enum Class: {cls.name} has mismatched attributes.")

enum_by_value_modifier

None

# Get the first class from the enum_field.py file
enum_field_cls = codebase.get_file("vendored/enum_field.py").classes[0]

# Gather all subclasses of the enum_field_cls, including itself
subclasses = [*enum_field_cls.subclasses, enum_field_cls]

# Iterate through each enum class
for enum_cls in subclasses:
    print(f"Processing enum class: {enum_cls.name}")

    # Retrieve the call sites for the current enum class
    call_sites = enum_cls.call_sites

    # Limit to the first 10 call sites for processing
    for call_site in call_sites[:10]:
        # Get the arguments of the call site
        args = call_site.args

        # Iterate through each argument
        for arg in args:
            print(f"Checking argument: {arg.name} with resolved value: {arg.resolved_value}")

            # If the argument is 'by_value' and its resolved value is 'True', change it to 'False'
            if arg.name == "by_value" and "True" in arg.resolved_value:
                print(f"Modifying argument '{arg.name}' from True to False")
                arg.edit("by_value=False")

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()

extract_inline_props_type

I have a bunch of React components like this:

function MyComp(props: { a: A }) {...}

// or

const Z = ({ x }: { x: boolean }) => ...

And I want to extract the props type like this:

type MyCompProps = { a: A };

function MyComp(props: MyCompProps) { ... }

type ZProps = { x: boolean }

const Z = ({ x }: ZProps) => ...

So, that is, I want to extract the inlined props into their own type, then have the props be of this type.

for file in codebase.files:
    # Iterate over all functions in the file
    for function in file.functions:
        # Check if the function is a React functional component
        if function.is_jsx:  # Assuming is_jsx indicates a function component
            # Check if the function has inline props definition
            if len(function.parameters) == 1 and isinstance(function.parameters[0].type, Dict):
                # Extract the inline prop type
                inline_props: TSObjectType = function.parameters[0].type.source
                # Create a new type definition for the props
                props_type_name = f"{function.name}Props"
                props_type_definition = f"type {props_type_name} = {inline_props};"

                # Set the new type for the parameter
                function.parameters[0].set_type_annotation(props_type_name)
                # Add the new type definition to the file
                function.insert_before("\n" + props_type_definition + "\n")

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.

find_and_rename_loading_state_variable

This finds all instances of useState being assigned to a variable called loading, i.e. const [loading, setLoading] = useState(false), and renames the variable loading locally.

for function in codebase.functions:
    # Check if it is a JSX function
    if function.is_jsx:
        # Iterate through all assignment statements in the function
        for assignment_stmnt in function.code_block.assignment_statements:
            # Check if the assignment is a call to useState
            if isinstance(assignment_stmnt.value, FunctionCall) and assignment_stmnt.value.name == "useState":
                # Check if the state variable name is 'loading'
                if assignment_stmnt.assignments[0].name == "loading":
                    # Rename it within the current scope
                    function.rename_local_variable("loading", "isLoading")

flag_and_rename

This skill uses both flag_ai and ai to flag and rename all functions that are incorrect or misleading. The first step is to use flag_ai to flag the functions that are incorrect or misleading. Then, the second step is to use ai to generate a new name for the function. Finally, the function is renamed in the codebase.

This is an example of how codebase.flag_ai and codebase.ai can be used together to achieve a more complex task.

flag_code

This skill uses flag_ai to flag all code that may have a potential bug. This is an example of Codebase AI being used to flag code that meets a certain criteria.

flag_plan_imports

This codemod looks through the codebase for all imports that import Plan (a DB model) and flags the import.

for file in codebase.files:
    # Iterate over all imports (not import statements)
    for imp in file.imports:
        # Check the name
        if imp.name == "Plan":
            # Flag it
            imp.flag()

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}")

fragment_to_shorthand_codemod

This codemod demonstrates the transition from React’s <Fragment /> to < />.

It accomplishes this by iterating over all jsx_elements in the codebase and doing element.set_name(…).

It also does prop filtration by name on these elements.

In this specific scenario, where it was originally implemented, the remaining excess imports are automatically removed by the linter.

You could, however, accomplish this manually by iterating over all Imports in impacted files and doing imp.remove().

for file in codebase.files:
    # Iterate through all JSX elements in the file
    for element in file.jsx_elements:
        # Check if the element is a Fragment
        if element.name == "Fragment" and "key" not in map(lambda x: x.name, element.props):
            # Replace Fragment with shorthand syntax
            element.set_name("")  # This effectively converts `<Fragment>` to `<>`

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.

hook_inspector

For a given filename, print out all the hooks that are imported and differentiate between internal (user-defined code) and external (e.g. something defined by React)

FILENAME = "src/layouts/mobileLayout/screens/MyLindiesScreen/views/TasksView.tsx"

print("🎣 🎣 🎣 FISHING FOR HOOKS 🎣 🎣 🎣")
print(f"filename: {FILENAME}\n")

# Grab a specific file
file = codebase.get_file(FILENAME)

# Initialize a set to store names of hooks
hooks = set()

# Check all import statements in the file
for import_statement in file.import_statements:
    # Iterate through individual imported symbols
    for imp in import_statement.imports:
        # Filter for hooks
        if imp.imported_symbol.name.startswith("use"):
            # Will be a `Function` if the user defined it
            if isinstance(imp.imported_symbol, Function):
                hooks.add(imp.imported_symbol.name)
                print(f"🪝🟢 Internal hook: {imp.imported_symbol.name}")
            # Otherwise, will be an `ExternalModule`
            elif isinstance(imp.imported_symbol, ExternalModule):
                hooks.add(imp.imported_symbol.name)
                print(f"🪝🔵 External hook: {imp.imported_symbol.name}")

if_statement_depth_logger

None

def check_if_depth(if_block: IfBlockStatement, current_depth):
    # If the current depth exceeds 1, log the if block
    if current_depth > 2:
        print(f"Depth: {current_depth}:\n_____\n{if_block.source}\n")
        if_block.insert_before("# 🚩 🚩 🚩", fix_indentation=True)
    # Recursively check nested if blocks
    for nested_if in if_block.consequence_block.if_blocks:
        check_if_depth(nested_if, current_depth + 1)

# Get all functions + methods
callables = codebase.functions + [m for c in codebase.classes for m in c.methods]

# Iterate over all functions in the file
for function in callables:
    # Check each if block in the function
    for if_block in function.code_block.if_blocks:
        check_if_depth(if_block, 1)

import_relationship_graph_visualizer

None

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

list_apps = ["apps/spend_request", "apps/spend_allocation", "apps/approval_policy"]
app_to_imp = {app: "" for app in list_apps}
num_imports = {app: 0 for app in list_apps}

# Add nodes for each app
for app in list_apps:
    app_name = app.split("/")[2]
    G.add_node(app_name, color="red")

# Iterate through each app directory
for app in list_apps:
    directory = codebase.get_directory(app)
    for file in directory.files:
        # Iterate over all import statements in the file
        for import_statement in file.import_statements:
            # Check if the import statement is importing an app
            for imp in import_statement.imports:
                if "app" in imp.name:
                    node_name = imp.import_statement.source.split()[1]
                    node_name = node_name.split(".")[3:]
                    node_name = ".".join(node_name)

                    app_to_imp[app] += imp.import_statement.source + "\n"
                    num_imports[app] += 1

                    app_name = app.split("/")[2]
                    G.add_node(app_name, code=app_to_imp[app], text=f"{num_imports[app]} imports")
                    G.add_edge(app_name, node_name)

# Remove nodes with only one connection
nodes_to_remove = [node for node, degree in G.degree() if degree == 1]
G.remove_nodes_from(nodes_to_remove)

codebase.visualize(G)

inconsistent_assignment_name_analyzer

This skill iterates through all functions in the codebase and identifies ones that are inconsistent in the value they get assigned when called as part of an assignment statement. For example, if a function foo() is frequently called and assigned to variables of different names, this would get printed out. Prints things out aesthetically

isolate_jsx_components

This codemod source code is designed to refactor a TypeScript codebase written with React (specifically .tsx files that contain JSX). Its primary purpose is to isolate each JSX component function in its individual file.

The code first scans those TypeScript files in the codebase that end with ‘.tsx’ extension, indicating they contain JSX. Inside these files, it identifies JSX component functions and maps them to their associated file names in a dictionary. To do this, the code checks whether a function is a JSX component and checks if the component function has a valid (not None) name.

Following that, it iterates over each file identified earlier and for each JSX component in the file, constructs a new file path. If a file doesn’t already exist at the new file path, it generates a new file at that path, imports React, and then shifts the JSX component to this new file. It takes care to maintain all dependencies in the process.

If a file already exists for a component, it is not overwritten, and a message is printed to the console. The refactoring strategy used here is ‘update_all_imports’ which means that the codemod will update all locations where the moved function is imported.

By doing this, the codemod effectively splits up bulkier React files into separated ones, with each file containing a single JSX component, leading to a more organized and manageable codebase.

components_by_file = {}

# Iterate through all files in the codebase
for file in codebase.files:
    # Check if the file is a TypeScript file containing JSX
    if file.filepath.endswith(".tsx"):
        # Initialize the list for this file
        components_by_file[file.filepath] = []
        # Iterate through all functions in the file
        for function in file.functions:
            # Check if the function is a JSX component
            if function.is_jsx:
                # Ensure the component name is valid (not None)
                if function.name:
                    # Add the component name to the list for this file
                    components_by_file[file.filepath].append(function.name)

# Iterate through the components by file
for filepath, components in components_by_file.items():
    # Get the directory of the current file
    component_dir = Path(filepath).parent
    # Iterate through each component in the list
    for component in components:
        # Create a new file path for the component
        new_file_path = component_dir / f"{component}.tsx"
        # Check if the file already exists
        if not codebase.has_file(str(new_file_path)):
            # Create a new file for the component
            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
            # Get the function corresponding to the component name
            function = codebase.get_file(filepath).get_function(component)
            function.move_to_file(new_file, include_dependencies=True, strategy="update_all_imports")
        else:
            print(f"File already exists for component: {component}")

jsx_elements_collector

This codemod script scans through the entire codebase, identifies files that contain JSX elements and collates these JSX elements into a dictionary. This dictionary maintains a list of JSX element names for each source file path.

The script initially sets an empty dictionary. For each file in the codebase, it iterates through the JSX elements it contains. If the file path is not already present as a key in the dictionary, it initializes it with an empty list. It then adds the JSX element’s name to this list, provided the name isn’t None.

Finally, the script iterates through the dictionary entries and prints the JSX elements grouped by file. Before printing, it filters out any None values from the list of JSX elements for each file. The output format is “File: [filepath], JSX Elements: [names of the JSX elements separated by commas]“.

In essence, this script conveniently organizes and displays JSX element names on a per-file basis.

jsx_elements_dict = {}

# Iterate through all files in the codebase
for file in codebase.files:
    # Check if the file contains JSX elements
    for element in file.jsx_elements:
        # If the file path is not in the dictionary, initialize it with an empty list
        if file.filepath not in jsx_elements_dict:
            jsx_elements_dict[file.filepath] = []
        # Append the element name to the list for the corresponding file path, if it's not None
        if element.name is not None:
            jsx_elements_dict[file.filepath].append(element.name)

# Print the collected JSX elements information grouped by file
for filepath, elements in jsx_elements_dict.items():
    # Filter out any None values before joining
    valid_elements = [name for name in elements if name is not None]
    print(f'File: {filepath}, JSX Elements: {', '.join(valid_elements)}')

jsx_use_state_type_inferer

This finds all useState calls within JSX components and tries to infer their type by looking at the value, then sets the generic passed to it.

for function in codebase.functions:
    # Filter for JSX functions
    if function.is_jsx:
        # Iterate through all assignments in the function
        for assignment in function.code_block.assignment_statements:
            # Filter out assignments with no right side
            if assignment.value:
                # Find assignments to `useState` function calls
                if isinstance(assignment.value, FunctionCall) and assignment.value.name == "useState":
                    fcall = assignment.value

                    # Find ones without a type
                    if "useState<" not in fcall.source:
                        # Extract initial value being passed to `useState`
                        value = fcall.args[0].value

                        # Add type annotations
                        if isinstance(value, String):
                            fcall.replace("useState", "useState<string>")
                        elif isinstance(value, Boolean):
                            fcall.replace("useState", "useState<bool>")
                        elif isinstance(value, Number):
                            fcall.replace("useState", "useState<number>")

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)

migrate_imports_to_backward_compatible_versions

None

# Utility function for migrating imports to backwards-compatible versions
def migrate_import(_import: PyImport, target: str, new: str) -> None:
    # Check if the import module name is matching what we are looking for
    if _import.module == target:
        print(f"Processing: {_import.source}")
        _import.set_import_module(new)

# From:
#   from math import abs
#   import math
# To:
#   from math_replacement import abs
#   import math_replacement
migrate_import(codebase.imports[0], "math", "math_replacement")

modify_test_files_for_uuid

This codemod script has a primary purpose of modifying test files in a codebase to replace hardcoded “space-1” strings with call to a uuid() function. The uuid function provides a universally unique identifier which is desirable over hardcoded values for testing environments for uniqueness and non-predictability.

Key features of the script:

1. Iterates over all files in the codebase.
2. Checks if the active file is a test file and contains the string “space-1”.
3. If both conditions met, it then verifies if it already has “uuid” import. If not, it adds the import statement from a specific module path.
4. It then searches for occurrences of the hardcoded string “space-1” in the test file and replaces them with a call to the uuid() function.

for file in codebase.files:
    # Check if the file is a test file and "space-1" string literal in the content (perf optimization, only look at relevant files)
    if "test" in file.filepath and '"space-1"' in file.content:
        # Check if the import already exists
        if not file.has_import("uuid"):
            # Add the import statement for uuid
            file.add_import_from_import_string("import { uuid } from 'gather-common-including-video/dist/src/public/uuid';")

        # Search for occurrences of the string "space-1"
        for occurrence in file.search('"space-1"'):
            # Replace the occurrence with a call to uuid()
            occurrence.edit("uuid()")

modularize_classes_in_codebase

This codemod skill script is designed to iterate over all files in a given codebase to improve the modularity and maintainability of the code. Its primary function is to move any classes that are extending from a superclass defined in the same file into their own separate files. However, it only processes files that are more than 500 lines long, skipping those that are shorter.

Starting its process, the script calculates the length of each file by splitting the content by line. If a file has fewer than 500 lines, it continues to the next. If a file is long enough, it examines each class within the file. For each class, it looks at the superclasses and checks if any of them are defined in the same file.

If a superclass is located in the same file as its child class, the script proceeds to move the child class into a new, separate file. It first builds the path for the new file, using the current directory and the name of the class, with the original file’s extension. It then creates a new file at this path in the codebase, initially with an empty content. The child class is then moved into this new file, using the “update_all_imports” strategy to ensure all references to the class are correctly updated. The script also provides feedback during its execution, printing out a message each time it moves a class to a new file.

This codemod skill script can be helpful in large codebases where class definitions may be tightly

for file in codebase.files:
    file_length = len(file.content.split("\n"))
    if file_length < 500:
        continue
    # Check all classes in the file
    for cls in file.classes:
        # Check if any of the superclasses are defined in the same file
        for superclass in cls.superclasses:
            # if the superclass is in the same file
            if superclass.file == cls.file:
                # move the child class into it's own file
                cur_dir = file.directory
                extension = file.extension
                cls_file_path = f"{cur_dir.path}/{cls.name}{extension}"
                cls_file = codebase.create_file(cls_file_path, content="")
                print(f":shoe: Moving {cls.name} to {cls_file_path}")
                cls.move_to_file(cls_file, include_dependencies=False, strategy="update_all_imports")

most_common_use_state_variables

This aesthetically prints out a sorted list of the most common values that are used in a useState as the state variable.

from collections import Counter

# Initialize a Counter to count occurrences of state variable names
state_variable_counts = Counter()

# Iterate through all functions in the codebase
for function in codebase.functions:
    # Check if it is a JSX function
    if function.is_jsx:
        # Iterate through all assignment statements in the function
        for assignment in function.code_block.assignment_statements:
            # Check if the assignment is a call to useState
            if isinstance(assignment.value, FunctionCall) and assignment.value.name == "useState":
                # Get the state variable name
                state_var_name = assignment.assignments[0].name
                # Increment the count for this state variable name
                state_variable_counts[state_var_name] += 1

# Print the counts of state variable names
print("#####[ Most Common State Variable Names ]#####")
for name, count in state_variable_counts.most_common():
    print(f"{count}: {name}")

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)

react_component_render_map

For each react component in a specific file, print out the names of all components that it might render, skipping anonymous ones and tags like

etc.

file = codebase.get_file("src/layouts/MyLindies/sections/AgentTable/components/TableHeader.tsx")

# Initialize a dictionary to store the mapping of components to their rendered child components
component_render_map = {}

# Check if the file contains any JSX components
for function in file.functions:
    # Filter to named JSX components
    if function.is_jsx and function.name:
        # Initialize a list to store child components for the current component
        child_components = []
        # Iterate through all JSX elements in the function
        for element in function.jsx_elements:
            # Filter out anonymous ones and lowercase ones
            if element.name and element.name[0].isupper():
                # Add new ones to child_components
                if element.name not in child_components:
                    child_components.append(element.name)
        # Save child component
        component_render_map[function.name] = child_components

# Print the resulting mapping
for component, children in component_render_map.items():
    print(f'⚛️ {component} renders: {', '.join([x for x in children])}')

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.

reinforce_data_schemas

The provided code is a Python script that dynamically analyzes and modifies the source code of a given Python program to introduce data classes for stronger data type enforcement.

In essence, it’s a Codemod utility that reinforces static typing in a codebase by introducing @dataclass definitions into functions and call sites.

Key points:

1. At the start of the script, some constant like the function name and parameter name to analyze are declared. NOTE: Overridable attribute types and defaults for those attributes can also be defined.

2. The script proceeds to examine the identified function and its parameters in the codebase. If the function or parameters don’t exist, it raises an error.

3. It then iterates over all call sites of the function, accumulating arguments that match a dictionary data type.

4. The keys in all these dictionary arguments are analyzed to identify unique keys and their frequency.

5. Each unique key across all the dictionaries is then typed by inferring the value type associated with the key in the different call sites.

6. A @dataclass definition is then created with attributes that correspond to the unique keys identified and their inferred or overridden types.

7. This @dataclass is injected into the source code, directly after the last import line.

8. The function parameter type is adjusted to use this new @dataclass.

9. Every call site of the function is then updated to initialize a new instance of this @dataclass in place of

FUNCTION_NAME = "create_procurement_spend_request_helper"
PARAM_NAME = "spend_request_vendor_args"

# Optional for all arguments (but recommended)
schema_attr_type_overrides = {
    "vendor_name": "str | None",
    "payee_id": "int | None",
}
schema_attr_defaults = {"vendor_name": "= None", "payee_id": "= None"}

function = codebase.get_function(FUNCTION_NAME)
if not function:
    raise ValueError(f"No such function: (f{FUNCTION_NAME})")

parameter = function.get_parameter(PARAM_NAME)
if not parameter:
    raise ValueError(f"No such parameter: (f{FUNCTION_NAME}) f{PARAM_NAME}")

print("##########################################")
print("########## 🐢 Reinforce Schemas ##########")
print("##########################################")
print(f"Function: {FUNCTION_NAME}")
print(f"Parameter: {PARAM_NAME}")
print()

# Dictionary to hold counts of each key in ramp_data
matching_args = []

# Iterate through all call sites of the function
for call_site in function.call_sites:
    # Get the argument for ramp_data
    arg = call_site.get_arg_by_parameter_name(PARAM_NAME)
    if arg and isinstance(arg.value, Dict):
        # Add it to matching args
        matching_args.append(arg)
        # Optional - flag the original args
        # arg.insert_after('# 👈 🚩🚩🚩\n', newline=False, fix_indentation=True)

print("\n### 🔬 Analytics ###")
print(f"📞 Found {len(function.call_sites)} call-sites")
print(f"🔎 Found {len(matching_args)} dict args")

# Possible top-level attributes that will live on this schema

def flatten(level):
    return [y for x in level for y in x]

# List of all keys as a dict
all_keys = flatten([x.value.keys() for x in matching_args])
unique_keys = set(all_keys)

# Display key counts
print("\n### 📊 Frequency of keys: ###")
import pandas as pd

print(pd.Series(all_keys).value_counts())

# Grab all the values that each key takes on

values = {key: [] for key in unique_keys}
for arg in matching_args:
    data = arg.value
    for key in data.keys():
        values[key].append(data[key])

def get_type(value) -> str:
    if isinstance(value, String):
        return "str"
    #    if isinstance(value, Boolean):
    elif "Boolean" in str(type(value)):
        return "bool"
    #    elif isinstance(value, Number):
    elif "Number" in str(type(value)):
        return "int"
    elif value.source == "None":
        return "None"
    return "Any"

def infer_schema_attr_type(key_values: list) -> list[str]:
    """Look at all values the key has taken on to infer the value"""
    attr_types = list(set([get_type(v) for v in key_values]))
    return " | ".join(attr_types)

def to_camel_case(snake_str: str) -> str:
    return "".join(x.capitalize() for x in snake_str.lower().split("_"))

def render_schema_attribute(name: str, values: list) -> str:
    """Goes from a list of unique types"""
    default = schema_attr_defaults.get(name, "")
    annotation = schema_attr_type_overrides.get(name) or f" = {infer_schema_attr_type(values)}"
    return f"{name}: {default}{annotation}"

def define_dataclass(dataclass_name: str, values: dict[str, list[str]]) -> str:
    """Creates a dataclass from a dict"""
    attributes = "\n".join([render_schema_attribute(k, v) for k, v in values.items()])
    return f"""@dataclass\nclass {dataclass_name}:\n{attributes}""".strip()

print("\n### 📦 @dataclass definition ###")
print("```")
dataclass_name = to_camel_case(f"{PARAM_NAME}_schema")
dataclass_source = define_dataclass(dataclass_name, values)
print(dataclass_source)
print("\n```\n")

print("### 🚀 Creating Dataclass ###")

# =====[ Add dataclass definition ]=====
file = function.file
last_import = file.imports[-1]
last_import.insert_after(f"\n\n{dataclass_source}")
print("  ✅ Added Dataclass definition")

# =====[ Replace parameter ]=====
if "| None" in parameter.source:
    parameter.set_type_annotation(dataclass_name + " | None")
else:
    parameter.set_type_annotation(dataclass_name)
print("  ✅ Updated dataclass type signature")

# =====[ Grab import line for new symbol ]=====
import_line = function.get_import_string()
import_line = f"from {function.file.import_module_name} import {dataclass_name}"

print("\n### 🚀 Editing Call-sites ###")

def make_schema(arg: Argument, call_site: FunctionCall):
    data = arg.value
    args = ", ".join([f"{key}={data[key].source}" for key in data.keys()])
    return f"{dataclass_name}({args})"

count = 0
# Iterate over all call sites
for call_site in function.call_sites:
    # Extract the relevant arg
    arg = call_site.get_arg_by_parameter_name(parameter.name)
    if arg:
        # Replace it with a schema
        arg.set_value(make_schema(arg, call_site))
        # Import the schema if not already existent
        if import_line not in call_site.file.content and call_site.file != function.file:
            call_site.file.add_import_from_import_string(import_line)
        count += 1

print(f"  ✅ Updated {count} call-sites")

relay_fragment_updater

None

# Configuration section for easy customization
# Set the parameters for the codemod below

TARGET_FILE_PATH = "src/app/DeveloperTools/DeveloperToolsFeatureFlags.tsx"
FUNCTION_NAME = "DeveloperToolsFeatureFlags"
TARGET_PROP_FIELD = "isKenDocs"
RELAY_PROP_FIELD = "pinRef"
RELAY_PROP_TYPE = "PinTitleDisplay_pin$key"
FRAGMENT_NAME = "PinDisplayDisplay_pin"

def update_relay_fragment(target_file_path: str, function_name: str, target_prop_field: str, relay_prop_field: str, relay_prop_type: str, fragment_name: str):
    # Log the start of the codemod
    print("🚀 Starting the codemod...")

    # Get the target file
    target_file = codebase.get_file(target_file_path)

    # Find the specific function to modify
    function_to_modify = target_file.get_function(function_name)

    # Define the mapping for prop fields to fragment schema
    propFieldsToFragmentSchemaMap = {target_prop_field: "pinTitle"}

    def buildFragment(fragment_name, sourceType, fields: List[str]):
        """Builds a GraphQL fragment string."""
        return f"fragment {fragment_name} on {sourceType} {{ \n" + "\n \t".join(fields) + "\n}"

    # Check if the function is found and is a React component
    if function_to_modify and function_to_modify.is_jsx:
        print("🔍 Modifying function parameters...")

        # Get the current parameters (props)
        current_params = function_to_modify.parameters

        # Modify the props to rename target prop
        for param in current_params:
            for usage in function_to_modify.code_block.get_variable_usages(param.name):
                if param.name in propFieldsToFragmentSchemaMap:
                    usage.replace(usage.source, propFieldsToFragmentSchemaMap[param.name])

            if param.name == target_prop_field:
                print(f"✏️ Renaming prop '{param.name}' to '{relay_prop_field}'...")
                param.set_name(relay_prop_field)
                param.set_type_annotation(relay_prop_type)

        # Add the import for useFragment
        print("📦 Adding import for useFragment...")
        target_file.add_import_from_import_string("import { useFragment } from 'react-relay'")

        # Create the new useFragment hook
        use_fragment_code = f"const data = useFragment(graphql`{buildFragment(fragment_name, relay_prop_type, list(propFieldsToFragmentSchemaMap.values()))}`, {relay_prop_field})"

        # Insert the useFragment hook at the beginning of the function body
        print("🔄 Inserting useFragment hook...")
        function_to_modify.insert_statements(use_fragment_code)

    # Log the completion of the codemod
    print("✅ Codemod completed successfully!")

# Call the function with specific parameters
update_relay_fragment(TARGET_FILE_PATH, FUNCTION_NAME, TARGET_PROP_FIELD, RELAY_PROP_FIELD, RELAY_PROP_TYPE, FRAGMENT_NAME)

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.

rename_use_state_setter

This finds all calls to useState hook in React where the setter is not just setXYZ, where XYZ is the name of the state variable, and renames it. Note that this does not handle the case where there’s a collision between the setter name and another function defined in the component.

for function in codebase.functions:
    # Check if it is a JSX function
    if function.is_jsx:
        # Iterate through all assignment statements in the function
        for assignment in function.code_block.assignment_statements:
            # Check if the assignment is a call to useState
            if isinstance(assignment.value, FunctionCall) and assignment.value.name == "useState":
                if len(assignment.assignments) == 2:
                    # Get the state variable name
                    state_var_name = assignment.assignments[0].name
                    # Get the setter name
                    setter_name = assignment.assignments[1].name
                    # Properly capitalize the state variable name
                    expected_name = "set" + state_var_name[0].upper() + state_var_name[1:]
                    # Check if the setter name does not match the expected pattern
                    if setter_name != expected_name:
                        # Flag
                        function.rename_local_variable(setter_name, expected_name)

replace_aliased_wildcard_imports

This takes all imports in a given file that are imported via wildcard alias (and filters out lodash) and then converts them to be explicit regular (non-wildcard) imports, making sure to update the imports accordingly. uses the usage.match.parent.attribute to get chained attributes from the wildcards.

print("#####[ Replace Aliased * Imports ]#####")
DIRECTORY = "src/chromeExtension/"
print(f"Directory: {DIRECTORY}\n")

# Get the specific directory
directory = codebase.get_directory(DIRECTORY)

# Iterate over all files in the specified directory
for file in directory.files:
    # Check all imports in the file
    for imp in file.imports:
        # If the import is a wildcard import
        if imp.is_wildcard_import() and not imp.name == "_":
            # Replace with explicit imports (this is a placeholder for the actual logic)
            print(f"⬇️ Import: {imp.source} ")
            # Iterate through usages in the file
            for usage in imp.usages:
                # Get the name of the parent (chained attributes)
                full_name = usage.match.parent
                # Replace with the name of the attribute itself
                print(f"  👏 Replacing attribute: {full_name.attribute}")
                # Create new name
                new_name = imp.name + full_name.attribute.source.capitalize()
                # Replace the usage
                full_name.edit(new_name)
                # Add import of the symbol
                symbol = imp.from_file.get_symbol(full_name.attribute.source)
                file.add_symbol_import(symbol, alias=new_name)

            # Remove the original (wildcard) import
            imp.remove()

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.

return_type_co_occurrence_matrix

Look at every union return type and find a co-occurrence matrix for the options, then print it out in a way that’s easy to read. We will use this to find pairs that go together frequently and make more abstract types.

from collections import defaultdict

# Dictionary to hold counts of return type co-occurrences
co_occurrence_matrix = defaultdict(lambda: defaultdict(int))

# Iterate through all functions in the codebase
for function in codebase.functions:
    # Get the return type of the function
    return_type = function.return_type
    if isinstance(return_type, UnionType):
        # Get all options in the union type
        options = [x.source for x in return_type]
        # Update the co-occurrence counts
        for i in range(len(options)):
            for j in range(i + 1, len(options)):
                co_occurrence_matrix[options[i]][options[j]] += 1
                co_occurrence_matrix[options[j]][options[i]] += 1

# Print the co-occurrence matrix in a more aesthetic way
print("🪐 Co-occurrence Matrix of Return Types:\n")
for return_type, counts in co_occurrence_matrix.items():
    sorted_counts = dict(sorted(counts.items(), key=lambda item: item[1], reverse=True))
    print(f"Return Type: {return_type}")
    for co_type, count in sorted_counts.items():
        print(f"  - Co-occurs with '{co_type}': {count} times")
    print()  # Add a newline for better separation

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")

skill_language_distribution_analyzer

This codemod script scrutinizes a codebase to survey and count the functional skills that are implemented in Python and TypeScript. For every class in the codebase, it explores the methods for a ‘@skill_impl’ decorator - a mark signifying an implemented skill.

If it finds such a decorator, it identifies the programming language used for the skill implementation by looking for ‘language=ProgrammingLanguage.PYTHON’ or ‘language=ProgrammingLanguage.TYPESCRIPT’ in the decorator’s source.

The script creates a dictionary that maps each class (considered as a skill) to the languages that it supports (Python or TypeScript). The language support for each skill is represented with a corresponding emoji - a snake (🐍) for Python and a square (🟦) for TypeScript.

The script also counts the number of skills implemented in each language and prints these numbers. Finally, it displays the classes (skills) that only support Python but not TypeScript.

This script is beneficial for evaluating the language distribution across different skills in a mixed-languages codebase, and for identifying skills that should be upgraded to multi-language support.

skills_with_languages = {}
# Initialize counters for each language
python_count = 0
typescript_count = 0

# Iterate over all classes in the codebase
for cls in codebase.classes:
    # Check each method in the class
    for method in cls.methods:
        # Check if the method has the @skill_impl decorator
        for decorator in method.decorators:
            if "@skill_impl" in decorator.source:
                # Extract the language from the decorator
                if "language=ProgrammingLanguage.PYTHON" in decorator.source:
                    skills_with_languages.setdefault(cls.name, []).append("🐍 Python")
                    python_count += 1
                elif "language=ProgrammingLanguage.TYPESCRIPT" in decorator.source:
                    skills_with_languages.setdefault(cls.name, []).append("🟦 TypeScript")
                    typescript_count += 1

# Print the counts of skills per language
print(f"Total Skills: {python_count + typescript_count}")
print(f"🐍 Python Skills: {python_count}")
print(f"🟦 TypeScript Skills: {typescript_count}")

# Log the results for skills that support both languages
print("Skills that support both Python and TypeScript:")
for sk, languages in skills_with_languages.items():
    if "🐍 Python" in languages and "🟦 TypeScript" not in languages:
        print(f"- {sk}")

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")')

sql_alchemy_eager_loading_optimizer

This codemod script is designed to analyze a codebase for SQLAlchemy relationships and optimize usage of eager loading in ORM queries. Purpose: Its primary goal is to capture relationships defined using db.relationship and to optimize the loading strategy from joinedload to selectinload when applicable. Functionality: 1. Relationship Extraction: The first step collects all relationships within the classes of the codebase, specifically checking for the uselist argument of db.relationship to determine if the relationship represents a singular or multiple entities. 2. Optimization of Load Strategies: The second step inspects function calls across files to identify usages of selectinload and joinedload. It optimizes any joinedload calls with uselist=True by changing them to selectinload, which enhances performance by reducing the number of queries. Key Features: - It logs captured relationships for easy reference. - Automatically adjusts loading strategies for improved database querying efficiency. - Ensures that the necessary imports for selectinload are added when changes are made. Overall, the script aims to streamline ORM queries in SQLAlchemy and improve application performance related to data loading strategies.

relationships = {}
for cls in codebase.classes:
    for attr in cls.attributes:
        # Iterate through each assignment of the attribute
        for assignment in attr.assignments:
            # Check if the assignment is a FunctionCall and ends with 'db.relationship'
            if isinstance(assignment.value, FunctionCall) and "db.relationship" in assignment.value:
                for arg in assignment.value.args:
                    # Check for the 'uselist' argument
                    if arg.name == "uselist":
                        relationships[f"{cls.name}.{attr.name}"] = {"uselist": arg.value.source}
                        break
                else:
                    # Default to 'False' if 'uselist' is not found
                    relationships[f"{cls.name}.{attr.name}"] = {"uselist": "False"}

# Log out Relationships
print("#####[ Relationships ]#####")
print(relationships)

# Step 2: Find usages of selectinload or joinedload
for file in codebase.files:
    added = False
    for call in dict.fromkeys(file.function_calls):
        # Check if the function call is either selectinload or joinedload
        if call.name in ("selectinload", "joinedload"):
            relationship = call.args[0].source
            fn = call.name
            # Check if the relationship is known
            if relationship not in relationships:
                print("unknown", relationship)
                continue

            # If it's a joinedload with uselist=True, change it to selectinload
            if fn == "joinedload" and relationships[relationship]["uselist"] == "True":
                added = True
                call.set_name("selectinload")
    if added:
        # Add import for selectinload if any changes were made
        file.add_import_from_import_string("from sqlalchemy.orm import selectinload")

sql_alchemy_mapped_type_enhancer

The provided codemod script is designed to update SQLAlchemy model definitions by converting traditional column types into their mapped type equivalents using Python’s typing system. Purpose: To enhance SQLAlchemy column definitions by adding type annotations for better type checking and clarity in Python projects. Functionality: - It scans a specified directory for SQLAlchemy model classes. - It identifies attributes defined using db.Column and checks their types. - It maps standard SQLAlchemy column types to the new Mapped type annotations using a predefined dictionary. - It supports nullable types by wrapping them in Optional. - It ensures that the necessary import statement for Mapped is added if it is not already present. Key Features: - Automatic conversion of SQLAlchemy column types to type-safe annotations. - Detection of nullable attributes and appropriate wrapping in Optional. - Preservation of existing structure while enhancing type safety. - Automated import management for seamless integration with existing files.

print("#####[ Add Mapped Types to SQLAlchemy ]#####")
DIRECTORY = "suma/apps/accounting"
print(f"Directory: {DIRECTORY}\n")

######[ MAPPING ]#####
column_type_to_mapped_type = {
    "db.Integer": "Mapped[int]",
    "Optional[db.Integer]": "Mapped[int | None]",
    "db.Boolean": "Mapped[bool]",
    "Optional[db.Boolean]": "Mapped[bool | None]",
    "Optional[db.DateTime]": "Mapped[datetime | None]",
    "db.Text": "Mapped[str]",
    "Optional[db.Text]": "Mapped[str | None]",
    "db.String": "Mapped[str]",
    "Optional[db.String]": "Mapped[str | None]",
    "db.Float": "Mapped[float]",
    "Optional[db.Float]": "Mapped[float | None]",
    "db.BigInteger": "Mapped[int]",
    "Optional[db.BigInteger]": "Mapped[int | None]",
    "UUID": "Mapped[str]",
    "Optional[UUID]": "Mapped[str | None]",
}

# Grab directory
directory = codebase.get_directory(DIRECTORY)

# Iterate over all classes
for cls in directory.classes:
    # Get all the attributes which are assignments
    for attribute in cls.attributes:
        if attribute.assignment.type:
            continue

        # Check if the attribute is a db.Column call
        assignment_value = attribute.assignment.value
        if not isinstance(assignment_value, FunctionCall) or assignment_value.name != "Column":
            continue

        db_column_call = assignment_value
        is_nullable = any(x.name == "nullable" and x.value == "True" for x in db_column_call.args)
        first_argument = db_column_call.args[0].source.split("(")[0]

        # Add `Optional` to first arg
        if is_nullable:
            first_argument = f"Optional[{first_argument}]"

        if first_argument not in column_type_to_mapped_type.keys():
            continue

        # Set type annotation
        new_type = column_type_to_mapped_type[first_argument]
        attribute.assignment.set_type_annotation(new_type)

        # Add the import
        if not cls.file.has_import("Mapped"):
            cls.file.add_import_from_import_string("from sqlalchemy.orm import Mapped\n")

temporal_activity_name_updater

Update temporal activity decorators.

1. Finds all temporal “activities” (functions decorated with @activity.defn(…))
2. For those that have a name= parameter in this decorator, replaces the value to add my_directory at the beginning if it starts with backfill

for file in codebase.files:
    # Performance optimization - only look at files with @activity
    if "@activity" not in file.content:
        continue

    # Iterate through all functions in the file
    for function in file.functions:
        # Iterate through decorators
        for decorator in function.decorators:
            # Find @activity decorators
            if "@activity.defn" in decorator.source:
                # Check to see if they have a function call
                if len(decorator.function_calls) > 0:
                    # Grab the `name` arg from this function call
                    call = decorator.function_calls[0]
                    name_arg = call.get_arg_by_parameter_name("name")
                    if name_arg:
                        # Edit this arg if it contains `backfill` at the beginning of the string
                        if '"backfill' in name_arg.value:
                            name_arg.value.replace('"backfill', '"my_directory.backfill')

transaction_canonical_refactorer

None

# Get the symbol for the function to import
get_merchant_name_for_transaction = codebase.get_symbol("get_merchant_name_for_transaction")

# Iterate over all files in the codebase
for file in codebase.files:
    # Iterate over all functions in the file
    for function in file.functions:
        # =====[ Function parameters ]=====
        # Check each parameter in the function
        for parameter in function.parameters:
            # If the parameter type is TransactionCanonical
            if parameter.type and parameter.type.source == "TransactionCanonical":
                # Find all instances of transaction_canonical.merchant.merchant_name
                matches = function.code_block.search(f"{parameter.name}.merchant.merchant_name")
                for match in matches:
                    print(f"🔵 Patching: {function.name} ({function.file.filepath})")
                    # Replace with get_merchant_name_for_transaction(transaction_canonical)
                    match.edit(f"get_merchant_name_for_transaction({parameter.name})")
                # Add the function import to the file
                if len(matches) > 0:
                    function.file.add_symbol_import(get_merchant_name_for_transaction)

        # =====[ Locally-assigned variables ]=====
        # Check local variable assignments in the function
        for assignment in function.code_block.local_var_assignments:
            # If the assignment type is TransactionCanonical
            if assignment.name == "transaction_canonical":
                # Find all instances of transaction_canonical.merchant.merchant_name
                matches = function.code_block.search(f"{assignment.name}.merchant.merchant_name")
                for match in matches:
                    print(f"🟠 Patching: {function.name} ({function.file.filepath}) (local var)")
                    # Replace with get_merchant_name_for_transaction(transaction_canonical)
                    match.edit(f"get_merchant_name_for_transaction({assignment.name})")
                # Add the function import to the file
                if len(matches) > 0:
                    function.file.add_symbol_import(get_merchant_name_for_transaction)

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_function_parameters_and_return_type

Sets the first 5 functions to have a new first parameter test, the second 5 functions to have a new last parameter test, and the next 5 functions to have return type null.

for function in codebase.functions[:5]:
    # Add an untyped parameter named 'test' to the function's parameters in the first spot
    function.parameters.insert(0, "test")

for function in codebase.functions[5:10]:
    # Add an untyped parameter named test at the end of the function's parameters
    function.parameters.append("test")

for function in codebase.functions[10:15]:
    # Set return type to null
    function.set_return_type("null")

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)

use_suspense_query_refactor

This codemod script is designed to refactor a codebase that uses the ‘useSuspenseQuery’ function from the @tanstack/react-query library to the ‘useSuspenseQueries’ function instead. It performs a few key steps:

1. Iterates over all files in the codebase, specifically focusing on files within the ‘src’ directory and ignoring files like ‘.eslint.rc’.

2. For each file, it checks if ‘useSuspenseQuery’ is present. If it isn’t, the file is skipped.

3. If ‘useSuspenseQuery’ is found, an import statement for ‘useQuery’ and ‘useSuspenseQueries’ from the @tanstack/react-query library is added to the file.

4. Then, it enters each function in the file and scans for assignment statements that use ‘useSuspenseQuery’. If any instances are found, it gathers the variable names and the arguments passed to ‘useSuspenseQuery’.

5. Once all instances of ‘useSuspenseQuery’ are collected, they are replaced with calls to ‘useSuspenseQueries’. The variable names and query arguments gathered earlier are used to structure the new function call.

6. Finally, the old ‘useSuspenseQuery’ calls are replaced with the newly constructed ‘useSuspenseQueries’ calls, and the refactoring is complete.

The main purpose of this codemod script is to automate the process of refactoring code to use the new ‘useSuspenseQueries’ function

import_str = "import { useQuery, useSuspenseQueries } from '@tanstack/react-query'"

# Iterate through all files in the codebase
for file in codebase.files:
    # Check if the file contains any useSuspenseQuery calls
    if "useSuspenseQuery" not in file.source:
        continue

    # Iterate through all functions in the file
    for function in file.functions:
        # Skip functions that do not call useSuspenseQuery
        if "useSuspenseQuery" not in function.source:
            continue

        results = []  # To store the left-hand side of assignments
        queries = []  # To store the arguments passed to useSuspenseQuery
        old_statements = []  # To keep track of old statements to be replaced

        # Iterate through assignment statements in the function
        for a in function.code_block.assignment_statements:
            # Ensure the right-hand side is a function call
            if not isinstance(a.right, FunctionCall):
                continue

            fcall = a.right
            # Check if the function call is useSuspenseQuery
            if fcall.name != "useSuspenseQuery":
                continue

            # Store the instance of the old useSuspenseQuery call
            old_statements.append(a)
            results.append(a.left.source)  # Collect the variable names
            queries.append(fcall.args[0].value.source)  # Collect the query arguments

        # If useSuspenseQuery was called at least once, convert to useSuspenseQueries
        if old_statements:
            # Add the import statement for useQuery and useSuspenseQueries
            print(f"Adding import to {file.filepath}")
            file.add_import_from_import_string(import_str)

            print(f"Converting useSuspenseQuery to useSuspenseQueries in {function.name}")
            new_query = f"const [{', '.join(results)}] = useSuspenseQueries({{queries: [{', '.join(queries)}]}})"
            old_statements[0].edit(new_query)
            for s in old_statements[1:]:
                s.remove()

wrap_nested_arrow_functions_with_use_callback

This wraps all nested arrow functions in JSX components with useCallback while constraining it to functions that are passed in as a prop to a JSX element in the current file.

def is_passed_as_prop(component: Function, callback: Function) -> bool:
    """Returns True iff the callback is passed as a prop in the outer component"""
    # Check all JSX elements in the function
    for element in function.jsx_elements:
        # Iterate through their props
        for prop in element.props:
            # See if any props are passed the callback
            if prop.value and prop.value.source == "{" + callback.name + "}":
                return True
    # Return False by default
    return False

for function in codebase.functions:
    # Check if it is a JSX function
    if function.is_jsx:
        # Iterate over all the nested functions in the JSX function
        for child in function.nested_functions:
            # Check if the child function is an arrow function
            if child.is_arrow:
                # Check if child function is ever used as a prop
                if is_passed_as_prop(function, child):
                    print(f"🔵 Callback: {child.name} (<{function.name} /> - {function.file.filepath})")
                    # Wrap the arrow function with useCallback
                    child.insert_before("useCallback(", newline=False, extended=False)
                    child.insert_after(", [])", newline=False)

                    # Add useCallback import if it doesn't exist
                    if not function.file.has_import("useCallback"):
                        function.file.add_import_from_import_string("import { useCallback } from 'react'")

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.