app/models/code.py

from pydantic import BaseModel, Field


class ImportInfo(BaseModel):
    module: str | None
    name: str
    alias: str | None = None


class CallInfo(BaseModel):
    name: str
    line: int
    qualified_name: str | None = None


class FunctionInfo(BaseModel):
    name: str

    qualified_name: str

    line: int

    end_line: int

    docstring: str | None

    source_code: str

    calls: list[CallInfo] = Field(default_factory=list)


class ClassInfo(BaseModel):

    name: str

    qualified_name: str

    line: int
    
    end_line: int

    source_code: str

    methods: list[FunctionInfo] = Field(default_factory=list)


class PythonFileInfo(BaseModel):
    path: str

    imports: list[ImportInfo] = Field(default_factory=list)

    classes: list[ClassInfo] = Field(default_factory=list)

    functions: list[FunctionInfo] = Field(default_factory=list)