Call graphs (CG) and control flow graphs (CFG) consist of nodes and edges. CG is interprocedural, nodes represent subroutines (methods, functions, ...), while edges represent the relationship caller-called between two subroutines (e.g., A->B, "A" is the caller subroutine, while "B" is the called subroutine). CFG is intraprocedural, nodes represe…

interprocedural control flow graph (ICFG)

import ast
import jedi
import os


class CodeAnalyzer(ast.NodeVisitor):
    def __init__(self, source_code, project_root, file_path, target_function):
        self.source_code = source_code
        self.file_path = file_path
        self.project_root = project_root
        self.target_function = target_function
        self.function_info = None
        self.class_info = None
        self.script: jedi.Script = jedi.Script(code=source_code, path=file_path)
        self.definitions = []

    def visit_FunctionDef(self, node):
        if node.name == self.target_function:
            self.function_info = {
                'name': node.name,
                'parameters': [(arg.arg, self.get_type_annotation(arg.annotation)) for arg in node.args.args],
                'returns': self.get_type_annotation(node.returns),
                'calls': [],
                'exception_handling': False,
            }
            self.definitions.append(self.function_info)
            self.generic_visit(node)

    def visit_ClassDef(self, node):
        if node.name == self.target_function:
            self.class_info = {
                'type': 'class',
                'name': node.name,
                'methods': [],
                'bases': [self.get_type_annotation(base) for base in node.bases],
            }
            for item in node.body:
                if isinstance(item, ast.FunctionDef):
                    method_info = self.extract_method_info(item)
                    self.class_info['methods'].append(method_info)
            self.definitions.append(self.class_info)
            self.generic_visit(node)

    def extract_method_info(self, node):
        method_info = {
            'name': node.name,
            'parameters': [(arg.arg, self.get_type_annotation(arg.annotation)) for arg in node.args.args],
            'returns': self.get_type_annotation(node.returns),
        }
        return method_info

    def visit_Call(self, node):
        if self.function_info is not None:
            call_name = self._get_call_name(node)
            self.function_info['calls'].append(call_name)
        self.generic_visit(node)

    def visit_ExceptHandler(self, node):
        if self.function_info is not None:
            self.function_info['exception_handling'] = True
        self.generic_visit(node)

    def _get_call_name(self, node):
        if isinstance(node, ast.Call):
            return self._get_call_name(node.func)
        elif isinstance(node, ast.Attribute):
            # Tratando recursivamente a cadeia de atributos até chegar ao primeiro Name node
            value_name = self._get_call_name(node.value)
            return f"{value_name}.{node.attr}"
        elif isinstance(node, ast.Name):
            return node.id
        return "Unknown"

    def get_type_annotation(self, annotation):
        if annotation is None:
            return 'unknown'
        else:
            # Simplified version, considering Python version compatibility
            if hasattr(annotation, 'id'):
                return annotation.id
            elif hasattr(annotation, 'name'):
                return annotation.name  # For ast.Name
            else:
                return 'complex_type'

    def get_type_definition_location(self, type_name):
        definitions = self.script.get_names()
        print(type_name)
        print(definitions)
        for definition in definitions:
            print(definition.name)
            if definition.name == type_name:
                print(
                    f"definition.module_path: {definition.module_path} definition.line: {definition.line} full_name: {definition.full_name}")
                adjusted_module_path = definition.full_name.replace(f".{type_name}", "").replace('.', '/') + '.py'
                analyze_function(self.project_root, f"{self.project_root}/{adjusted_module_path}", type_name)

                return definition.module_path, definition.line
        return None, None

    # def get_type_definition_location(self, type_name):
    #     definitions = self.script.goto_definitions(name=type_name)
    #     if definitions:
    #         definition = definitions[0]
    #         return definition.module_path, definition.line
    #     return None, None


def analyze_function(project_root, file_path, target_function):
    print(f"target_function: {target_function} file_path: {file_path} project_root: {project_root}")
    with open(file_path, 'r') as file:
        source_code = file.read()

    analyzer = CodeAnalyzer(source_code, project_root, file_path, target_function)
    tree = ast.parse(source_code)
    analyzer.visit(tree)

    print(f"analyzer.class_info: {analyzer.class_info}")
    print(f"analyzer.function_info: {analyzer.function_info}")
    print(f"self.definitions: {analyzer.definitions}")
    if analyzer.function_info:
        print(f"analyzer.function_info: {analyzer.function_info}")
        print(analyzer.function_info['parameters'])
        for param, param_type in analyzer.function_info['parameters']:
            if param_type not in ['dict', 'bool', 'int', 'str', 'float', 'complex',
                                  'unknown']:
                module_path, line = analyzer.get_type_definition_location(param_type)
                if 'paths' not in analyzer.function_info:
                    analyzer.function_info['paths'] = {}
                analyzer.function_info['paths'][param] = f"{param_type} (Defined at: {module_path}, line {line})"
    return analyzer.function_info


def format_analysis(info):
    if info is None:
        return "Function not found."

    output = f"Function Name: {info['name']}\nParameters:\n"
    for param, param_type in info['parameters']:
        output += f"  {param}: {param_type} ({info['paths'].get(param)})\n"
    output += f"Returns: {info['returns']}\nCalls:\n"
    for call in info['calls']:
        output += f"  {call}\n"
    output += f"Exception Handling: {info['exception_handling']}\n"
    return output


# Exemplo de uso
project_root = "/Users/joaogabriellima/Documents/Work/Katapult/lms-stack/lms-platform"
target_name = "log_address"
file_path = "lms/apps/api/utils.py"

info = analyze_function(project_root, file_path, target_name)

analysis_output = format_analysis(info)
print("*" * 10)
print(analysis_output)