diff --git a/src/termdiagram/models/diagrams.py b/src/termdiagram/models/diagrams.py new file mode 100644 index 0000000..83dfcf6 --- /dev/null +++ b/src/termdiagram/models/diagrams.py @@ -0,0 +1,81 @@ +from typing import List, Optional, Dict, Any +from dataclasses import dataclass, field +from enum import Enum + + +class SymbolType(Enum): + MODULE = "module" + CLASS = "class" + FUNCTION = "function" + METHOD = "method" + + +@dataclass +class MethodSymbol: + name: str + params: List[str] + return_type: Optional[str] = None + line_number: int = 0 + docstring: Optional[str] = None + + +@dataclass +class FunctionSymbol: + name: str + params: List[str] + return_type: Optional[str] = None + line_number: int = 0 + docstring: Optional[str] = None + + +@dataclass +class ClassSymbol: + name: str + bases: List[str] + methods: List[MethodSymbol] = field(default_factory=list) + line_number: int = 0 + docstring: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + return { + "name": self.name, + "bases": self.bases, + "methods": [ + { + "name": m.name, + "params": m.params, + "return_type": m.return_type, + "line_number": m.line_number, + } + for m in self.methods + ], + "line_number": self.line_number, + } + + +@dataclass +class Module: + name: str + path: str + classes: List[ClassSymbol] = field(default_factory=list) + functions: List[FunctionSymbol] = field(default_factory=list) + imports: List[str] = field(default_factory=list) + line_count: int = 0 + + def to_dict(self) -> Dict[str, Any]: + return { + "name": self.name, + "path": self.path, + "classes": [c.to_dict() for c in self.classes], + "functions": [ + { + "name": f.name, + "params": f.params, + "return_type": f.return_type, + "line_number": f.line_number, + } + for f in self.functions + ], + "imports": self.imports, + "line_count": self.line_count, + }