Add visualization code files to analyzers/visualization directory

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",
+]

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

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)

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)