apisec-code-bolt 0.1.0__py3-none-any.whl

apisec_code_bolt/parsing/python/visitors.py

@@ -0,0 +1,1544 @@
+"""
+LibCST visitors for Python code analysis.
+
+This module contains visitors that traverse Python CST and extract:
+- Function/method definitions
+- Class definitions
+- Import statements
+- Function calls
+- Assignments
+- Decorators
+
+These are low-level extraction visitors. Higher-level semantic analysis
+is done in the parser module.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import libcst as cst
+
+from ...core.types import CodeLocation
+from .expression_utils import (
+    collect_name_nodes,
+    detect_concatenation,
+    detect_container_type,
+    detect_format_call,
+    detect_fstring,
+    extract_value_metadata,
+    is_classvar_annotation,
+)
+
+# Re-export dataclasses so existing ``from .visitors import ...`` still works.
+from .extraction_types import (  # noqa: F401 – re-exports
+    ExtractedArgument,
+    ExtractedAssignment,
+    ExtractedCall,
+    ExtractedClass,
+    ExtractedControlFlowBlock,
+    ExtractedDecorator,
+    ExtractedField,
+    ExtractedFunction,
+    ExtractedImport,
+    ExtractedParameter,
+    ExtractedReturn,
+)
+
+# =============================================================================
+# Position Tracking Mixin
+# =============================================================================
+
+
+class PositionProvider:
+    """Mixin for tracking positions in source code."""
+
+    def __init__(self, source: str) -> None:
+        self._source = source
+        self._line_offsets = self._compute_line_offsets(source)
+
+    @staticmethod
+    def _compute_line_offsets(source: str) -> list[int]:
+        """Compute byte offset of each line."""
+        offsets = [0]
+        for i, char in enumerate(source):
+            if char == "\n":
+                offsets.append(i + 1)
+        return offsets
+
+    def _get_position(self, node: cst.CSTNode) -> tuple[int, int, int, int]:
+        """Get (line, column, end_line, end_column) for a node."""
+        # LibCST tracks positions if we have a wrapper
+        try:
+            pos = self._wrapper.resolve(cst.metadata.PositionProvider)[node]
+            return (
+                pos.start.line,
+                pos.start.column,
+                pos.end.line,
+                pos.end.column,
+            )
+        except (KeyError, AttributeError):
+            return (0, 0, 0, 0)
+
+
+# =============================================================================
+# Main Extraction Visitor
+# =============================================================================
+
+
+class PythonExtractor(cst.CSTVisitor):
+    """
+    Visitor that extracts structural information from Python CST.
+
+    Extracts:
+    - Functions and methods
+    - Classes
+    - Imports
+    - Function calls
+    - Assignments
+    """
+
+    def __init__(
+        self,
+        source: str,
+        file_path: Path | None = None,
+        module_name: str | None = None,
+    ) -> None:
+        self._source = source
+        self._file_path = file_path
+        self._module_name = module_name or ""
+        self._wrapper: cst.MetadataWrapper | None = None
+
+        # Extraction results
+        self.functions: list[ExtractedFunction] = []
+        self.classes: list[ExtractedClass] = []
+        self.imports: list[ExtractedImport] = []
+        self.calls: list[ExtractedCall] = []
+        self.assignments: list[ExtractedAssignment] = []
+        self.module_docstring: str | None = None
+        self.module_variables: list[str] = []
+
+        # Context tracking
+        self._current_class: ExtractedClass | None = None
+        self._current_function: ExtractedFunction | None = None
+        self._class_stack: list[ExtractedClass] = []
+        self._function_stack: list[ExtractedFunction] = []
+
+        # Control flow context tracking
+        self._in_loop: int = 0  # Nesting depth
+        self._in_conditional: int = 0
+        self._in_try: bool = False
+        self._in_except: bool = False
+        self._in_finally: bool = False
+        self._in_with: bool = False
+        self._in_comprehension: bool = False
+
+        # Control flow blocks for current function
+        self._current_control_flow: list[ExtractedControlFlowBlock] = []
+
+        # Position tracking
+        self._positions: dict[int, tuple[int, int, int, int]] = {}
+
+    def set_metadata_wrapper(self, wrapper: cst.MetadataWrapper) -> None:
+        """Set the metadata wrapper for position tracking."""
+        self._wrapper = wrapper
+
+    # =========================================================================
+    # Import Extraction
+    # =========================================================================
+
+    def visit_Import(self, node: cst.Import) -> bool:
+        """Extract 'import x' statements."""
+        for name_item in node.names if isinstance(node.names, Sequence) else []:
+            if isinstance(name_item, cst.ImportAlias):
+                module_name = self._get_dotted_name(name_item.name)
+                alias = None
+                if name_item.asname and isinstance(name_item.asname, cst.AsName):
+                    alias = self._node_to_code(name_item.asname.name)
+
+                self.imports.append(
+                    ExtractedImport(
+                        module=module_name,
+                        names=[(module_name, alias)],
+                        is_from_import=False,
+                        line=self._get_line(node),
+                    )
+                )
+        return False
+
+    def visit_ImportFrom(self, node: cst.ImportFrom) -> bool:
+        """Extract 'from x import y' statements."""
+        # Calculate relative import level
+        relative_level = 0
+        if node.relative:
+            for dot in node.relative:
+                if isinstance(dot, cst.Dot):
+                    relative_level += 1
+
+        # Get module name
+        module = ""
+        if node.module:
+            module = self._get_dotted_name(node.module)
+
+        is_relative = relative_level > 0
+
+        # Get imported names
+        names: list[tuple[str, str | None]] = []
+        if isinstance(node.names, cst.ImportStar):
+            names = [("*", None)]
+        elif isinstance(node.names, Sequence):
+            for name_item in node.names:
+                if isinstance(name_item, cst.ImportAlias):
+                    name = self._get_dotted_name(name_item.name)
+                    alias = None
+                    if name_item.asname and isinstance(name_item.asname, cst.AsName):
+                        alias = self._node_to_code(name_item.asname.name)
+                    names.append((name, alias))
+
+        self.imports.append(
+            ExtractedImport(
+                module=module,
+                names=names,
+                is_from_import=True,
+                is_relative=is_relative,
+                relative_level=relative_level,
+                line=self._get_line(node),
+            )
+        )
+        return False
+
+    # =========================================================================
+    # Function Extraction
+    # =========================================================================
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
+        """Extract function definitions."""
+        func = self._extract_function(node)
+
+        if self._current_class:
+            func.owner_type = self._current_class.name
+            if any(d.name in ("staticmethod", "classmethod") for d in func.decorators):
+                func.binding = "static"
+            else:
+                func.binding = "instance"
+            self._current_class.methods.append(func)
+        else:
+            self.functions.append(func)
+
+        # Push onto stack for nested context
+        self._function_stack.append(func)
+        self._current_function = func
+
+        # Reset control flow tracking for this function
+        self._current_control_flow = []
+
+        return True  # Continue visiting to extract calls/assignments in body
+
+    def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
+        """Leave function definition."""
+        # Store control flow info before leaving
+        if self._current_function and self._current_control_flow:
+            self._current_function.control_flow_info = self._build_control_flow_info()
+
+        if self._function_stack:
+            self._function_stack.pop()
+        self._current_function = self._function_stack[-1] if self._function_stack else None
+
+        # Restore control flow tracking for enclosing function
+        self._current_control_flow = []
+
+    def _build_control_flow_info(self) -> dict[str, Any]:
+        """Build control flow info dictionary from collected blocks."""
+        if_blocks: list[dict[str, Any]] = []
+        loops: list[dict[str, Any]] = []
+        try_blocks: list[dict[str, Any]] = []
+        with_blocks: list[dict[str, Any]] = []
+        comprehensions: list[dict[str, Any]] = []
+
+        for block in self._current_control_flow:
+            if block.block_type == "if":
+                if_blocks.append(
+                    {
+                        "start_line": block.start_line,
+                        "end_line": block.end_line,
+                        "has_elif": block.has_elif,
+                        "has_else": block.has_else,
+                        "elif_lines": block.elif_lines,
+                        "else_line": block.else_line,
+                    }
+                )
+            elif block.block_type in ("for", "while"):
+                loops.append(
+                    {
+                        "loop_type": block.block_type,
+                        "start_line": block.start_line,
+                        "end_line": block.end_line,
+                        "has_break": block.has_break,
+                        "has_continue": block.has_continue,
+                    }
+                )
+            elif block.block_type == "try":
+                try_blocks.append(
+                    {
+                        "try_start": block.start_line,
+                        "try_end": block.end_line,
+                        "except_blocks": [
+                            {"start_line": s, "end_line": e} for s, e in block.except_blocks
+                        ],
+                        "finally_block": (
+                            {
+                                "start_line": block.finally_block[0],
+                                "end_line": block.finally_block[1],
+                            }
+                            if block.finally_block
+                            else None
+                        ),
+                    }
+                )
+            elif block.block_type == "with":
+                with_blocks.append(
+                    {
+                        "start_line": block.start_line,
+                        "end_line": block.end_line,
+                        "context_expr": block.context_expr,
+                        "with_items": block.with_items,
+                    }
+                )
+            elif block.block_type == "comprehension":
+                comprehensions.append(
+                    {
+                        "line": block.start_line,
+                    }
+                )
+
+        return {
+            "if_blocks": if_blocks,
+            "loops": loops,
+            "try_blocks": try_blocks,
+            "with_blocks": with_blocks,
+            "comprehensions": comprehensions,
+        }
+
+    def _extract_function(self, node: cst.FunctionDef) -> ExtractedFunction:
+        """Extract all information from a function definition."""
+        name = node.name.value
+
+        # Build qualified name
+        qualified_parts = [self._module_name] if self._module_name else []
+        for cls in self._class_stack:
+            qualified_parts.append(cls.name)
+        qualified_parts.append(name)
+        qualified_name = ".".join(filter(None, qualified_parts))
+
+        # Extract parameters
+        params = self._extract_parameters(node.params)
+
+        # Extract return annotation
+        return_annotation = None
+        if node.returns:
+            return_annotation = self._node_to_code(node.returns.annotation)
+
+        # Extract decorators
+        decorators = self._extract_decorators(node.decorators)
+
+        # Extract docstring
+        docstring = self._extract_docstring(node.body)
+
+        # Get position
+        line, col, end_line, _ = self._get_node_position(node)
+
+        # Extract local variables
+        local_vars = self._extract_local_variables(node.body)
+
+        # Capture body source for content hashing
+        body_source = self._node_to_code(node.body)
+
+        return ExtractedFunction(
+            name=name,
+            qualified_name=qualified_name,
+            parameters=params,
+            return_annotation=return_annotation,
+            decorators=decorators,
+            is_async=isinstance(node.asynchronous, cst.Asynchronous),
+            docstring=docstring,
+            body_source=body_source or None,
+            line=line,
+            end_line=end_line,
+            column=col,
+            local_variables=local_vars,
+        )
+
+    def _extract_parameters(self, params: cst.Parameters) -> list[ExtractedParameter]:
+        """Extract function parameters."""
+        result: list[ExtractedParameter] = []
+
+        # Positional-only parameters (before /)
+        for param in params.posonly_params:
+            result.append(self._extract_single_param(param, is_positional_only=True))
+
+        # Regular parameters
+        for param in params.params:
+            result.append(self._extract_single_param(param))
+
+        # *args
+        if params.star_arg and isinstance(params.star_arg, cst.Param):
+            result.append(self._extract_single_param(params.star_arg, is_variadic=True))
+
+        # Keyword-only parameters (after *)
+        for param in params.kwonly_params:
+            result.append(self._extract_single_param(param, is_keyword_only=True))
+
+        # **kwargs
+        if params.star_kwarg:
+            result.append(self._extract_single_param(params.star_kwarg, is_keyword_variadic=True))
+
+        return result
+
+    def _extract_single_param(
+        self,
+        param: cst.Param,
+        is_variadic: bool = False,
+        is_keyword_variadic: bool = False,
+        is_positional_only: bool = False,
+        is_keyword_only: bool = False,
+    ) -> ExtractedParameter:
+        """Extract a single parameter."""
+        name = param.name.value
+
+        annotation = None
+        if param.annotation:
+            annotation = self._node_to_code(param.annotation.annotation)
+
+        default = None
+        if param.default:
+            default = self._node_to_code(param.default)
+
+        return ExtractedParameter(
+            name=name,
+            annotation=annotation,
+            default=default,
+            is_variadic=is_variadic,
+            is_keyword_variadic=is_keyword_variadic,
+            is_positional_only=is_positional_only,
+            is_keyword_only=is_keyword_only,
+        )
+
+    def _extract_local_variables(self, body: cst.BaseSuite) -> list[str]:
+        """Extract local variable names from function body."""
+        variables: list[str] = []
+
+        if isinstance(body, cst.IndentedBlock):
+            for stmt in body.body:
+                if isinstance(stmt, cst.SimpleStatementLine):
+                    for item in stmt.body:
+                        if isinstance(item, (cst.Assign, cst.AnnAssign)):
+                            vars_in_stmt = self._get_assignment_targets(item)
+                            variables.extend(vars_in_stmt)
+
+        return list(set(variables))
+
+    def _get_assignment_targets(self, node: cst.Assign | cst.AnnAssign) -> list[str]:
+        """Get variable names from an assignment."""
+        names: list[str] = []
+
+        if isinstance(node, cst.AnnAssign):
+            if isinstance(node.target, cst.Name):
+                names.append(node.target.value)
+        elif isinstance(node, cst.Assign):
+            for target in node.targets:
+                if isinstance(target.target, cst.Name):
+                    names.append(target.target.value)
+                elif isinstance(target.target, cst.Tuple):
+                    # Unpacking: a, b = ...
+                    for element in target.target.elements:
+                        if isinstance(element.value, cst.Name):
+                            names.append(element.value.value)
+
+        return names
+
+    # =========================================================================
+    # Class Extraction
+    # =========================================================================
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
+        """Extract class definitions."""
+        cls = self._extract_class(node)
+
+        if self._current_class:
+            # Nested class - not common but handle it
+            pass
+        else:
+            self.classes.append(cls)
+
+        # Push onto stack
+        self._class_stack.append(cls)
+        self._current_class = cls
+
+        return True  # Continue to extract methods
+
+    def leave_ClassDef(self, node: cst.ClassDef) -> None:
+        """Leave class definition."""
+        if self._class_stack:
+            self._class_stack.pop()
+        self._current_class = self._class_stack[-1] if self._class_stack else None
+
+    def _extract_class(self, node: cst.ClassDef) -> ExtractedClass:
+        """Extract all information from a class definition."""
+        name = node.name.value
+
+        # Build qualified name
+        qualified_parts = [self._module_name] if self._module_name else []
+        for cls in self._class_stack:
+            qualified_parts.append(cls.name)
+        qualified_parts.append(name)
+        qualified_name = ".".join(filter(None, qualified_parts))
+
+        # Extract base classes
+        bases: list[str] = []
+        if node.bases:
+            for arg in node.bases:
+                if isinstance(arg, cst.Arg):
+                    bases.append(self._node_to_code(arg.value))
+
+        # Extract decorators
+        decorators = self._extract_decorators(node.decorators)
+
+        # Check for special decorators
+        is_dataclass = any(d.name == "dataclass" for d in decorators)
+        is_pydantic = any("BaseModel" in b or "BaseSettings" in b for b in bases) or any(
+            d.name in {"validator", "field_validator"} for d in decorators
+        )
+
+        # Extract docstring
+        docstring = self._extract_docstring(node.body)
+
+        # Extract fields (for Pydantic/dataclass)
+        fields = self._extract_class_fields(node.body)
+
+        # Get position
+        line, col, end_line, _ = self._get_node_position(node)
+
+        return ExtractedClass(
+            name=name,
+            qualified_name=qualified_name,
+            bases=bases,
+            decorators=decorators,
+            docstring=docstring,
+            is_dataclass=is_dataclass,
+            is_pydantic_model=is_pydantic,
+            fields=fields,
+            line=line,
+            end_line=end_line,
+            column=col,
+        )
+
+    _FRAMEWORK_INTERNAL_NAMES = frozenset(
+        {
+            "model_config",
+            "model_fields",
+            "model_computed_fields",
+            "model_extra",
+            "model_fields_set",
+            "__private_attributes__",
+            "__class_vars__",
+            "__validators__",
+            "__pre_root_validators__",
+            "__post_root_validators__",
+        }
+    )
+
+    def _extract_class_fields(self, body: cst.BaseSuite) -> list[ExtractedField]:
+        """Extract class-level field definitions."""
+        fields: list[ExtractedField] = []
+
+        if not isinstance(body, cst.IndentedBlock):
+            return fields
+
+        for stmt in body.body:
+            if isinstance(stmt, cst.SimpleStatementLine):
+                for item in stmt.body:
+                    if isinstance(item, cst.AnnAssign):
+                        if isinstance(item.target, cst.Name):
+                            field_name = item.target.value
+
+                            if field_name in self._FRAMEWORK_INTERNAL_NAMES:
+                                continue
+
+                            annotation = self._node_to_code(item.annotation.annotation)
+
+                            if is_classvar_annotation(item.annotation.annotation):
+                                continue
+
+                            default = None
+                            field_info: dict[str, Any] = {}
+
+                            if item.value:
+                                default_code = self._node_to_code(item.value)
+                                default = default_code
+
+                                # Check for Field() calls
+                                if isinstance(item.value, cst.Call):
+                                    callee = self._node_to_code(item.value.func)
+                                    if callee in {"Field", "field"}:
+                                        field_info = self._extract_field_call_args(item.value)
+
+                            fields.append(
+                                ExtractedField(
+                                    name=field_name,
+                                    annotation=annotation,
+                                    default=default,
+                                    field_info=field_info,
+                                )
+                            )
+
+        return fields
+
+    def _extract_field_call_args(self, call: cst.Call) -> dict[str, Any]:
+        """Extract arguments from Field() call."""
+        result: dict[str, Any] = {}
+
+        for arg in call.args:
+            if arg.keyword:
+                key = arg.keyword.value
+                value = self._node_to_code(arg.value)
+                result[key] = value
+
+        return result
+
+    # =========================================================================
+    # Decorator Extraction
+    # =========================================================================
+
+    def _extract_decorators(self, decorators: Sequence[cst.Decorator]) -> list[ExtractedDecorator]:
+        """Extract decorator information."""
+        result: list[ExtractedDecorator] = []
+
+        for dec in decorators:
+            extracted = self._extract_single_decorator(dec)
+            if extracted:
+                result.append(extracted)
+
+        return result
+
+    def _extract_single_decorator(self, decorator: cst.Decorator) -> ExtractedDecorator | None:
+        """Extract a single decorator."""
+        dec_node = decorator.decorator
+
+        # Handle @name
+        if isinstance(dec_node, cst.Name):
+            return ExtractedDecorator(
+                name=dec_node.value,
+                full_name=dec_node.value,
+                arguments={},
+                positional_args=[],
+                location=self._make_location(decorator),
+                raw_source=self._node_to_code(decorator),
+            )
+
+        # Handle @name(...)
+        if isinstance(dec_node, cst.Call):
+            func = dec_node.func
+
+            # Get name
+            if isinstance(func, cst.Name):
+                name = func.value
+                full_name = func.value
+            elif isinstance(func, cst.Attribute):
+                name = func.attr.value
+                full_name = self._node_to_code(func)
+            else:
+                return None
+
+            # Extract arguments
+            kwargs: dict[str, Any] = {}
+            positional: list[Any] = []
+
+            for arg in dec_node.args:
+                if arg.keyword:
+                    # Keyword argument
+                    key = arg.keyword.value
+                    value = self._extract_literal_or_code(arg.value)
+                    kwargs[key] = value
+                else:
+                    # Positional argument
+                    value = self._extract_literal_or_code(arg.value)
+                    positional.append(value)
+
+            return ExtractedDecorator(
+                name=name,
+                full_name=full_name,
+                arguments=kwargs,
+                positional_args=positional,
+                location=self._make_location(decorator),
+                raw_source=self._node_to_code(decorator),
+            )
+
+        # Handle @obj.method or @obj.method(...)
+        if isinstance(dec_node, cst.Attribute):
+            full_name = self._node_to_code(dec_node)
+            name = dec_node.attr.value
+
+            return ExtractedDecorator(
+                name=name,
+                full_name=full_name,
+                arguments={},
+                positional_args=[],
+                location=self._make_location(decorator),
+                raw_source=self._node_to_code(decorator),
+            )
+
+        return None
+
+    # =========================================================================
+    # Call Extraction
+    # =========================================================================
+
+    def visit_Call(self, node: cst.Call) -> bool:
+        """Extract function calls."""
+        call = self._extract_call(node)
+        if call:
+            self.calls.append(call)
+        return True  # Continue visiting nested calls
+
+    def _extract_call(self, node: cst.Call) -> ExtractedCall | None:
+        """Extract a function call."""
+        # Get callee expression
+        callee = self._node_to_code(node.func)
+
+        # Determine if method call and get receiver
+        is_method_call = False
+        receiver = None
+        if isinstance(node.func, cst.Attribute):
+            is_method_call = True
+            receiver = self._node_to_code(node.func.value)
+
+        # Extract arguments
+        arguments = self._extract_call_arguments(node.args)
+
+        # Get position
+        line, col, end_line, _ = self._get_node_position(node)
+
+        # Get enclosing function
+        in_function = self._current_function.qualified_name if self._current_function else None
+
+        return ExtractedCall(
+            callee=callee,
+            arguments=arguments,
+            line=line,
+            column=col,
+            end_line=end_line,
+            in_function=in_function,
+            is_method_call=is_method_call,
+            receiver=receiver,
+            # Control flow context
+            in_loop=self._in_loop > 0,
+            in_conditional=self._in_conditional > 0,
+            in_try=self._in_try,
+            in_except=self._in_except,
+            in_finally=self._in_finally,
+            in_with=self._in_with,
+            in_comprehension=self._in_comprehension,
+            loop_depth=self._in_loop,
+            conditional_depth=self._in_conditional,
+        )
+
+    def _extract_call_arguments(self, args: Sequence[cst.Arg]) -> list[ExtractedArgument]:
+        """Extract call arguments with full structural evidence."""
+        result: list[ExtractedArgument] = []
+        pos = 0
+
+        for arg in args:
+            keyword = None
+            position = None
+            is_starred = False
+            is_double_starred = False
+
+            if arg.keyword:
+                keyword = arg.keyword.value
+            elif arg.star == "":
+                position = pos
+                pos += 1
+
+            if arg.star == "*":
+                is_starred = True
+            elif arg.star == "**":
+                is_double_starred = True
+
+            value_source = self._node_to_code(arg.value)
+            literal_info = self._extract_literal_info(arg.value)
+            name_info = self._extract_name_info(arg.value)
+
+            is_fstr = detect_fstring(arg.value)
+            is_concat = detect_concatenation(arg.value)
+            is_fmt = detect_format_call(arg.value)
+            container = detect_container_type(arg.value)
+            src_vars = collect_name_nodes(arg.value)
+
+            is_call_res = isinstance(arg.value, cst.Call)
+            called_func = None
+            if is_call_res:
+                called_func = self._node_to_code(arg.value.func)
+
+            result.append(
+                ExtractedArgument(
+                    position=position,
+                    keyword=keyword,
+                    value_source=value_source,
+                    is_literal=literal_info[0],
+                    literal_value=literal_info[1],
+                    literal_type=literal_info[2],
+                    is_name=name_info[0],
+                    name_value=name_info[1],
+                    is_starred=is_starred,
+                    is_double_starred=is_double_starred,
+                    is_string_interpolation=is_fstr,
+                    is_concatenation=is_concat,
+                    is_format_call=is_fmt,
+                    container_type=container,
+                    source_variables=src_vars,
+                    is_call_result=is_call_res,
+                    called_function=called_func,
+                )
+            )
+
+        return result
+
+    def _extract_literal_info(self, node: cst.BaseExpression) -> tuple[bool, Any, str | None]:
+        """Extract literal value if node is a literal."""
+        if isinstance(node, cst.Integer):
+            return True, int(node.value), "int"
+        if isinstance(node, cst.Float):
+            return True, float(node.value), "float"
+        if isinstance(node, (cst.SimpleString, cst.FormattedString, cst.ConcatenatedString)):
+            # Extract string value
+            try:
+                value = self._extract_string_value(node)
+                return True, value, "str"
+            except Exception:
+                return True, self._node_to_code(node), "str"
+        if isinstance(node, cst.Name):
+            if node.value == "True":
+                return True, True, "bool"
+            if node.value == "False":
+                return True, False, "bool"
+            if node.value == "None":
+                return True, None, "None"
+        if isinstance(node, cst.List):
+            return True, self._node_to_code(node), "list"
+        if isinstance(node, cst.Dict):
+            return True, self._node_to_code(node), "dict"
+        if isinstance(node, cst.Tuple):
+            return True, self._node_to_code(node), "tuple"
+
+        return False, None, None
+
+    def _extract_string_value(self, node: cst.BaseExpression) -> str:
+        """Extract the actual string value from a string node."""
+        if isinstance(node, cst.SimpleString):
+            # Remove quotes and handle escape sequences
+            raw = node.value
+            # Determine quote style
+            if raw.startswith('"""') or raw.startswith("'''"):
+                return raw[3:-3]
+            elif raw.startswith('"') or raw.startswith("'"):
+                return raw[1:-1]
+            # Handle prefixes like r"...", f"...", etc.
+            for i, c in enumerate(raw):
+                if c in "\"'":
+                    if raw[i : i + 3] in ('"""', "'''"):
+                        return raw[i + 3 : -3]
+                    return raw[i + 1 : -1]
+        return self._node_to_code(node)
+
+    def _extract_name_info(self, node: cst.BaseExpression) -> tuple[bool, str | None]:
+        """Extract name if node is a simple name reference."""
+        if isinstance(node, cst.Name) and node.value not in {"True", "False", "None"}:
+            return True, node.value
+        return False, None
+
+    # =========================================================================
+    # Assignment Extraction
+    # =========================================================================
+
+    def visit_Assign(self, node: cst.Assign) -> bool:
+        """Extract simple assignments, including tuple unpacking."""
+        for target in node.targets:
+            # Handle tuple/list unpacking: a, b = value or [a, b] = value
+            if isinstance(target.target, (cst.Tuple, cst.List)):
+                self._extract_unpacking_assignment(target.target, node.value)
+            else:
+                assignment = self._extract_assignment(target.target, node.value)
+                if assignment:
+                    self.assignments.append(assignment)
+        return True
+
+    def _extract_unpacking_assignment(
+        self,
+        target: cst.Tuple | cst.List,
+        value: cst.BaseExpression,
+    ) -> None:
+        """
+        Extract assignments from tuple/list unpacking.
+
+        Handles:
+        - a, b = 1, 2
+        - a, b = func()
+        - a, *rest, b = [1, 2, 3, 4]
+        """
+        # Get all target elements
+        elements = target.elements
+        value_source = self._node_to_code(value)
+
+        for i, element in enumerate(elements):
+            # Handle StarredElement for *rest patterns
+            if isinstance(element, cst.StarredElement):
+                inner = element.value
+            else:
+                inner = element.value if hasattr(element, "value") else element
+
+            # Extract the target name
+            target_name = self._extract_assignment_target(inner)
+            if not target_name:
+                continue
+
+            # Get position
+            line = self._get_line(inner)
+
+            # Get enclosing function
+            in_function = self._current_function.qualified_name if self._current_function else None
+
+            # Check if value is a call (for the whole assignment)
+            is_call = isinstance(value, cst.Call)
+            called_function = None
+            if is_call:
+                called_function = self._node_to_code(value.func)
+
+            source_vars, is_method, is_fstr = extract_value_metadata(value)
+            assignment = ExtractedAssignment(
+                target=target_name,
+                value_source=f"{value_source}[{i}]",
+                annotation=None,
+                line=line,
+                in_function=in_function,
+                is_literal=False,
+                is_call=is_call,
+                called_function=called_function,
+                is_name=False,
+                referenced_name=None,
+                source_variables=source_vars,
+                is_method_call=is_method,
+                is_string_interpolation=is_fstr,
+            )
+            self.assignments.append(assignment)
+
+            # Track module-level variables
+            if not self._current_function and not self._current_class:
+                self.module_variables.append(target_name)
+
+    def visit_AnnAssign(self, node: cst.AnnAssign) -> bool:
+        """Extract annotated assignments."""
+        annotation = self._node_to_code(node.annotation.annotation)
+        value = node.value
+
+        assignment = self._extract_assignment(node.target, value, annotation)
+        if assignment:
+            self.assignments.append(assignment)
+        return True
+
+    def _extract_assignment(
+        self,
+        target: cst.BaseExpression,
+        value: cst.BaseExpression | None,
+        annotation: str | None = None,
+    ) -> ExtractedAssignment | None:
+        """Extract an assignment."""
+        target_name = self._extract_assignment_target(target)
+        if not target_name:
+            return None
+
+        value_source = self._node_to_code(value) if value else ""
+
+        # Check if value is a call
+        is_call = isinstance(value, cst.Call) if value else False
+        called_function = None
+        if is_call and value:
+            called_function = self._node_to_code(value.func)
+
+        # Check if value is a name reference
+        is_name = isinstance(value, cst.Name) if value else False
+        referenced_name = value.value if is_name and isinstance(value, cst.Name) else None
+
+        # Check if literal
+        is_literal = False
+        if value:
+            is_literal, _, _ = self._extract_literal_info(value)
+
+        # Get position
+        line = self._get_line(target)
+
+        # Get enclosing function
+        in_function = self._current_function.qualified_name if self._current_function else None
+
+        # Track module-level variables
+        if not self._current_function and not self._current_class:
+            self.module_variables.append(target_name)
+
+        source_vars, is_method, is_fstr = extract_value_metadata(value)
+
+        return ExtractedAssignment(
+            target=target_name,
+            value_source=value_source,
+            annotation=annotation,
+            line=line,
+            in_function=in_function,
+            is_literal=is_literal,
+            is_call=is_call,
+            called_function=called_function,
+            is_name=is_name,
+            referenced_name=referenced_name,
+            source_variables=source_vars,
+            is_method_call=is_method,
+            is_string_interpolation=is_fstr,
+        )
+
+    def _extract_assignment_target(self, target: cst.BaseExpression) -> str | None:
+        """
+        Extract the target name from an assignment target.
+
+        Handles:
+        - Simple names: x = value
+        - Attribute access: self.x = value, obj.attr = value
+        - Subscript: d[key] = value (returns d)
+        """
+        # Simple name assignment: x = value
+        if isinstance(target, cst.Name):
+            return target.value
+
+        # Attribute assignment: self.x = value, obj.attr = value
+        if isinstance(target, cst.Attribute):
+            # Return the full dotted name
+            return self._get_dotted_name(target)
+
+        # Subscript assignment: d[key] = value
+        if isinstance(target, cst.Subscript):
+            # Return the base name
+            if isinstance(target.value, cst.Name):
+                return target.value.value
+            elif isinstance(target.value, cst.Attribute):
+                return self._get_dotted_name(target.value)
+
+        return None
+
+    # =========================================================================
+    # Control Flow Extraction
+    # =========================================================================
+
+    def visit_If(self, node: cst.If) -> bool:
+        """Track if statements."""
+        self._in_conditional += 1
+
+        if self._current_function:
+            start_line, _, end_line, _ = self._get_node_position(node)
+
+            # Detect elif and else branches
+            has_elif = False
+            has_else = False
+            elif_lines: list[int] = []
+            else_line: int | None = None
+
+            # Check for else/elif
+            if node.orelse:
+                if isinstance(node.orelse, cst.If):
+                    has_elif = True
+                    elif_line, _, _, _ = self._get_node_position(node.orelse)
+                    elif_lines.append(elif_line)
+                elif isinstance(node.orelse, cst.Else):
+                    has_else = True
+                    else_line_pos, _, _, _ = self._get_node_position(node.orelse)
+                    else_line = else_line_pos
+
+            block = ExtractedControlFlowBlock(
+                block_type="if",
+                start_line=start_line,
+                end_line=end_line,
+                has_elif=has_elif,
+                has_else=has_else,
+                elif_lines=elif_lines,
+                else_line=else_line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_If(self, node: cst.If) -> None:
+        """Leave if statement."""
+        self._in_conditional = max(0, self._in_conditional - 1)
+
+    def visit_IfExp(self, node: cst.IfExp) -> bool:
+        """Track ternary/conditional expressions (x if cond else y)."""
+        self._in_conditional += 1
+
+        if self._current_function:
+            line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="if",
+                start_line=line,
+                end_line=end_line or line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_IfExp(self, node: cst.IfExp) -> None:
+        """Leave ternary/conditional expression."""
+        self._in_conditional = max(0, self._in_conditional - 1)
+
+    def visit_For(self, node: cst.For) -> bool:
+        """Track for loops."""
+        self._in_loop += 1
+
+        if self._current_function:
+            start_line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="for",
+                start_line=start_line,
+                end_line=end_line,
+                has_break=False,  # Will be updated by visit_Break
+                has_continue=False,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_For(self, node: cst.For) -> None:
+        """Leave for loop."""
+        self._in_loop = max(0, self._in_loop - 1)
+
+    def visit_While(self, node: cst.While) -> bool:
+        """Track while loops."""
+        self._in_loop += 1
+
+        if self._current_function:
+            start_line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="while",
+                start_line=start_line,
+                end_line=end_line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_While(self, node: cst.While) -> None:
+        """Leave while loop."""
+        self._in_loop = max(0, self._in_loop - 1)
+
+    def visit_Try(self, node: cst.Try) -> bool:
+        """Track try blocks."""
+        self._in_try = True
+
+        if self._current_function:
+            start_line, _, end_line, _ = self._get_node_position(node)
+
+            # Extract except blocks
+            except_blocks: list[tuple[int, int]] = []
+            for handler in node.handlers:
+                exc_start, _, exc_end, _ = self._get_node_position(handler)
+                except_blocks.append((exc_start, exc_end))
+
+            # Extract finally block
+            finally_block: tuple[int, int] | None = None
+            if node.finalbody:
+                fin_start, _, fin_end, _ = self._get_node_position(node.finalbody)
+                finally_block = (fin_start, fin_end)
+
+            block = ExtractedControlFlowBlock(
+                block_type="try",
+                start_line=start_line,
+                end_line=end_line,
+                except_blocks=except_blocks,
+                finally_block=finally_block,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_Try(self, node: cst.Try) -> None:
+        """Leave try block."""
+        self._in_try = False
+
+    def visit_ExceptHandler(self, node: cst.ExceptHandler) -> bool:
+        """Track except handlers."""
+        self._in_except = True
+        return True
+
+    def leave_ExceptHandler(self, node: cst.ExceptHandler) -> None:
+        """Leave except handler."""
+        self._in_except = False
+
+    def visit_Finally(self, node: cst.Finally) -> bool:
+        """Track finally blocks."""
+        self._in_finally = True
+        return True
+
+    def leave_Finally(self, node: cst.Finally) -> None:
+        """Leave finally block."""
+        self._in_finally = False
+
+    def visit_With(self, node: cst.With) -> bool:
+        """Track with blocks."""
+        self._in_with = True
+
+        if self._current_function:
+            start_line, _, end_line, _ = self._get_node_position(node)
+
+            # Extract all context-manager expressions
+            context_expr = None
+            with_items: list[str] = []
+            if node.items:
+                first_item = node.items[0]
+                context_expr = self._node_to_code(first_item.item)
+                for item in node.items:
+                    expr = self._node_to_code(item.item)
+                    if expr:
+                        with_items.append(expr)
+
+            block = ExtractedControlFlowBlock(
+                block_type="with",
+                start_line=start_line,
+                end_line=end_line,
+                context_expr=context_expr,
+                with_items=with_items,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_With(self, node: cst.With) -> None:
+        """Leave with block."""
+        self._in_with = False
+
+    def visit_ListComp(self, node: cst.ListComp) -> bool:
+        """Track list comprehensions."""
+        self._in_comprehension = True
+
+        if self._current_function:
+            line, _, _, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="comprehension",
+                start_line=line,
+                end_line=line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_ListComp(self, node: cst.ListComp) -> None:
+        """Leave list comprehension."""
+        self._in_comprehension = False
+
+    def visit_SetComp(self, node: cst.SetComp) -> bool:
+        """Track set comprehensions."""
+        self._in_comprehension = True
+
+        if self._current_function:
+            line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="comprehension",
+                start_line=line,
+                end_line=end_line or line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_SetComp(self, node: cst.SetComp) -> None:
+        """Leave set comprehension."""
+        self._in_comprehension = False
+
+    def visit_DictComp(self, node: cst.DictComp) -> bool:
+        """Track dict comprehensions."""
+        self._in_comprehension = True
+
+        if self._current_function:
+            line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="comprehension",
+                start_line=line,
+                end_line=end_line or line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_DictComp(self, node: cst.DictComp) -> None:
+        """Leave dict comprehension."""
+        self._in_comprehension = False
+
+    def visit_GeneratorExp(self, node: cst.GeneratorExp) -> bool:
+        """Track generator expressions."""
+        self._in_comprehension = True
+
+        if self._current_function:
+            line, _, end_line, _ = self._get_node_position(node)
+            block = ExtractedControlFlowBlock(
+                block_type="comprehension",
+                start_line=line,
+                end_line=end_line or line,
+            )
+            self._current_control_flow.append(block)
+
+        return True
+
+    def leave_GeneratorExp(self, node: cst.GeneratorExp) -> None:
+        """Leave generator expression."""
+        self._in_comprehension = False
+
+    def visit_Yield(self, node: cst.Yield) -> bool:
+        """Track yield statements and analyze what's being yielded."""
+        if self._current_function:
+            self._current_function.has_yield = True
+
+            # For generators, track what's yielded similarly to returns
+            # We treat yields as contributing to the "return type" inference
+            if node.value:
+                ret_info = self._extract_yield_info(node)
+                if ret_info:
+                    self._current_function.return_statements.append(ret_info)
+
+        return True  # Continue to extract calls in yield expressions
+
+    def _extract_yield_info(self, node: cst.Yield) -> ExtractedReturn | None:
+        """Extract detailed information about a yield statement."""
+        if node.value is None:
+            return None
+
+        line, _, _, _ = self._get_node_position(node)
+        value = node.value
+
+        # Check for call yield: yield func()
+        if isinstance(value, cst.Call):
+            call_name = self._node_to_code(value.func)
+            return ExtractedReturn(
+                line=line,
+                returns_call=True,
+                call_name=call_name,
+                expression_text=self._node_to_code(value),
+            )
+
+        # Check for variable yield: yield x
+        if isinstance(value, cst.Name):
+            return ExtractedReturn(
+                line=line,
+                returns_variable=True,
+                variable_name=value.value,
+            )
+
+        # Check for literal yield
+        is_literal, lit_value, lit_type = self._extract_literal_info(value)
+        if is_literal:
+            return ExtractedReturn(
+                line=line,
+                returns_literal=True,
+                literal_value=lit_value,
+                literal_type=lit_type,
+            )
+
+        # Generic expression
+        return ExtractedReturn(
+            line=line,
+            returns_expression=True,
+            expression_text=self._node_to_code(value),
+        )
+
+    def visit_Return(self, node: cst.Return) -> bool:
+        """Track return statements and analyze what's being returned."""
+        if self._current_function:
+            self._current_function.has_return = True
+
+            # Extract detailed return info
+            ret_info = self._extract_return_info(node)
+            if ret_info:
+                self._current_function.return_statements.append(ret_info)
+
+        return True  # Continue to extract calls in return expressions
+
+    def _extract_return_info(self, node: cst.Return) -> ExtractedReturn:
+        """Extract detailed information about a return statement."""
+        line, _, _, _ = self._get_node_position(node)
+
+        # Check if returning None or bare return
+        if node.value is None:
+            return ExtractedReturn(
+                line=line,
+                returns_none=True,
+            )
+
+        value = node.value
+
+        # Check for return None explicitly
+        if isinstance(value, cst.Name) and value.value == "None":
+            return ExtractedReturn(
+                line=line,
+                returns_none=True,
+            )
+
+        # Check for call return: return func()
+        if isinstance(value, cst.Call):
+            call_name = self._node_to_code(value.func)
+            return ExtractedReturn(
+                line=line,
+                returns_call=True,
+                call_name=call_name,
+                expression_text=self._node_to_code(value),
+            )
+
+        # Check for variable return: return x
+        if isinstance(value, cst.Name):
+            return ExtractedReturn(
+                line=line,
+                returns_variable=True,
+                variable_name=value.value,
+            )
+
+        # Check for attribute return: return self.x or return obj.attr
+        if isinstance(value, cst.Attribute):
+            return ExtractedReturn(
+                line=line,
+                returns_variable=True,
+                variable_name=self._get_dotted_name(value),
+            )
+
+        # Check for literal return
+        is_literal, lit_value, lit_type = self._extract_literal_info(value)
+        if is_literal:
+            return ExtractedReturn(
+                line=line,
+                returns_literal=True,
+                literal_value=lit_value,
+                literal_type=lit_type,
+            )
+
+        # Check for comprehension return
+        if isinstance(value, (cst.ListComp, cst.SetComp, cst.DictComp, cst.GeneratorExp)):
+            return ExtractedReturn(
+                line=line,
+                returns_comprehension=True,
+                expression_text=self._node_to_code(value),
+            )
+
+        # Check for lambda return
+        if isinstance(value, cst.Lambda):
+            return ExtractedReturn(
+                line=line,
+                returns_lambda=True,
+                expression_text=self._node_to_code(value),
+            )
+
+        # Generic expression
+        return ExtractedReturn(
+            line=line,
+            returns_expression=True,
+            expression_text=self._node_to_code(value),
+        )
+
+    # =========================================================================
+    # Module-level Extraction
+    # =========================================================================
+
+    def visit_Module(self, node: cst.Module) -> bool:
+        """Extract module-level information."""
+        # Extract module docstring
+        if node.body:
+            first_stmt = node.body[0]
+            if isinstance(first_stmt, cst.SimpleStatementLine):
+                if first_stmt.body and isinstance(first_stmt.body[0], cst.Expr):
+                    expr = first_stmt.body[0]
+                    if isinstance(expr.value, (cst.SimpleString, cst.ConcatenatedString)):
+                        self.module_docstring = self._extract_string_value(expr.value)
+
+        return True
+
+    # =========================================================================
+    # Utility Methods
+    # =========================================================================
+
+    def _node_to_code(self, node: cst.CSTNode | None) -> str:
+        """Convert CST node back to source code."""
+        if node is None:
+            return ""
+        try:
+            return node.code if hasattr(node, "code") else cst.parse_module("").code_for_node(node)
+        except Exception:
+            # Fallback: use the module's code_for_node
+            try:
+                module = cst.parse_module("")
+                return module.code_for_node(node)
+            except Exception:
+                return str(node)
+
+    def _get_dotted_name(self, node: cst.BaseExpression) -> str:
+        """Get a dotted name from an expression (e.g., a.b.c)."""
+        if isinstance(node, cst.Name):
+            return node.value
+        if isinstance(node, cst.Attribute):
+            base = self._get_dotted_name(node.value)
+            return f"{base}.{node.attr.value}"
+        return self._node_to_code(node)
+
+    def _extract_docstring(self, body: cst.BaseSuite) -> str | None:
+        """Extract docstring from function/class body."""
+        if not isinstance(body, cst.IndentedBlock):
+            return None
+
+        if not body.body:
+            return None
+
+        first_stmt = body.body[0]
+        if isinstance(first_stmt, cst.SimpleStatementLine):
+            if first_stmt.body and isinstance(first_stmt.body[0], cst.Expr):
+                expr = first_stmt.body[0]
+                if isinstance(expr.value, (cst.SimpleString, cst.ConcatenatedString)):
+                    return self._extract_string_value(expr.value)
+
+        return None
+
+    def _extract_literal_or_code(self, node: cst.BaseExpression) -> Any:
+        """Extract literal value or return source code."""
+        is_lit, value, _ = self._extract_literal_info(node)
+        if is_lit:
+            return value
+        return self._node_to_code(node)
+
+    def _make_location(self, node: cst.CSTNode) -> CodeLocation:
+        """Make a CodeLocation from a node."""
+        line, col, end_line, end_col = self._get_node_position(node)
+        return CodeLocation(
+            file=self._file_path or Path("unknown"),
+            line=line,
+            column=col,
+            end_line=end_line,
+            end_column=end_col,
+        )
+
+    def _get_node_position(self, node: cst.CSTNode) -> tuple[int, int, int, int]:
+        """Get position of a node."""
+        if self._wrapper:
+            try:
+                pos = self._wrapper.resolve(cst.metadata.PositionProvider).get(node)
+                if pos:
+                    return (
+                        pos.start.line,
+                        pos.start.column,
+                        pos.end.line,
+                        pos.end.column,
+                    )
+            except Exception:
+                pass
+        return (0, 0, 0, 0)
+
+    def _get_line(self, node: cst.CSTNode) -> int:
+        """Get just the line number of a node."""
+        return self._get_node_position(node)[0]