niishantth/codegraph

1"""Pydantic models for the knowledge graph entities and relationships."""2 3from __future__ import annotations4 5from enum import Enum6from typing import Any, Optional7 8from pydantic import BaseModel, Field9 10 11class EntityType(str, Enum):12    FILE = "file"13    MODULE = "module"14    CLASS = "class"15    FUNCTION = "function"16    METHOD = "method"17    VARIABLE = "variable"18    IMPORT = "import"19    CALL = "call"20    ERROR = "error"21 22 23class RelationType(str, Enum):24    CONTAINS = "contains"25    IMPORTS = "imports"26    CALLS = "calls"27    INHERITS = "inherits"28    DEFINES = "defines"29    USES = "uses"30    RAISES = "raises"31    HANDLES = "handles"32    RETURN_TYPE = "return_type"33    ARG_TYPE = "arg_type"34 35 36class Position(BaseModel):37    line: int38    column: int39 40 41class Span(BaseModel):42    start: Position43    end: Position44 45 46class Entity(BaseModel):47    id: str48    type: EntityType49    name: str50    file_path: str51    span: Optional[Span] = None52    meta: dict[str, Any] = Field(default_factory=dict)53 54    def __hash__(self) -> int:55        return hash(self.id)56 57    def __eq__(self, other: object) -> bool:58        if isinstance(other, Entity):59            return self.id == other.id60        return False61 62 63class Relationship(BaseModel):64    source: str65    target: str66    type: RelationType67    meta: dict[str, Any] = Field(default_factory=dict)68 69 70class ErrorMapping(BaseModel):71    error_message: str72    error_type: str73    file_path: Optional[str] = None74    line: Optional[int] = None75    column: Optional[int] = None76    matched_entity: Optional[Entity] = None77    related_entities: list[Entity] = Field(default_factory=list)78    suggested_fix: Optional[str] = None79