# src/codegraphcontext/tools/languages/python.py

import os

import tempfile

import nbformat

from nbconvert import PythonExporter

from pathlib import Path

from typing import Any, Dict, Optional, Tuple

import ast

import logging

import warnings

from codegraphcontext.utils.debug_log import debug_log, info_logger, error_logger, warning_logger, debug_logger

from codegraphcontext.utils.tree_sitter_manager import execute_query

# Suppress verbose traitlets/nbconvert DEBUG logs

logging.getLogger('traitlets').setLevel(logging.WARNING)

logging.getLogger('nbconvert').setLevel(logging.WARNING)

# Suppress IPython UserWarning from nbconvert

warnings.filterwarnings('ignore', message='.*IPython is needed to transform IPython syntax.*')

PY_QUERIES = {

"imports": """

(import_statement name: (_) @import)

(import_from_statement) @from_import_stmt

""",

"classes": """

(class_definition

name: (identifier) @name

superclasses: (argument_list)? @superclasses

body: (block) @body)

""",

"functions": """

(function_definition

name: (identifier) @name

parameters: (parameters) @parameters

body: (block) @body

return_type: (_)? @return_type)

""",

"calls": """

(call

function: (identifier) @name)

(call

function: (attribute attribute: (identifier) @name) @full_call)

""",

"variables": """

(assignment

left: (identifier) @name)

""",

"lambda_assignments": """

(assignment

left: (identifier) @name

right: (lambda) @lambda_node)

""",

"docstrings": """

(expression_statement (string) @docstring)

""",

"dict_method_refs": """

(dictionary

(pair

key: (_) @key

value: (attribute) @method_ref))

""",

}

class PythonTreeSitterParser:

"""A Python-specific parser using tree-sitter, encapsulating language-specific logic."""

def __init__(self, generic_parser_wrapper):

self.generic_parser_wrapper = generic_parser_wrapper

self.language_name = generic_parser_wrapper.language_name

self.language = generic_parser_wrapper.language

self.parser = generic_parser_wrapper.parser

def _get_node_text(self, node) -> str:

return node.text.decode('utf-8')

def _get_parent_context(self, node, types=('function_definition', 'class_definition')):

curr = node.parent

while curr:

if curr.type in types:

name_node = curr.child_by_field_name('name')

return self._get_node_text(name_node) if name_node else None, curr.type, curr.start_point[0] + 1

curr = curr.parent

return None, None, None

def _calculate_complexity(self, node):

from codegraphcontext.tools.indexing.constants import MAX_AST_DEPTH

complexity_nodes = {

"if_statement", "for_statement", "while_statement", "except_clause",

"with_statement", "boolean_operator", "list_comprehension",

"generator_expression", "case_clause"

}

count = 1

skipped = False

def traverse(n, depth=0):

nonlocal count, skipped

if depth > MAX_AST_DEPTH:

skipped = True

return

if n.type in complexity_nodes:

count += 1

for child in n.children:

traverse(child, depth + 1)

traverse(node)

if skipped:

warning_logger(

f"AST depth exceeded {MAX_AST_DEPTH} levels; "

"complexity count may be underestimated."

)

return count

def _get_docstring(self, body_node):

if body_node and body_node.child_count > 0:

first_child = body_node.children[0]

if first_child.type == 'expression_statement' and first_child.children[0].type == 'string':

try:

return ast.literal_eval(self._get_node_text(first_child.children[0]))

except (ValueError, SyntaxError):

return self._get_node_text(first_child.children[0])

return None

def parse(self, path: Path, is_dependency: bool = False, is_notebook: bool = False, index_source: bool = False) -> Dict:

"""Parses a file and returns its structure in a standardized dictionary format."""

original_file_path = path

temp_py_file = None

source_code = None

self.index_source = index_source

try:

if is_notebook:

info_logger(f"Converting notebook {path} to temporary Python file.")

with open(path, 'r', encoding='utf-8') as f:

notebook_node = nbformat.read(f, as_version=4)

exporter = PythonExporter()

python_code, _ = exporter.from_notebook_node(notebook_node)

with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.py', encoding='utf-8') as tf:

tf.write(python_code)

temp_py_file = Path(tf.name)

# The file to be parsed is now the temporary file

path = temp_py_file

with open(path, "r", encoding="utf-8") as f:

source_code = f.read()

tree = self.parser.parse(bytes(source_code, "utf8"))

root_node = tree.root_node

functions = self._find_functions(root_node)

functions.extend(self._find_lambda_assignments(root_node, index_source))

classes = self._find_classes(root_node)

imports = self._find_imports(root_node)

function_calls = self._find_calls(root_node)

self._attach_module_context(functions, function_calls, root_node, index_source)

variables = self._find_variables(root_node)

return {

"path": str(original_file_path), # Always return the original path

"functions": functions,

"classes": classes,

"variables": variables,

"imports": imports,

"function_calls": function_calls,

"is_dependency": is_dependency,

"lang": self.language_name,

}

except Exception as e:

error_logger(f"Failed to parse {original_file_path}: {e}")

return {"path": str(original_file_path), "error": str(e)}

finally:

if temp_py_file and temp_py_file.exists():

os.remove(temp_py_file)

info_logger(f"Removed temporary file: {temp_py_file}")

def _attach_module_context(self, functions, function_calls, root_node, index_source: bool = False):

"""Represent module-level executable code as Python's <module> frame."""

module_level_calls = [

call for call in function_calls

if not call.get("context") or call["context"][0] is None

]

if not module_level_calls:

return

has_module_frame = any(

func.get("name") == "<module>" and func.get("line_number") == 1

for func in functions

)

if not has_module_frame:

module_func = {

"name": "<module>",

"line_number": 1,

"end_line": root_node.end_point[0] + 1,

"args": [],

"cyclomatic_complexity": 1,

"context": None,

"context_type": "module",

"class_context": None,

"decorators": [],

"lang": self.language_name,

"is_dependency": False,

}

if index_source:

module_func["source"] = self._get_node_text(root_node)

module_func["docstring"] = self._get_docstring(root_node)

functions.append(module_func)

for call in module_level_calls:

call["context"] = ("<module>", "module", 1)

def _find_lambda_assignments(self, root_node, index_source: bool = False):

functions = []

query_str = PY_QUERIES.get('lambda_assignments')

if not query_str: return []

for match in execute_query(self.language, query_str, root_node):

capture_name = match[1]

node = match[0]

if capture_name == 'name':

assignment_node = node.parent

lambda_node = assignment_node.child_by_field_name('right')

name = self._get_node_text(node)

params_node = lambda_node.child_by_field_name('parameters')

context, context_type, _ = self._get_parent_context(assignment_node)

class_context, _, _ = self._get_parent_context(assignment_node, types=('class_definition',))

func_data = {

"name": name,

"line_number": node.start_point[0] + 1,

"end_line": assignment_node.end_point[0] + 1,

"args": [p for p in [self._get_node_text(p) for p in params_node.children if p.type == 'identifier'] if p] if params_node else [],

"cyclomatic_complexity": 1,

"context": context,

"context_type": context_type,

"class_context": class_context,

"decorators": [],

"lang": self.language_name,

"is_dependency": False,

}

if self.index_source:

func_data["source"] = self._get_node_text(assignment_node)

func_data["docstring"] = None

functions.append(func_data)

return functions

def _find_functions(self, root_node, index_source: bool = False):

functions = []

query_str = PY_QUERIES['functions']

for match in execute_query(self.language, query_str, root_node):

capture_name = match[1]

node = match[0]

if capture_name == 'name':

func_node = node.parent

name = self._get_node_text(node)

params_node = func_node.child_by_field_name('parameters')

body_node = func_node.child_by_field_name('body')

# In tree-sitter-python, decorators are children of decorated_definition, which is the parent of function_definition

decorators = []

if func_node.parent and func_node.parent.type == 'decorated_definition':

decorators = [self._get_node_text(child) for child in func_node.parent.children if child.type == 'decorator']

else:

decorators = [self._get_node_text(child) for child in func_node.children if child.type == 'decorator']

context, context_type, _ = self._get_parent_context(func_node)

class_context, _, _ = self._get_parent_context(func_node, types=('class_definition',))

args = []

if params_node:

for p in params_node.children:

arg_text = None

if p.type == 'identifier':

# Simple parameter: def foo(x)

arg_text = self._get_node_text(p)

elif p.type == 'default_parameter':

# Parameter with default: def foo(x=5)

name_node = p.child_by_field_name('name')

if name_node:

arg_text = self._get_node_text(name_node)

elif p.type == 'typed_parameter':

# Typed parameter: def foo(x: int)

name_node = p.child_by_field_name('name')

if name_node:

arg_text = self._get_node_text(name_node)

elif p.type == 'typed_default_parameter':

# Typed parameter with default: def foo(x: int = 5) or def foo(x: str = typer.Argument(...))

name_node = p.child_by_field_name('name')

if name_node:

arg_text = self._get_node_text(name_node)

elif p.type == 'list_splat_pattern' or p.type == 'dictionary_splat_pattern':

# *args or **kwargs

arg_text = self._get_node_text(p)

if arg_text:

args.append(arg_text)

func_data = {

"name": name,

"line_number": node.start_point[0] + 1,

"end_line": func_node.end_point[0] + 1,

"args": args,

"cyclomatic_complexity": self._calculate_complexity(func_node),

"context": context,

"context_type": context_type,

"class_context": class_context,

"decorators": [d for d in decorators if d],

"lang": self.language_name,

"is_dependency": False,

}

if self.index_source:

func_data["source"] = self._get_node_text(func_node)

func_data["docstring"] = self._get_docstring(body_node)

functions.append(func_data)

return functions

def _find_classes(self, root_node, index_source: bool = False):

classes = []

query_str = PY_QUERIES['classes']

for match in execute_query(self.language, query_str, root_node):

capture_name = match[1]

node = match[0]

if capture_name == 'name':

class_node = node.parent

name = self._get_node_text(node)

body_node = class_node.child_by_field_name('body')

superclasses_node = class_node.child_by_field_name('superclasses')

bases = []

if superclasses_node:

bases = [self._get_node_text(child) for child in superclasses_node.children if child.type in ('identifier', 'attribute')]

decorators = [self._get_node_text(child) for child in class_node.children if child.type == 'decorator']

context, _, _ = self._get_parent_context(class_node)

class_data = {

"name": name,

"line_number": node.start_point[0] + 1,

"end_line": class_node.end_point[0] + 1,

"bases": [b for b in bases if b],

"context": context,

"decorators": [d for d in decorators if d],

"lang": self.language_name,

"is_dependency": False,

}

if self.index_source:

class_data["source"] = self._get_node_text(class_node)

class_data["docstring"] = self._get_docstring(body_node)

classes.append(class_data)

return classes

def _find_imports(self, root_node):

imports = []

seen_modules = set()

query_str = PY_QUERIES['imports']

for node, capture_name in execute_query(self.language, query_str, root_node):

if capture_name in ('import', 'from_import_stmt'):

# For 'import_statement'

if capture_name == 'import':

node_text = self._get_node_text(node)

alias = None

if ' as ' in node_text:

parts = node_text.split(' as ')

full_name = parts[0].strip()

alias = parts[1].strip()

else:

full_name = node_text.strip()

if full_name in seen_modules:

continue

seen_modules.add(full_name)

import_data = {

"name": full_name,

"full_import_name": full_name,

"line_number": node.start_point[0] + 1,

"alias": alias,

"context": self._get_parent_context(node)[:2],

"lang": self.language_name,

"is_dependency": False,

}

imports.append(import_data)

# For 'import_from_statement'

elif capture_name == 'from_import_stmt':

module_name_node = node.child_by_field_name('module_name')

if not module_name_node: continue

module_name = self._get_node_text(module_name_node)

# Handle 'from ... import ...'

import_list_node = node.child_by_field_name('name')

if import_list_node:

for child in import_list_node.children:

imported_name = None

alias = None

if child.type == 'aliased_import':

name_node = child.child_by_field_name('name')

alias_node = child.child_by_field_name('alias')

if name_node: imported_name = self._get_node_text(name_node)

if alias_node: alias = self._get_node_text(alias_node)

elif child.type == 'dotted_name' or child.type == 'identifier':

imported_name = self._get_node_text(child)

if imported_name:

full_import_name = f"{module_name}.{imported_name}"

if full_import_name in seen_modules:

continue

seen_modules.add(full_import_name)

imports.append({

"name": imported_name,

"full_import_name": full_import_name,

"line_number": child.start_point[0] + 1,

"alias": alias,

"context": self._get_parent_context(child)[:2],

"lang": self.language_name,

"is_dependency": False,

})

return imports

def _extract_call_name(self, function_node) -> str:

if function_node.type == "identifier":

return self._get_node_text(function_node)

if function_node.type == "attribute":

attr = function_node.child_by_field_name("attribute")

if attr:

return self._get_node_text(attr)

return self._get_node_text(function_node).split(".")[-1]

def _call_args(self, call_node) -> list:

args = []

arguments_node = call_node.child_by_field_name("arguments")

if arguments_node:

for arg in arguments_node.children:

arg_text = self._get_node_text(arg)

if arg_text and arg_text not in ("(", ")", ","):

args.append(arg_text)

return args

def _record_call(self, call_node, enclosing_caller, calls):

function_node = call_node.child_by_field_name("function")

if not function_node:

return None

called_name = self._extract_call_name(function_node)

if enclosing_caller:

context = (enclosing_caller, "nested_call", call_node.start_point[0] + 1)

else:

context = self._get_parent_context(call_node)

calls.append(

{

"name": called_name,

"full_name": self._get_node_text(function_node),

"line_number": call_node.start_point[0] + 1,

"args": self._call_args(call_node),

"inferred_obj_type": None,

"context": context,

"class_context": self._get_parent_context(

call_node, types=("class_definition",)

)[:2],

"lang": self.language_name,

"is_dependency": False,

}

)

return called_name

def _walk_call_tree(self, node, enclosing_caller, calls):

if node is None:

return

if node.type == "call":

called_name = self._record_call(node, enclosing_caller, calls)

if called_name:

arguments_node = node.child_by_field_name("arguments")

if arguments_node:

for child in arguments_node.children:

self._walk_call_tree(child, called_name, calls)

return

for child in node.children:

self._walk_call_tree(child, enclosing_caller, calls)

def _find_calls(self, root_node):

calls = []

self._walk_call_tree(root_node, None, calls)

# Dictionary-based method references (indirect calls)

dict_method_calls = self._find_dict_method_references(root_node)

calls.extend(dict_method_calls)

return calls

def _find_dict_method_references(self, root_node):

"""

Detects indirect function calls through dictionary mappings.

Example pattern:

tool_map = {

"add_code": self.add_code_to_graph_tool,

"find_code": self.find_code_tool,

}

handler = tool_map.get(tool_name)

if handler:

handler(**args)

This creates CALLS relationships from the context function to all

methods referenced in the dictionary.

"""

calls = []

query_str = PY_QUERIES.get('dict_method_refs')

if not query_str:

return calls

# Track dictionaries that contain method references

dict_assignments = {} # dict_var_name -> list of method references

for node, capture_name in execute_query(self.language, query_str, root_node):

if capture_name == 'method_ref':

# Found a method reference in a dictionary value

# Navigate up to find the assignment

dict_node = node.parent # pair node

while dict_node and dict_node.type != 'dictionary':

dict_node = dict_node.parent

if dict_node:

# Find the assignment node

assignment_node = dict_node.parent

if assignment_node and assignment_node.type == 'assignment':

# Get the variable name being assigned

left_node = assignment_node.child_by_field_name('left')

if left_node:

var_name = self._get_node_text(left_node)

method_ref = self._get_node_text(node)

# Extract just the method name (remove 'self.')

method_name = method_ref.split('.')[-1] if '.' in method_ref else method_ref

if var_name not in dict_assignments:

dict_assignments[var_name] = {

'methods': [],

'context': self._get_parent_context(assignment_node),

'line_number': assignment_node.start_point[0] + 1

}

dict_assignments[var_name]['methods'].append({

'name': method_name,

'full_name': method_ref,

'line_number': node.start_point[0] + 1

})

# Now create call relationships for each method in the dictionaries

# The context is the function where the dictionary is defined

for dict_var, data in dict_assignments.items():

context, context_type, context_line = data['context']

class_context, _, _ = (None, None, None)

for method_info in data['methods']:

call_data = {

"name": method_info['name'],

"full_name": method_info['full_name'],

"line_number": method_info['line_number'],

"args": [], # We don't know the args at this point

"inferred_obj_type": None,

"context": (context, context_type, context_line),

"class_context": (class_context, None),

"lang": self.language_name,

"is_dependency": False,

"is_indirect_call": True, # Mark as indirect for debugging

}

calls.append(call_data)

return calls

def _find_variables(self, root_node):

variables = []

query_str = PY_QUERIES['variables']

for match in execute_query(self.language, query_str, root_node):

capture_name = match[1]

node = match[0]

if capture_name == 'name':

assignment_node = node.parent

# Skip lambda assignments, they are handled by _find_lambda_assignments

right_node = assignment_node.child_by_field_name('right')

if right_node and right_node.type == 'lambda':

continue

name = self._get_node_text(node)

value = self._get_node_text(right_node) if right_node else None

type_node = assignment_node.child_by_field_name('type')

type_text = self._get_node_text(type_node) if type_node else None

context, _, _ = self._get_parent_context(node)

class_context, _, _ = self._get_parent_context(node, types=('class_definition',))

variable_data = {

"name": name,

"line_number": node.start_point[0] + 1,

"value": value,

"type": type_text,

"context": context,

"class_context": class_context,

"lang": self.language_name,

"is_dependency": False,

}

variables.append(variable_data)

return variables

def pre_scan_python(files: list[Path], parser_wrapper) -> dict:

"""Scans Python files to create a map of class/function names to their file paths."""

imports_map = {}

query_str = """

(class_definition name: (identifier) @name)

(function_definition name: (identifier) @name)

"""

for path in files:

temp_py_file = None

try:

source_to_parse = ""

if path.suffix == '.ipynb':

with open(path, 'r', encoding='utf-8') as f:

notebook_node = nbformat.read(f, as_version=4)

exporter = PythonExporter()

python_code, _ = exporter.from_notebook_node(notebook_node)

with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.py', encoding='utf-8') as tf:

tf.write(python_code)

temp_py_file = Path(tf.name)

with open(temp_py_file, "r", encoding="utf-8") as f:

source_to_parse = f.read()

else:

with open(path, "r", encoding="utf-8") as f:

source_to_parse = f.read()

tree = parser_wrapper.parser.parse(bytes(source_to_parse, "utf8"))

for capture, _ in execute_query(parser_wrapper.language, query_str, tree.root_node):

name = capture.text.decode('utf-8')

if name not in imports_map:

imports_map[name] = []

imports_map[name].append(str(path.resolve()))

except Exception as e:

warning_logger(f"Tree-sitter pre-scan failed for {path}: {e}")

finally:

if temp_py_file and temp_py_file.exists():

os.remove(temp_py_file)

return imports_map