diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py index e69de29bb..28f5a9e30 100644 --- a/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py +++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py @@ -0,0 +1,28 @@ +# Import visualization modules +from codegen_on_oss.analyzers.visualization.organize import ( + MoveSymbolDemonstration, + MoveSymbolsWithDependencies, + MoveSymbolToFileWithDependencies, + MoveSymbolWithAddBackEdgeStrategy, + MoveSymbolWithUpdatedImports, + SplitFunctionsIntoSeparateFiles, +) +from codegen_on_oss.analyzers.visualization.viz_call_graph import ( + CallGraphFilter, + CallGraphFromNode, + CallPathsBetweenNodes, +) +from codegen_on_oss.analyzers.visualization.viz_dead_code import DeadCode + +__all__ = [ + "CallGraphFilter", + "CallGraphFromNode", + "CallPathsBetweenNodes", + "DeadCode", + "MoveSymbolDemonstration", + "MoveSymbolToFileWithDependencies", + "MoveSymbolWithAddBackEdgeStrategy", + "MoveSymbolWithUpdatedImports", + "MoveSymbolsWithDependencies", + "SplitFunctionsIntoSeparateFiles", +] diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/organize.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/organize.py new file mode 100644 index 000000000..adf074f58 --- /dev/null +++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/organize.py @@ -0,0 +1,553 @@ +from abc import ABC +from typing import Any, Dict, List, Optional, Union, Callable as PyCallable + +from codegen.sdk.core.codebase import CodebaseType, TSCodebaseType +from codegen.shared.enums.programming_language import ProgrammingLanguage + +from tests.shared.skills.decorators import skill, skill_impl +from tests.shared.skills.skill import Skill +from tests.shared.skills.skill_test import ( + SkillTestCase, + SkillTestCasePyFile, + SkillTestCaseTSFile, +) + +SplitFunctionsIntoSeparateFilesPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +NON_FUNCTION = 'This is not a function' +def function1(): + print("This is function 1") + +def function2(): + print("This is function 2") + +def function3(): + print("This is function 3") +""", + output=""" + NON_FUNCTION = 'This is not a function' +""", + filepath="path/to/file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def function1(): + print("This is function 1") +""", + filepath="function1.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def function2(): + print("This is function 2") +""", + filepath="function2.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def function3(): + print("This is function 3") +""", + filepath="function3.py", + ), +]) + + +@skill( + prompt="""Generate a code snippet that retrieves a Python file from a codebase, iterates through its functions, + creates a new file for each function using the function's name, and moves the function to the newly created file.""", + guide=True, + uid="5cead96b-7922-49db-b6dd-d48fb51680d2", +) +class SplitFunctionsIntoSeparateFiles(Skill, ABC): + """This code snippet retrieves a Python file from the codebase and iterates through its functions. For each + function, it creates a new file named after the function and moves the function's definition to the newly created + file. + """ + + @staticmethod + @skill_impl(test_cases=[SplitFunctionsIntoSeparateFilesPyTestCase], language=ProgrammingLanguage.PYTHON) + def skill_func(codebase: CodebaseType) -> None: + # Retrieve the Python file from the codebase + file = codebase.get_file("path/to/file.py") + # Iterate through the functions in the file + for function in file.functions: + # Create a new file for each function using the function's name + new_file = codebase.create_file(function.name + ".py") + # Move the function to the newly created file + function.move_to_file(new_file) + + +MoveSymbolDemonstrationPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +def my_function(): + print("This is my function") + +def another_function(): + my_function() +""", + output=""" +from path.to.dst.location import my_function + +def another_function(): + my_function() +""", + filepath="path/to/source_file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def my_function(): + print("This is my function") +""", + filepath="path/to/dst/location.py", + ), +]) + +MoveSymbolDemonstrationTSTestCase = SkillTestCase([ + SkillTestCaseTSFile( + input=""" +function myFunction() { + console.log("This is my function"); +} + +function anotherFunction() { + myFunction(); +} +""", + output=""" +import { myFunction } from 'path/to/dst/location'; + +function anotherFunction() { + myFunction(); +} +""", + filepath="path/to/source_file.ts", + ), + SkillTestCaseTSFile( + input="", + output=""" +export function myFunction() { + console.log("This is my function"); +} +""", + filepath="path/to/dst/location.ts", + ), +]) + + +@skill( + prompt="Generate a code snippet that demonstrates how to move a symbol from one file to another in a codebase.", + guide=True, + uid="1f0182b7-d3c6-4cde-8ffd-d1bbe31e51be", +) +class MoveSymbolDemonstration(Skill, ABC): + """This code snippet demonstrates how to move a symbol from one file to another in a codebase.""" + + @staticmethod + @skill_impl(test_cases=[MoveSymbolDemonstrationPyTestCase], language=ProgrammingLanguage.PYTHON) + def python_skill_func(codebase: CodebaseType) -> None: + source_file = codebase.get_file("path/to/source_file.py") + # =====[ Code Snippet ]===== + # Get the symbol + symbol_to_move = source_file.get_symbol("my_function") + # Pick a destination file + dst_file = codebase.get_file("path/to/dst/location.py") + # Move the symbol, move all of its dependencies with it (remove from old file), and add an import of symbol into old file + symbol_to_move.move_to_file(dst_file, include_dependencies=True, strategy="add_back_edge") + + @staticmethod + @skill_impl(test_cases=[MoveSymbolDemonstrationTSTestCase], language=ProgrammingLanguage.TYPESCRIPT) + def typescript_skill_func(codebase: CodebaseType) -> None: + source_file = codebase.get_file("path/to/source_file.ts") + # =====[ Code Snippet ]===== + # Get the symbol + symbol_to_move = source_file.get_symbol("myFunction") + # Pick a destination file + dst_file = codebase.get_file("path/to/dst/location.ts") + # Move the symbol, move all of its dependencies with it (remove from old file), and add an import of symbol into old file + symbol_to_move.move_to_file(dst_file, include_dependencies=True, strategy="add_back_edge") + + +MoveSymbolWithUpdatedImportsPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +def symbol_to_move(): + print("This symbol will be moved") + +def use_symbol(): + symbol_to_move() +""", + output=""" +from new_file import symbol_to_move + +def use_symbol(): + symbol_to_move() +""", + filepath="original_file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def symbol_to_move(): + print("This symbol will be moved") +""", + filepath="new_file.py", + ), +]) + +MoveSymbolWithUpdatedImportsTSTestCase = SkillTestCase([ + SkillTestCaseTSFile( + input=""" +function symbolToMove() { + console.log("This symbol will be moved"); +} + +function useSymbol() { + symbolToMove(); +} +""", + output=""" +import { symbolToMove } from 'new_file'; + +function useSymbol() { + symbolToMove(); +} +""", + filepath="original_file.ts", + ), + SkillTestCaseTSFile( + input="", + output=""" +export function symbolToMove() { + console.log("This symbol will be moved"); +} +""", + filepath="new_file.ts", + ), +]) + + +@skill( + prompt="""Generate a code snippet that demonstrates how to use a method called `move_to_file` on an object named + `symbol_to_move`. The method should take two parameters: `dest_file`, which represents the destination file path, + and `strategy`, which should be set to the string value "update_all_imports.""", + guide=True, + uid="d24a61b5-212e-4567-87b0-f6ab586b42c1", +) +class MoveSymbolWithUpdatedImports(Skill, ABC): + """Moves the symbol to the specified destination file using the given strategy. The default strategy is to update + all imports. + """ + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolWithUpdatedImportsPyTestCase], + language=ProgrammingLanguage.PYTHON, + ) + def python_skill_func(codebase: CodebaseType) -> None: + symbol_to_move = codebase.get_symbol("symbol_to_move") + dst_file = codebase.create_file("new_file.py") + symbol_to_move.move_to_file(dst_file, strategy="update_all_imports") + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolWithUpdatedImportsTSTestCase], + language=ProgrammingLanguage.TYPESCRIPT, + ) + def typescript_skill_func(codebase: TSCodebaseType) -> None: + symbol_to_move = codebase.get_symbol("symbolToMove") + dst_file = codebase.create_file("new_file.ts") + symbol_to_move.move_to_file(dst_file, strategy="update_all_imports") + + +MoveSymbolWithAddBackEdgeStrategyPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +def symbol_to_move(): + print("This symbol will be moved") + +def use_symbol(): + symbol_to_move() +""", + output=""" +from new_file import symbol_to_move + +def use_symbol(): + symbol_to_move() +""", + filepath="original_file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def symbol_to_move(): + print("This symbol will be moved") +""", + filepath="new_file.py", + ), +]) + +MoveSymbolWithAddBackEdgeStrategyTSTestCase = SkillTestCase([ + SkillTestCaseTSFile( + input=""" +function symbolToMove() { + console.log("This symbol will be moved"); +} + +function useSymbol() { + symbolToMove(); +} +""", + output=""" +import { symbolToMove } from 'new_file'; + +function useSymbol() { + symbolToMove(); +} +""", + filepath="original_file.ts", + ), + SkillTestCaseTSFile( + input="", + output=""" +export function symbolToMove() { + console.log("This symbol will be moved"); +} +""", + filepath="new_file.ts", + ), +]) + + +@skill( + prompt="""Generate a code snippet that calls a method named 'move_to_file' on an object named 'symbol_to_move'. + The method should take two arguments: 'dest_file' and a keyword argument 'strategy' with the value + 'add_back_edge'.""", + guide=True, + uid="f6c21eea-a9f5-4c30-b797-ff8fc3646d00", +) +class MoveSymbolWithAddBackEdgeStrategy(Skill, ABC): + """Moves the symbol to the specified destination file using the given strategy. The default strategy is to add a + back edge during the move. + """ + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolWithAddBackEdgeStrategyPyTestCase], + language=ProgrammingLanguage.PYTHON, + ) + def skill_func(codebase: CodebaseType) -> None: + symbol_to_move = codebase.get_symbol("symbol_to_move") + dst_file = codebase.create_file("new_file.py") + symbol_to_move.move_to_file(dst_file, strategy="add_back_edge") + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolWithAddBackEdgeStrategyTSTestCase], + language=ProgrammingLanguage.TYPESCRIPT, + ) + def typescript_skill_func(codebase: TSCodebaseType) -> None: + symbol_to_move = codebase.get_symbol("symbolToMove") + dst_file = codebase.create_file("new_file.ts") + symbol_to_move.move_to_file(dst_file, strategy="add_back_edge") + + +MoveSymbolToFileWithDependenciesPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +def dependency_function(): + print("I'm a dependency") + +def my_symbol(): + dependency_function() + print("This is my symbol") + +def use_symbol(): + my_symbol() +""", + output=""" +from new_file import my_symbol + +def use_symbol(): + my_symbol() +""", + filepath="original_file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def dependency_function(): + print("I'm a dependency") + +def my_symbol(): + dependency_function() + print("This is my symbol") +""", + filepath="new_file.py", + ), +]) + +MoveSymbolToFileWithDependenciesTSTestCase = SkillTestCase([ + SkillTestCaseTSFile( + input=""" +function dependencyFunction() { + console.log("I'm a dependency"); +} + +function mySymbol() { + dependencyFunction(); + console.log("This is my symbol"); +} + +function useSymbol() { + mySymbol(); +} +""", + output=""" +import { mySymbol } from 'new_file'; + +function useSymbol() { + mySymbol(); +} +""", + filepath="original_file.ts", + ), + SkillTestCaseTSFile( + input="", + output=""" +function dependencyFunction() { + console.log("I'm a dependency"); +} + +export function mySymbol() { + dependencyFunction(); + console.log("This is my symbol"); +} +""", + filepath="new_file.ts", + ), +]) + + +@skill( + prompt="""Generate a code snippet that demonstrates how to use a method called `move_to_file` on an object named + `my_symbol`. The method should take two parameters: `dest_file`, which specifies the destination file, + and `include_dependencies`, which is a boolean parameter set to `True`.""", + guide=True, + uid="0665e746-fa10-4d63-893f-be305202bab2", +) +class MoveSymbolToFileWithDependencies(Skill, ABC): + """Moves the symbol to the specified destination file. + + If include_dependencies is set to True, any dependencies associated with the symbol will also be moved to the + destination file. + """ + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolToFileWithDependenciesPyTestCase], + language=ProgrammingLanguage.PYTHON, + ) + def skill_func(codebase: CodebaseType) -> None: + my_symbol = codebase.get_symbol("my_symbol") + dst_file = codebase.create_file("new_file.py") + my_symbol.move_to_file(dst_file, include_dependencies=True) + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolToFileWithDependenciesTSTestCase], + language=ProgrammingLanguage.TYPESCRIPT, + ) + def typescript_skill_func(codebase: TSCodebaseType) -> None: + my_symbol = codebase.get_symbol("mySymbol") + dst_file = codebase.create_file("new_file.ts") + my_symbol.move_to_file(dst_file, include_dependencies=True) + + +MoveSymbolsWithDependenciesPyTestCase = SkillTestCase([ + SkillTestCasePyFile( + input=""" +def dependency_function(): + print("I'm a dependency") + +def my_function(): + dependency_function() + print("This is my function") + +class MyClass: + def __init__(self): + self.value = dependency_function() + +def use_symbols(): + my_function() + obj = MyClass() +""", + output=""" +from path.to.destination_file import my_function, MyClass + +def use_symbols(): + my_function() + obj = MyClass() +""", + filepath="path/to/source_file.py", + ), + SkillTestCasePyFile( + input="", + output=""" +def dependency_function(): + print("I'm a dependency") + +def my_function(): + dependency_function() + print("This is my function") + +class MyClass: + def __init__(self): + self.value = dependency_function() +""", + filepath="path/to/destination_file.py", + ), +]) + + +@skill( + prompt="""Generate a Python code snippet that creates a list of symbols to move from a source file to a + destination file. The symbols should include a function named 'my_function' and a class named 'MyClass' from the + source file. Then, iterate over the list of symbols and move each symbol to the destination file, ensuring to + include dependencies and update all imports.""", + guide=True, + uid="0895acd3-3788-44a6-8450-d1a5c9cea564", +) +class MoveSymbolsWithDependencies(Skill, ABC): + """Moves specified symbols from the source file to the destination file. + + This code snippet retrieves a function and a class from the source file and stores them in a list. It then + iterates over this list, moving each symbol to the destination file while including dependencies and updating all + imports accordingly. + """ + + @staticmethod + @skill_impl( + test_cases=[MoveSymbolsWithDependenciesPyTestCase], + language=ProgrammingLanguage.PYTHON, + ) + def skill_func(codebase: CodebaseType) -> None: + # Retrieve the source and destination files + source_file = codebase.get_file("path/to/source_file.py") + dest_file = codebase.get_file("path/to/destination_file.py") + # Create a list of symbols to move + symbols_to_move = [ + source_file.get_function("my_function"), + source_file.get_class("MyClass"), + ] + # Move each symbol to the destination file + for symbol in symbols_to_move: + symbol.move_to_file( + dest_file, include_dependencies=True, strategy="update_all_imports" + ) diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_call_graph.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_call_graph.py new file mode 100644 index 000000000..ea8a4da3d --- /dev/null +++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_call_graph.py @@ -0,0 +1,347 @@ +from abc import ABC +from typing import Any, Dict, List, Optional, Union + +import networkx as nx +from codegen.sdk.core.class_definition import Class +from codegen.sdk.core.codebase import CodebaseType +from codegen.sdk.core.detached_symbols.function_call import FunctionCall +from codegen.sdk.core.external_module import ExternalModule +from codegen.sdk.core.function import Function +from codegen.sdk.core.interfaces.callable import Callable +from codegen.shared.enums.programming_language import ProgrammingLanguage + +from tests.shared.skills.decorators import skill, skill_impl +from tests.shared.skills.skill import Skill +from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile + +CallGraphFromNodeTest = SkillTestCase( + [ + SkillTestCasePyFile( + input=""" +def function_to_trace(): + Y() + Z() + +def Y(): + A() + +def Z(): + B() + +def A(): + pass + +def B(): + C() + +def C(): + pass +""", + filepath="example.py", + ) + ], + graph=True, +) + + +@skill( + eval_skill=False, + prompt="Show me a visualization of the call graph from X", + uid="81e8fbb7-a00a-4e74-b9c2-24f79d24d389", +) +class CallGraphFromNode(Skill, ABC): + """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")` + """ + + @staticmethod + @skill_impl(test_cases=[CallGraphFromNodeTest], language=ProgrammingLanguage.PYTHON) + @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT) + def skill_func(codebase: CodebaseType) -> None: + # Create a directed graph + G = nx.DiGraph() + + # ===== [ Whether to Graph External Modules] ===== + GRAPH_EXERNAL_MODULE_CALLS = False + + # ===== [ Maximum Recursive Depth ] ===== + MAX_DEPTH = 5 + + def create_downstream_call_trace(parent: Union[FunctionCall, Function, None] = None, depth: int = 0) -> None: + """Creates call graph for parent + + This function recurses through the call graph of a function and creates a visualization + + Args: + parent (FunctionCallDefinition| Function): The function for which a call graph will be created. + depth (int): The current depth of the recursive stack. + + """ + # if the maximum recursive depth has been exceeded return + if MAX_DEPTH <= depth: + return + if isinstance(parent, FunctionCall): + src_call, src_func = parent, parent.function_definition + else: + src_call, src_func = parent, parent + + # Iterate over all call paths of the symbol + for call in src_func.function_calls: + # the symbol being called + func = call.function_definition + + # ignore direct recursive calls + if func.name == src_func.name: + continue + + # if the function being called is not from an external module + if not isinstance(func, ExternalModule): + # add `call` to the graph and an edge from `src_call` to `call` + G.add_node(call) + G.add_edge(src_call, call) + + # recursive call to function call + create_downstream_call_trace(call, depth + 1) + elif GRAPH_EXERNAL_MODULE_CALLS: + # add `call` to the graph and an edge from `src_call` to `call` + G.add_node(call) + G.add_edge(src_call, call) + + # ===== [ Function To Be Traced] ===== + function_to_trace = codebase.get_function("function_to_trace") + + # Set starting node + G.add_node(function_to_trace, color="yellow") + + # Add all the children (and sub-children) to the graph + create_downstream_call_trace(function_to_trace) + + # Visualize the graph + codebase.visualize(G) + + +CallGraphFilterTest = SkillTestCase( + [ + SkillTestCasePyFile( + input=""" +class MyClass: + def get(self): + self.helper_method() + return "GET request" + + def post(self): + self.helper_method() + return "POST request" + + def patch(self): + return "PATCH request" + + def delete(self): + return "DELETE request" + + def helper_method(self): + pass + + def other_method(self): + self.helper_method() + return "This method should not be included" + +def external_function(): + instance = MyClass() + instance.get() + instance.post() + instance.other_method() +""", + filepath="path/to/file.py", + ), + SkillTestCasePyFile( + input=""" +from path.to.file import MyClass + +def function_to_trace(): + instance = MyClass() + assert instance.get() == "GET request" + assert instance.post() == "POST request" + assert instance.patch() == "PATCH request" + assert instance.delete() == "DELETE request" +""", + filepath="path/to/file1.py", + ), + ], + graph=True, +) + + +@skill( + eval_skill=False, + prompt="Show me a visualization of the call graph from MyClass and filter out test files and include only the methods that have the name post, get, patch, delete", + uid="fc1f3ea0-46e7-460a-88ad-5312d4ca1a12", +) +class CallGraphFilter(Skill, ABC): + """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 class declarations + and includes only methods with specific names (post, get, patch, delete). + 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 + """ + + @staticmethod + @skill_impl(test_cases=[CallGraphFilterTest], language=ProgrammingLanguage.PYTHON) + @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT) + def skill_func(codebase: CodebaseType) -> None: + # Create a directed graph + G = nx.DiGraph() + + # Get the symbol for my_class + func_to_trace = codebase.get_function("function_to_trace") + + # Add the main symbol as a node + G.add_node(func_to_trace, color="red") + + # ===== [ Maximum Recursive Depth ] ===== + MAX_DEPTH = 5 + + SKIP_CLASS_DECLARATIONS = True + + cls = codebase.get_class("MyClass") + + # Define a recursive function to traverse function calls + def create_filtered_downstream_call_trace(parent: Union[FunctionCall, Function], current_depth: int, max_depth: int) -> None: + if current_depth > max_depth: + return + + # if parent is of type Function + if isinstance(parent, Function): + # set both src_call, src_func to parent + src_call, src_func = parent, parent + else: + # get the first callable of parent + src_call, src_func = parent, parent.function_definition + + # Iterate over all call paths of the symbol + for call in src_func.function_calls: + # the symbol being called + func = call.function_definition + + if SKIP_CLASS_DECLARATIONS and isinstance(func, Class): + continue + + # if the function being called is not from an external module and is not defined in a test file + if not isinstance(func, ExternalModule) and not func.file.filepath.startswith("test"): + # add `call` to the graph and an edge from `src_call` to `call` + metadata: Dict[str, Any] = {} + if isinstance(func, Function) and func.is_method and func.name in ["post", "get", "patch", "delete"]: + name = f"{func.parent_class.name}.{func.name}" + metadata = {"color": "yellow", "name": name} + G.add_node(call, **metadata) + G.add_edge(src_call, call, symbol=cls) # Add edge from current to successor + + # Recursively add successors of the current symbol + create_filtered_downstream_call_trace(call, current_depth + 1, max_depth) + + # Start the recursive traversal + create_filtered_downstream_call_trace(func_to_trace, 1, MAX_DEPTH) + + # Visualize the graph + codebase.visualize(G) + + +CallPathsBetweenNodesTest = SkillTestCase( + [ + SkillTestCasePyFile( + input=""" +def start_func(): + intermediate_func() +def intermediate_func(): + end_func() + +def end_func(): + pass +""", + filepath="example.py", + ) + ], + graph=True, +) + + +@skill( + eval_skill=False, + prompt="Show me a visualization of the call paths between start_class and end_class", + uid="aa3f70c3-ac1c-4737-a8b8-7ba89e3c5671", +) +class CallPathsBetweenNodes(Skill, ABC): + """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 + """ + + @staticmethod + @skill_impl( + test_cases=[CallPathsBetweenNodesTest], language=ProgrammingLanguage.PYTHON + ) + @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT) + def skill_func(codebase: CodebaseType) -> None: + # Create a directed graph + G = nx.DiGraph() + + # Get the start and end functions + start_func = codebase.get_function("start_func") + end_func = codebase.get_function("end_func") + + # Add the start and end functions as nodes + G.add_node(start_func, color="green") + G.add_node(end_func, color="red") + + # Create a dictionary to store all functions and their calls + function_calls: Dict[Function, List[Function]] = {} + + # Get all functions in the codebase + all_functions = codebase.get_all_functions() + + # Build the function call graph + for func in all_functions: + function_calls[func] = [] + for call in func.function_calls: + called_func = call.function_definition + if isinstance(called_func, Function): + function_calls[func].append(called_func) + G.add_edge(func, called_func) + + # Find all paths between start_func and end_func + paths = list(nx.all_simple_paths(G, start_func, end_func)) + + # Create a new graph with only the paths between start_func and end_func + path_graph = nx.DiGraph() + path_graph.add_node(start_func, color="green") + path_graph.add_node(end_func, color="red") + + # Add all nodes and edges in the paths + for path in paths: + for i in range(len(path) - 1): + path_graph.add_node(path[i]) + path_graph.add_node(path[i + 1]) + path_graph.add_edge(path[i], path[i + 1]) + + # Visualize the path graph + codebase.visualize(path_graph) diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_dead_code.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_dead_code.py new file mode 100644 index 000000000..b810ae38a --- /dev/null +++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/viz_dead_code.py @@ -0,0 +1,195 @@ +from abc import ABC +from typing import Any, Dict, List, Optional, Set, Union + +import networkx as nx +from codegen.sdk.core.codebase import CodebaseType +from codegen.sdk.core.function import Function +from codegen.sdk.core.import_resolution import Import +from codegen.sdk.core.symbol import Symbol +from codegen.shared.enums.programming_language import ProgrammingLanguage + +from tests.shared.skills.decorators import skill, skill_impl +from tests.shared.skills.skill import Skill +from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile + +PyDeadCodeTest = SkillTestCase( + [ + SkillTestCasePyFile( + input=""" +# Live code +def used_function(): + return "I'm used!" + +class UsedClass: + def used_method(self): + return "I'm a used method!" + +# Dead code +def unused_function(): + return "I'm never called!" + +class UnusedClass: + def unused_method(self): + return "I'm never used!" + +# Second-order dead code +def second_order_dead(): + unused_function() + UnusedClass().unused_method() + +# More live code +def another_used_function(): + return used_function() + +# Main execution +def main(): + print(used_function()) + print(UsedClass().used_method()) + print(another_used_function()) + +if __name__ == "__main__": + main() +""", + filepath="example.py", + ), + SkillTestCasePyFile( + input=""" +# This file should be ignored by the DeadCode skill + +from example import used_function, UsedClass + +def test_used_function(): + assert used_function() == "I'm used!" + +def test_used_class(): + assert UsedClass().used_method() == "I'm a used method!" +""", + filepath="test_example.py", + ), + SkillTestCasePyFile( + input=""" +# This file contains a decorated function that should be ignored + +from functools import lru_cache + +@lru_cache +def cached_function(): + return "I'm cached!" + +# This function is dead code but should be ignored due to decoration +@deprecated +def old_function(): + return "I'm old but decorated!" + +# This function is dead code and should be detected +def real_dead_code(): + return "I'm really dead!" +""", + filepath="decorated_functions.py", + ), + ], + graph=True, +) + + +@skill( + eval_skill=False, + prompt="Show me a visualization of the call graph from my_class and filter out test files and include only the methods that have the name post, get, patch, delete", + uid="ec5e98c9-b57f-43f8-8b3c-af1b30bb91e6", +) +class DeadCode(Skill, ABC): + """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. + """ + + @staticmethod + @skill_impl(test_cases=[PyDeadCodeTest], language=ProgrammingLanguage.PYTHON) + @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT) + def skill_func(codebase: CodebaseType) -> None: + # Create a directed graph + G = nx.DiGraph() + + # Get all functions in the codebase + all_functions = codebase.get_all_functions() + + # Create a set to track used functions + used_functions: Set[Function] = set() + + # Find the entry point function (e.g., main function or any function that's called from outside) + entry_points = [] + for func in all_functions: + # Check if the function is imported elsewhere + if func.usages: + entry_points.append(func) + used_functions.add(func) + + # Recursively mark all functions that are called from entry points + def mark_used_functions(func: Function) -> None: + for call in func.function_calls: + called_func = call.function_definition + if isinstance(called_func, Function) and called_func not in used_functions: + used_functions.add(called_func) + mark_used_functions(called_func) + + # Mark all functions that are called from entry points + for entry_point in entry_points: + mark_used_functions(entry_point) + + # Find dead code (functions that are not used) + dead_functions = [func for func in all_functions if func not in used_functions] + + # Add all functions to the graph + for func in all_functions: + if func in used_functions: + G.add_node(func, color="green", status="used") + else: + G.add_node(func, color="red", status="unused") + + # Add edges for function calls + for func in all_functions: + for call in func.function_calls: + called_func = call.function_definition + if isinstance(called_func, Function): + G.add_edge(func, called_func) + + # Visualize the graph + codebase.visualize(G) + + @staticmethod + def _process_dependencies(dead_code: List[Function], graph: nx.DiGraph) -> None: + """Process dependencies of dead code functions. + + Args: + dead_code: List of functions identified as dead code + graph: NetworkX graph to visualize the dead code + """ + # Identify second-order dead code (functions only called by dead code) + second_order_dead: Set[Function] = set() + + # Check each dead function's calls + for dead_func in dead_code: + for call in dead_func.function_calls: + called_func = call.function_definition + if isinstance(called_func, Function): + # Check if this function is only called by dead code + is_second_order = True + for usage in called_func.symbol_usages: + # If used by a function not in dead_code, it's not second-order dead + if usage.parent_function not in dead_code: + is_second_order = False + break + + if is_second_order and called_func not in dead_code: + second_order_dead.add(called_func) + # Add to graph as second-order dead code + graph.add_node(called_func, color="orange") + + # Add edge to show the call relationship + graph.add_edge(dead_func, called_func)