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:
if "@activity" not in file.content:
continue
for function in file.functions:
for decorator in function.decorators:
if "@activity.defn" in decorator.source:
if len(decorator.function_calls) > 0:
call = decorator.function_calls[0]
name_arg = call.get_arg_by_parameter_name("name")
if name_arg:
if isinstance(name_arg.value, String):
constant_name = str(name_arg.value).split(".")[-1].upper().replace("-", "_").replace(" ", "_")
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)
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})")
if not file.has_import(f"{constant_name}"):
file.add_import_from_import_string(f"from .constants import {constant_name}")
name_arg.set_value(constant_name)
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.
decorator_symbol = codebase.get_symbol("decorator_function")
for file in codebase.files:
for function in file.functions:
if function.name.startswith("foo"):
if not file.has_import("decorator_function") and decorator_symbol.file != file:
file.add_symbol_import(decorator_symbol)
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.
decorator_symbol = codebase.get_symbol("my_decorator")
for file in codebase.files:
if not file.has_import(decorator_symbol.name) and decorator_symbol.file != file:
file.add_symbol_import(decorator_symbol)
for function in file.functions:
if function == decorator_symbol:
continue
function.add_decorator(f"@{decorator_symbol.name}", skip_if_exists=True)
for cls in file.classes:
for method in cls.methods:
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 function in codebase.functions:
if function.name.startswith("foo"):
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:
for method in cls.methods:
if any(["@skill_impl" in d.source for d in method.decorators]):
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.get_all_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.get_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.
call_graph_from_node
This skill shows a visualization of the call graph from a given node.
By default, the call graph would use yellow for the starting node, but it could change depending on the user request.
call_paths_between_nodes
This skill generates and visualizes a call graph between two specified functions or symbols.
It starts from a given function/symbol and iteratively traverses through its usages and 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.
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"}
for imported in codebase.imports:
if imported.module == "typing" and imported.name in import_replacements:
imported.remove()
for symbol in imported.get_symbol_usages():
for usage in symbol.find(imported.name, exact=True):
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
for v in codebase.global_vars:
if isinstance(v.value, Dict):
v.set_value(f"Schema({", ".join(f"{k}={v}" for k, v in v.value.items())})")
enum_attribute_mismatch_checker
None
for cls in codebase.classes:
if cls.is_subclass_of("Enum"):
has_mismatched_attributes = False
for attr in cls.attributes:
if attr.name != attr.value:
has_mismatched_attributes = True
break
if has_mismatched_attributes:
print(f"Enum Class: {cls.name} has mismatched attributes.")
enum_by_value_modifier
None
enum_field_cls = codebase.get_file("vendored/enum_field.py").classes[0]
subclasses = [*enum_field_cls.get_subclasses(), enum_field_cls]
for enum_cls in subclasses:
print(f"Processing enum class: {enum_cls.name}")
call_sites = enum_cls.call_sites
for call_site in call_sites[:10]:
args = call_site.args
for arg in args:
print(f"Checking argument: {arg.name} with resolved value: {arg.resolved_value}")
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")
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()
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:
for imp in file.imports:
if imp.name == "Plan":
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 = {}
for cls in codebase.classes:
if cls.is_subclass_of("BaseModel") and "from app.models.base import BaseModel" in cls.file.content:
foreign_key_mapping[cls.name] = []
for attr in cls.attributes:
if "ForeignKey" in attr.source:
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]
if target_table not in foreign_key_mapping[cls.name]:
foreign_key_mapping[cls.name].append(target_table)
G = nx.DiGraph()
for model, targets in foreign_key_mapping.items():
model_name = model.replace("Model", "")
G.add_node(model_name)
for target in targets:
G.add_node(target)
G.add_edge(model_name, target)
codebase.visualize(G)
in_degrees = G.in_degree()
in_degree_list = [(node, degree) for node, degree in in_degrees]
sorted_in_degrees = sorted(in_degree_list, key=lambda x: x[1], reverse=True)
for node, degree in sorted_in_degrees:
print(f"Node: {node}, In-Degree: {degree}")
if degree == 0:
G.nodes[node]["color"] = "red"
models_mapping_to_task = []
for model, targets in foreign_key_mapping.items():
if "Task" in targets:
models_mapping_to_task.append(model)
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 current_depth > 2:
print(f"Depth: {current_depth}:\n_____\n{if_block.source}\n")
if_block.insert_before("# 🚩 🚩 🚩", fix_indentation=True)
for nested_if in if_block.consequence_block.if_blocks:
check_if_depth(nested_if, current_depth + 1)
callables = codebase.functions + [m for c in codebase.classes for m in c.methods]
for function in callables:
for if_block in function.code_block.if_blocks:
check_if_depth(if_block, 1)
import_relationship_graph_visualizer
None
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}
for app in list_apps:
app_name = app.split("/")[2]
G.add_node(app_name, color="red")
for app in list_apps:
directory = codebase.get_directory(app)
for file in directory.files:
for import_statement in file.import_statements:
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)
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.get_all_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
def migrate_import(_import: PyImport, target: str, new: str) -> None:
if _import.module == target:
print(f"Processing: {_import.source}")
_import.set_import_module(new)
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:
- Iterates over all files in the codebase.
- Checks if the active file is a test file and contains the string “space-1”.
- 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.
- 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
for file in codebase.files:
if "dataclasses" not in file.filepath and "dataclasses" not in file.name:
for cls in file.classes:
if "@dataclass" in cls.source:
new_filename = "dataclasses.py"
if not codebase.has_file(new_filename):
dst_file = codebase.create_file(new_filename, "")
else:
dst_file = codebase.get_file(new_filename)
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 file in codebase.files:
if file.name == "enums.py":
continue
enum_classes = [cls for cls in file.classes if cls.is_subclass_of("Enum")]
if enum_classes:
parent_dir = Path(file.filepath).parent
new_filepath = str(parent_dir / "enums.py")
if len(file.symbols) == len(enum_classes):
file.update_filepath(new_filepath)
else:
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 enum_class in enum_classes:
enum_class.move_to_file(dst_file)
move_foo_functions_skill
Moves all functions starting with ‘foo’ to a file named foo
foo_file = codebase.get_file("foo.py") if codebase.has_file("foo.py") else codebase.create_file("foo.py")
for function in codebase.functions:
if function.name.startswith("foo"):
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:
-
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.
-
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.
-
It then iterates over all call sites of the function, accumulating arguments that match a dictionary data type.
-
The keys in all these dictionary arguments are analyzed to identify unique keys and their frequency.
-
Each unique key across all the dictionaries is then typed by inferring the value type associated with the key in the different call sites.
-
A @dataclass definition is then created with attributes that correspond to the unique keys identified and their inferred or overridden types.
-
This @dataclass is injected into the source code, directly after the last import line.
-
The function parameter type is adjusted to use this new @dataclass.
-
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"
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()
matching_args = []
for call_site in function.call_sites:
arg = call_site.get_arg_by_parameter_name(PARAM_NAME)
if arg and isinstance(arg.value, Dict):
matching_args.append(arg)
print("\n### 🔬 Analytics ###")
print(f"📞 Found {len(function.call_sites)} call-sites")
print(f"🔎 Found {len(matching_args)} dict args")
def flatten(level):
return [y for x in level for y in x]
all_keys = flatten([x.value.keys() for x in matching_args])
unique_keys = set(all_keys)
print("\n### 📊 Frequency of keys: ###")
import pandas as pd
print(pd.Series(all_keys).value_counts())
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"
elif "Boolean" in str(type(value)):
return "bool"
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 ###")
file = function.file
last_import = file.imports[-1]
last_import.insert_after(f"\n\n{dataclass_source}")
print(" ✅ Added Dataclass definition")
if "| None" in parameter.source:
parameter.set_type_annotation(dataclass_name + " | None")
else:
parameter.set_type_annotation(dataclass_name)
print(" ✅ Updated dataclass type signature")
import_line = function.get_import_string()
import_line = f"from {function.file.get_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
for call_site in function.call_sites:
arg = call_site.get_arg_by_parameter_name(parameter.name)
if arg:
arg.set_value(make_schema(arg, call_site))
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.get_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.get_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
co_occurrence_matrix = defaultdict(lambda: defaultdict(int))
for function in codebase.functions:
return_type = function.return_type
if isinstance(return_type, UnionType):
options = [x.source for x in return_type]
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("🪐 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()
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 = {}
python_count = 0
typescript_count = 0
for cls in codebase.classes:
for method in cls.methods:
for decorator in method.decorators:
if "@skill_impl" in decorator.source:
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(f"Total Skills: {python_count + typescript_count}")
print(f"🐍 Python Skills: {python_count}")
print(f"🟦 TypeScript Skills: {typescript_count}")
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:
for assignment in attr.assignments:
if isinstance(assignment.value, FunctionCall) and "db.relationship" in assignment.value:
for arg in assignment.value.args:
if arg.name == "uselist":
relationships[f"{cls.name}.{attr.name}"] = {"uselist": arg.value.source}
break
else:
relationships[f"{cls.name}.{attr.name}"] = {"uselist": "False"}
print("#####[ Relationships ]#####")
print(relationships)
for file in codebase.files:
added = False
for call in dict.fromkeys(file.function_calls):
if call.name in ("selectinload", "joinedload"):
relationship = call.args[0].source
fn = call.name
if relationship not in relationships:
print("unknown", relationship)
continue
if fn == "joinedload" and relationships[relationship]["uselist"] == "True":
added = True
call.set_name("selectinload")
if added:
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")
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]",
}
directory = codebase.get_directory(DIRECTORY)
for cls in directory.classes:
for attribute in cls.attributes:
if attribute.assignment.type:
continue
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]
if is_nullable:
first_argument = f"Optional[{first_argument}]"
if first_argument not in column_type_to_mapped_type.keys():
continue
new_type = column_type_to_mapped_type[first_argument]
attribute.assignment.set_type_annotation(new_type)
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.
- Finds all temporal “activities” (functions decorated with @activity.defn(…))
- 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:
if "@activity" not in file.content:
continue
for function in file.functions:
for decorator in function.decorators:
if "@activity.defn" in decorator.source:
if len(decorator.function_calls) > 0:
call = decorator.function_calls[0]
name_arg = call.get_arg_by_parameter_name("name")
if name_arg:
if '"backfill' in name_arg.value:
name_arg.value.replace('"backfill', '"my_directory.backfill')
transaction_canonical_refactorer
None
get_merchant_name_for_transaction = codebase.get_symbol("get_merchant_name_for_transaction")
for file in codebase.files:
for function in file.functions:
for parameter in function.parameters:
if parameter.type and parameter.type.source == "TransactionCanonical":
matches = function.code_block.search(f"{parameter.name}.merchant.merchant_name")
for match in matches:
print(f"🔵 Patching: {function.name} ({function.file.filepath})")
match.edit(f"get_merchant_name_for_transaction({parameter.name})")
if len(matches) > 0:
function.file.add_symbol_import(get_merchant_name_for_transaction)
for assignment in function.code_block.local_var_assignments:
if assignment.name == "transaction_canonical":
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)")
match.edit(f"get_merchant_name_for_transaction({assignment.name})")
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 function in codebase.functions:
for with_statement in function.code_block.with_statements:
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.
optional_type_pattern = re.compile(r"Optional\[(.*?)]")
def update_optional_parameter_type_hints(function: PyFunction):
for parameter in function.parameters:
if parameter.is_typed:
old_type = parameter.type
if "Optional[" in old_type:
new_type = optional_type_pattern.sub(r"\1 | None", old_type)
parameter.set_type_annotation(new_type)
def update_optional_return_type_hints(function: PyFunction):
if function.return_type:
old_return_type = function.return_type.source
if "Optional[" in old_return_type:
new_return_type = optional_type_pattern.sub(r"\1 | None", old_return_type)
function.return_type.edit(new_return_type)
for function in codebase.functions:
update_optional_parameter_type_hints(function)
update_optional_return_type_hints(function)
for cls in codebase.classes:
for method in cls.methods:
update_optional_parameter_type_hints(method)
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:
-
Iterates over all files in the codebase, specifically focusing on files within the ‘src’ directory and ignoring
files like ‘.eslint.rc’.
-
For each file, it checks if ‘useSuspenseQuery’ is present. If it isn’t, the file is skipped.
-
If ‘useSuspenseQuery’ is found, an import statement for ‘useQuery’ and ‘useSuspenseQueries’ from the
@tanstack/react-query library is added to the file.
-
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’.
-
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.
-
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.
