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

apisec_code_bolt/parsing/python/type_resolver.py

@@ -0,0 +1,730 @@
+"""
+Type and schema resolution for Python code analysis.
+
+This module handles:
+- Resolving type annotations to their definitions
+- Resolving Pydantic model schemas including nested models
+- Building type dependency graphs
+- Resolving imports to their source definitions
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .visitors import ExtractedClass
+
+
+# =============================================================================
+# Resolved Type Information
+# =============================================================================
+
+
+@dataclass
+class ResolvedType:
+    """A resolved type with full information."""
+
+    name: str  # Simple name
+    qualified_name: str  # Full qualified name
+    module: str  # Module where defined
+
+    # Type category
+    is_builtin: bool = False
+    is_generic: bool = False
+    is_optional: bool = False
+    is_list: bool = False
+    is_dict: bool = False
+    is_union: bool = False
+    is_pydantic_model: bool = False
+    is_dataclass: bool = False
+    is_enum: bool = False
+
+    # Generic type arguments
+    type_args: list[ResolvedType] = field(default_factory=list)
+
+    # For complex types
+    origin_type: str | None = None  # e.g., "list" for list[int]
+
+
+@dataclass
+class ResolvedField:
+    """A resolved field in a model."""
+
+    name: str
+    type_annotation: str
+    resolved_type: ResolvedType | None = None
+
+    required: bool = True
+    default: Any = None
+    default_factory: str | None = None
+
+    # Field constraints
+    constraints: dict[str, Any] = field(default_factory=dict)
+
+    # Nested model reference
+    nested_model: str | None = None
+
+
+@dataclass
+class ResolvedModel:
+    """A fully resolved Pydantic model or dataclass."""
+
+    name: str
+    qualified_name: str
+    module: str
+
+    fields: list[ResolvedField] = field(default_factory=list)
+
+    # Inheritance
+    base_classes: list[str] = field(default_factory=list)
+
+    # Nested models (models used as field types)
+    nested_models: list[str] = field(default_factory=list)
+
+    # Model configuration
+    config: dict[str, Any] = field(default_factory=dict)
+
+    is_pydantic: bool = False
+    is_dataclass: bool = False
+
+
+# =============================================================================
+# Built-in Types
+# =============================================================================
+
+
+BUILTIN_TYPES = frozenset(
+    {
+        "int",
+        "str",
+        "float",
+        "bool",
+        "bytes",
+        "None",
+        "list",
+        "dict",
+        "set",
+        "frozenset",
+        "tuple",
+        "List",
+        "Dict",
+        "Set",
+        "FrozenSet",
+        "Tuple",
+        "Optional",
+        "Union",
+        "Any",
+        "Type",
+        "Callable",
+        "Sequence",
+        "Mapping",
+        "Iterable",
+        "Iterator",
+        "Literal",
+        "Final",
+        "ClassVar",
+        "TypeVar",
+        "Annotated",
+        "Generic",
+    }
+)
+
+
+TYPING_MODULE_TYPES = frozenset(
+    {
+        "List",
+        "Dict",
+        "Set",
+        "FrozenSet",
+        "Tuple",
+        "Optional",
+        "Union",
+        "Any",
+        "Type",
+        "Callable",
+        "Sequence",
+        "Mapping",
+        "Iterable",
+        "Iterator",
+        "Literal",
+        "Final",
+        "ClassVar",
+        "TypeVar",
+        "Annotated",
+        "Generic",
+        "Protocol",
+    }
+)
+
+
+# =============================================================================
+# Type Annotation Parser
+# =============================================================================
+
+
+class TypeAnnotationParser:
+    """
+    Parses type annotation strings into structured types.
+
+    Handles:
+    - Simple types: int, str, MyClass
+    - Generic types: list[int], dict[str, int]
+    - Union types: int | str, Union[int, str]
+    - Optional types: Optional[int], int | None
+    - Nested generics: list[dict[str, list[int]]]
+    """
+
+    # Pattern for generic types: TypeName[Args]
+    GENERIC_PATTERN = re.compile(r"^(\w+(?:\.\w+)*)\s*\[(.*)\]$", re.DOTALL)
+
+    # Pattern for union types using |
+    UNION_PATTERN = re.compile(r"\s*\|\s*")
+
+    def parse(self, annotation: str) -> ResolvedType:
+        """Parse a type annotation string."""
+        annotation = annotation.strip()
+
+        # Handle None
+        if annotation == "None":
+            return ResolvedType(
+                name="None",
+                qualified_name="None",
+                module="builtins",
+                is_builtin=True,
+            )
+
+        # Handle | union syntax
+        if " | " in annotation or annotation.startswith("Union["):
+            return self._parse_union(annotation)
+
+        # Handle Optional
+        if annotation.startswith("Optional["):
+            return self._parse_optional(annotation)
+
+        # Handle generic types
+        generic_match = self.GENERIC_PATTERN.match(annotation)
+        if generic_match:
+            return self._parse_generic(generic_match.group(1), generic_match.group(2))
+
+        # Simple type
+        return self._create_simple_type(annotation)
+
+    def _parse_union(self, annotation: str) -> ResolvedType:
+        """Parse a union type."""
+        # Extract types from Union[...] or A | B syntax
+        if annotation.startswith("Union["):
+            inner = annotation[6:-1]
+            type_strs = self._split_type_args(inner)
+        else:
+            type_strs = self.UNION_PATTERN.split(annotation)
+
+        type_args = [self.parse(t.strip()) for t in type_strs]
+
+        # Check if this is really Optional (Union with None)
+        has_none = any(t.name == "None" for t in type_args)
+        non_none_types = [t for t in type_args if t.name != "None"]
+
+        if has_none and len(non_none_types) == 1:
+            # This is Optional[X]
+            result = non_none_types[0]
+            result.is_optional = True
+            return result
+
+        return ResolvedType(
+            name="Union",
+            qualified_name="typing.Union",
+            module="typing",
+            is_union=True,
+            is_generic=True,
+            type_args=type_args,
+            origin_type="Union",
+        )
+
+    def _parse_optional(self, annotation: str) -> ResolvedType:
+        """Parse Optional[X] type."""
+        inner = annotation[9:-1]  # Remove "Optional[" and "]"
+        inner_type = self.parse(inner)
+        inner_type.is_optional = True
+        return inner_type
+
+    def _parse_generic(self, base_type: str, args_str: str) -> ResolvedType:
+        """Parse a generic type like list[int] or dict[str, int]."""
+        # Parse type arguments
+        arg_strs = self._split_type_args(args_str)
+        type_args = [self.parse(arg.strip()) for arg in arg_strs]
+
+        # Determine type properties
+        base_lower = base_type.lower()
+        is_list = base_lower in {"list", "sequence", "iterable"}
+        is_dict = base_lower in {"dict", "mapping"}
+        is_builtin = base_type in BUILTIN_TYPES or base_type in TYPING_MODULE_TYPES
+
+        module = "builtins" if base_lower in {"list", "dict", "set", "tuple"} else "typing"
+        if not is_builtin:
+            module = ""  # User-defined type
+
+        return ResolvedType(
+            name=base_type,
+            qualified_name=f"{module}.{base_type}" if module else base_type,
+            module=module,
+            is_builtin=is_builtin,
+            is_generic=True,
+            is_list=is_list,
+            is_dict=is_dict,
+            type_args=type_args,
+            origin_type=base_type,
+        )
+
+    def _create_simple_type(self, type_name: str) -> ResolvedType:
+        """Create a simple (non-generic) type."""
+        is_builtin = type_name in BUILTIN_TYPES
+        module = "builtins" if is_builtin else ""
+
+        return ResolvedType(
+            name=type_name,
+            qualified_name=f"{module}.{type_name}" if module else type_name,
+            module=module,
+            is_builtin=is_builtin,
+        )
+
+    def _split_type_args(self, args_str: str) -> list[str]:
+        """
+        Split type arguments respecting nested brackets.
+
+        e.g., "str, list[int], dict[str, int]" -> ["str", "list[int]", "dict[str, int]"]
+        """
+        result = []
+        current = []
+        depth = 0
+
+        for char in args_str:
+            if char == "[":
+                depth += 1
+                current.append(char)
+            elif char == "]":
+                depth -= 1
+                current.append(char)
+            elif char == "," and depth == 0:
+                result.append("".join(current).strip())
+                current = []
+            else:
+                current.append(char)
+
+        if current:
+            result.append("".join(current).strip())
+
+        return [r for r in result if r]
+
+
+# =============================================================================
+# Type Resolver
+# =============================================================================
+
+
+class TypeResolver:
+    """
+    Resolves type references to their definitions.
+
+    Maintains:
+    - Import mapping (name -> module)
+    - Class definitions (name -> ExtractedClass)
+    - Resolved type cache
+    """
+
+    def __init__(self) -> None:
+        self._parser = TypeAnnotationParser()
+
+        # Import tracking
+        self._imports: dict[str, str] = {}  # name -> module
+        self._from_imports: dict[str, tuple[str, str]] = {}  # name -> (module, original_name)
+
+        # Class definitions
+        self._classes: dict[str, ExtractedClass] = {}  # qualified_name -> class
+
+        # File to module mapping
+        self._file_modules: dict[Path, str] = {}
+
+        # Resolution cache
+        self._cache: dict[str, ResolvedType] = {}
+
+    def add_import(self, imp) -> None:
+        """
+        Register an import for resolution.
+
+        Handles both ExtractedImport (names is list of tuples) and
+        ParsedImport (names is list of strings with separate alias field).
+        """
+        # Handle different import formats
+        names_list = imp.names
+
+        if imp.is_from_import:
+            for item in names_list:
+                # Handle tuple format (name, alias) from ExtractedImport
+                if isinstance(item, tuple):
+                    name, alias = item
+                    key = alias or name
+                    self._from_imports[key] = (imp.module, name)
+                else:
+                    # Handle string format from ParsedImport
+                    name = item
+                    key = name
+                    self._from_imports[key] = (imp.module, name)
+        else:
+            for item in names_list:
+                if isinstance(item, tuple):
+                    name, alias = item
+                    key = alias or name
+                    self._imports[key] = name
+                else:
+                    name = item
+                    # For regular import, the alias is stored separately
+                    alias = getattr(imp, "alias", None)
+                    key = alias or name
+                    self._imports[key] = name
+
+    def add_class(self, cls: ExtractedClass, file_path: Path | None = None) -> None:
+        """Register a class definition."""
+        self._classes[cls.qualified_name] = cls
+        self._classes[cls.name] = cls  # Also by simple name
+
+    def set_file_module(self, file_path: Path, module_name: str) -> None:
+        """Set module name for a file."""
+        self._file_modules[file_path] = module_name
+
+    def resolve_type(self, annotation: str) -> ResolvedType:
+        """Resolve a type annotation to full type information."""
+        if annotation in self._cache:
+            return self._cache[annotation]
+
+        resolved = self._parser.parse(annotation)
+
+        # Resolve user-defined types
+        if not resolved.is_builtin and not resolved.is_generic:
+            self._resolve_user_type(resolved)
+
+        # Recursively resolve type arguments
+        for i, arg in enumerate(resolved.type_args):
+            if not arg.is_builtin:
+                resolved.type_args[i] = self.resolve_type(arg.name)
+
+        self._cache[annotation] = resolved
+        return resolved
+
+    def _resolve_user_type(self, resolved: ResolvedType) -> None:
+        """Resolve a user-defined type to its definition."""
+        name = resolved.name
+
+        # Check from imports
+        if name in self._from_imports:
+            module, original_name = self._from_imports[name]
+            resolved.module = module
+            resolved.qualified_name = f"{module}.{original_name}"
+
+        # Check regular imports
+        elif name in self._imports:
+            resolved.module = self._imports[name]
+            resolved.qualified_name = self._imports[name]
+
+        # Check registered classes
+        if name in self._classes:
+            cls = self._classes[name]
+            resolved.is_pydantic_model = cls.is_pydantic_model
+            resolved.is_dataclass = cls.is_dataclass
+            resolved.qualified_name = cls.qualified_name
+
+    def resolve_model(self, cls: ExtractedClass) -> ResolvedModel:
+        """Fully resolve a Pydantic model or dataclass."""
+        resolved_fields: list[ResolvedField] = []
+        nested_models: list[str] = []
+
+        for f in cls.fields:
+            resolved_field = ResolvedField(
+                name=f.name,
+                type_annotation=f.annotation or "Any",
+            )
+
+            # Resolve the field type
+            if f.annotation:
+                resolved_type = self.resolve_type(f.annotation)
+                resolved_field.resolved_type = resolved_type
+
+                # Check for nested models
+                nested = self._find_nested_models(resolved_type)
+                nested_models.extend(nested)
+                if nested:
+                    resolved_field.nested_model = nested[0]
+
+            # Parse field constraints from default
+            if f.default:
+                self._parse_field_constraints(resolved_field, f.default, f.field_info)
+            else:
+                resolved_field.required = True
+
+            resolved_fields.append(resolved_field)
+
+        return ResolvedModel(
+            name=cls.name,
+            qualified_name=cls.qualified_name,
+            module=cls.qualified_name.rsplit(".", 1)[0] if "." in cls.qualified_name else "",
+            fields=resolved_fields,
+            base_classes=cls.bases,
+            nested_models=list(set(nested_models)),
+            is_pydantic=cls.is_pydantic_model,
+            is_dataclass=cls.is_dataclass,
+        )
+
+    def _find_nested_models(self, resolved_type: ResolvedType) -> list[str]:
+        """Find nested model references in a type."""
+        nested: list[str] = []
+
+        # Check if this type itself is a model
+        if resolved_type.is_pydantic_model or resolved_type.is_dataclass:
+            nested.append(resolved_type.qualified_name or resolved_type.name)
+
+        # Check type arguments recursively
+        for arg in resolved_type.type_args:
+            nested.extend(self._find_nested_models(arg))
+
+        return nested
+
+    def _parse_field_constraints(
+        self,
+        field: ResolvedField,
+        default_value: str,
+        field_info: dict[str, Any],
+    ) -> None:
+        """Parse field constraints from default value and Field() info."""
+        # Check if it's a Field() call
+        if default_value.startswith("Field("):
+            field.required = "..." in default_value or "default" not in field_info
+
+            # Copy constraints from field_info
+            for key, value in field_info.items():
+                if key in {"default", "default_factory"}:
+                    if key == "default":
+                        field.default = value
+                        field.required = False
+                    else:
+                        field.default_factory = value
+                        field.required = False
+                else:
+                    field.constraints[key] = value
+        elif default_value == "...":
+            field.required = True
+        elif default_value != "":
+            field.required = False
+            field.default = default_value
+
+    def get_all_nested_models(self, model_name: str) -> list[ResolvedModel]:
+        """
+        Get all nested models recursively for a given model.
+
+        Returns list of all models that are referenced by the given model,
+        including transitively nested models.
+        """
+        visited: set[str] = set()
+        result: list[ResolvedModel] = []
+
+        self._collect_nested_models(model_name, visited, result)
+
+        return result
+
+    def _collect_nested_models(
+        self,
+        model_name: str,
+        visited: set[str],
+        result: list[ResolvedModel],
+    ) -> None:
+        """Recursively collect nested models."""
+        if model_name in visited:
+            return
+        visited.add(model_name)
+
+        # Find the class
+        cls = self._classes.get(model_name)
+        if not cls:
+            return
+
+        # Resolve the model
+        resolved = self.resolve_model(cls)
+        result.append(resolved)
+
+        # Recursively collect nested models
+        for nested_name in resolved.nested_models:
+            self._collect_nested_models(nested_name, visited, result)
+
+
+# =============================================================================
+# Schema Builder
+# =============================================================================
+
+
+class SchemaBuilder:
+    """
+    Builds complete schema representations for API request/response bodies.
+
+    Handles:
+    - Pydantic model schemas
+    - Nested model resolution
+    - Field type resolution
+    - Constraint extraction
+    """
+
+    def __init__(self, resolver: TypeResolver) -> None:
+        self._resolver = resolver
+
+    def build_body_schema(
+        self,
+        model_name: str,
+        classes: dict[str, ExtractedClass],
+    ) -> dict[str, Any]:
+        """
+        Build a complete schema for a request/response body model.
+
+        Returns a JSON-schema-like dict representation.
+        """
+        # Register all classes
+        for cls in classes.values():
+            self._resolver.add_class(cls)
+
+        # Check if model exists
+        if model_name not in classes:
+            return {"type": "unknown", "name": model_name}
+
+        cls = classes[model_name]
+        return self._build_model_schema(cls, set())
+
+    def _build_model_schema(
+        self,
+        cls: ExtractedClass,
+        visited: set[str],
+    ) -> dict[str, Any]:
+        """Build schema for a single model."""
+        if cls.qualified_name in visited:
+            return {"$ref": cls.qualified_name}
+        visited.add(cls.qualified_name)
+
+        schema: dict[str, Any] = {
+            "type": "object",
+            "name": cls.name,
+            "qualified_name": cls.qualified_name,
+            "properties": {},
+            "required": [],
+        }
+
+        for f in cls.fields:
+            field_schema = self._build_field_schema(f, visited)
+            schema["properties"][f.name] = field_schema
+
+            # Check if required
+            if f.default is None and not f.field_info.get("default"):
+                schema["required"].append(f.name)
+
+        return schema
+
+    def _build_field_schema(
+        self,
+        f: ExtractedField,  # noqa: F821
+        visited: set[str],
+    ) -> dict[str, Any]:
+        """Build schema for a single field."""
+        field_schema: dict[str, Any] = {}
+
+        if f.annotation:
+            resolved = self._resolver.resolve_type(f.annotation)
+            field_schema = self._type_to_schema(resolved, visited)
+
+        # Add constraints
+        if f.field_info:
+            for key, value in f.field_info.items():
+                if key not in {"default", "default_factory"}:
+                    field_schema[key] = value
+
+        # Add default
+        if f.default:
+            field_schema["default"] = f.default
+
+        return field_schema
+
+    def _type_to_schema(
+        self,
+        resolved: ResolvedType,
+        visited: set[str],
+    ) -> dict[str, Any]:
+        """Convert a resolved type to JSON-schema representation."""
+        # Handle builtins
+        if resolved.is_builtin:
+            return self._builtin_to_schema(resolved)
+
+        # Handle Optional
+        if resolved.is_optional:
+            inner_schema = self._type_to_schema(
+                resolved.type_args[0] if resolved.type_args else resolved,
+                visited,
+            )
+            return {"anyOf": [inner_schema, {"type": "null"}]}
+
+        # Handle Union
+        if resolved.is_union:
+            return {"anyOf": [self._type_to_schema(arg, visited) for arg in resolved.type_args]}
+
+        # Handle List
+        if resolved.is_list and resolved.type_args:
+            return {
+                "type": "array",
+                "items": self._type_to_schema(resolved.type_args[0], visited),
+            }
+
+        # Handle Dict
+        if resolved.is_dict and len(resolved.type_args) >= 2:
+            return {
+                "type": "object",
+                "additionalProperties": self._type_to_schema(resolved.type_args[1], visited),
+            }
+
+        # Handle user-defined models
+        if resolved.is_pydantic_model or resolved.is_dataclass:
+            # Check if we have the class definition
+            cls_name = resolved.qualified_name or resolved.name
+            if cls_name in self._resolver._classes:
+                cls = self._resolver._classes[cls_name]
+                return self._build_model_schema(cls, visited)
+            return {"$ref": cls_name}
+
+        # Unknown type
+        return {"type": "unknown", "name": resolved.name}
+
+    def _builtin_to_schema(self, resolved: ResolvedType) -> dict[str, Any]:
+        """Convert builtin type to JSON schema."""
+        type_mapping = {
+            "str": {"type": "string"},
+            "int": {"type": "integer"},
+            "float": {"type": "number"},
+            "bool": {"type": "boolean"},
+            "None": {"type": "null"},
+            "bytes": {"type": "string", "format": "binary"},
+            "Any": {},
+        }
+
+        name = resolved.name.lower()
+        if name in type_mapping:
+            return type_mapping[name]
+
+        # Handle generic builtins
+        if resolved.is_list:
+            items = self._type_to_schema(resolved.type_args[0], set()) if resolved.type_args else {}
+            return {"type": "array", "items": items}
+
+        if resolved.is_dict:
+            additional = (
+                self._type_to_schema(resolved.type_args[1], set())
+                if len(resolved.type_args) > 1
+                else {}
+            )
+            return {"type": "object", "additionalProperties": additional}
+
+        return {"type": resolved.name}