Initial upload: TermDiagram v0.1.0

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,
+        }