Refactor models to Pydantic

pyproject.toml

@@ -6,9 +6,8 @@ authors = ["pytm Team"]
 license = "MIT License"
 
 [tool.poetry.dependencies]
-python = "^3.9 || ^3.10 || ^3.11"
-pydal = "~20200714.1"
-legacy-cgi = { version = "^2.0", markers = "python_version >= '3.13'" }
+python = "^3.9 || ^3.10 || ^3.11 || ^3.12 || ^3.13 || ^3.14"
+pydantic = "^2.10.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.5"

pytm/__init__.py

@@ -22,35 +22,45 @@
     "SetOfProcesses",
     "Threat",
     "TM",
+    "Controls",
+    "var",
 ]
 
 import sys
 
 from .json import load, loads
-from .pytm import (
-    TM,
-    Action,
-    Actor,
-    Assumption,
-    Boundary,
-    Classification,
-    Data,
-    Dataflow,
-    Datastore,
-    DatastoreType,
-    Element,
-    ExternalEntity,
-    Finding,
-    Lambda,
-    LLM,
-    Lifetime,
-    Process,
-    Server,
-    SetOfProcesses,
-    Threat,
-    TLSVersion,
-    var,
-)
+from .pytm import var
+# Import from new Pydantic models
+from .enums import Action, Classification, DatastoreType, Lifetime, TLSVersion
+from .base import Assumption, Controls
+from .element import Element
+from .data import Data
+from .threat import Threat
+from .finding import Finding
+from .asset import Asset, Lambda, LLM, Server, ExternalEntity
+from .datastore import Datastore
+from .actor import Actor
+from .process import Process, SetOfProcesses
+from .dataflow import Dataflow
+from .boundary import Boundary
+from .tm import TM
+
+# Rebuild models to resolve forward references
+Element.model_rebuild()
+Data.model_rebuild()
+Finding.model_rebuild()
+Asset.model_rebuild()
+Lambda.model_rebuild()
+LLM.model_rebuild()
+Server.model_rebuild()
+ExternalEntity.model_rebuild()
+Datastore.model_rebuild()
+Actor.model_rebuild()
+Process.model_rebuild()
+SetOfProcesses.model_rebuild()
+Dataflow.model_rebuild()
+Boundary.model_rebuild()
+TM.model_rebuild()
 
 
 def pdoc_overrides():
@@ -62,9 +72,11 @@ def pdoc_overrides():
         for i in dir(klass):
             if i in ("check", "dfd", "seq"):
                 result[f"{name}.{i}"] = False
-            attr = getattr(klass, i, {})
-            if isinstance(attr, var) and attr.doc != "":
-                result[f"{name}.{i}"] = attr.doc
+            model_fields = getattr(klass, "model_fields", {})
+            if i in model_fields:
+                description = model_fields[i].description
+                if description:
+                    result[f"{name}.{i}"] = description
     return result
 
 

pytm/actor.py

@@ -0,0 +1,103 @@
+"""Actor model - represents entities that initiate actions."""
+
+from typing import TYPE_CHECKING, List
+from pydantic import Field, field_validator
+
+from .element import Element
+from .base import DataSet
+
+if TYPE_CHECKING:
+    from .data import Data
+    from .dataflow import Dataflow
+
+
+class Actor(Element):
+    """An entity usually initiating actions.
+
+    Actors represent users or external systems that initiate
+    interactions with the system being modeled.
+
+    Attributes:
+        port (int): Default TCP port for outgoing data flows
+        protocol (str): Default network protocol for outgoing data flows
+        data (DataSet): pytm.Data object(s) in outgoing data flows
+        inputs (List[Dataflow]): Incoming Dataflows
+        outputs (List[Dataflow]): Outgoing Dataflows
+        isAdmin (bool): Indicates whether the actor has administrative privileges
+    """
+
+    port: int = Field(
+        default=-1, description="Default TCP port for outgoing data flows"
+    )
+    protocol: str = Field(
+        default="", description="Default network protocol for outgoing data flows"
+    )
+    data: DataSet = Field(
+        default_factory=DataSet,
+        description="pytm.Data object(s) in outgoing data flows",
+    )
+    inputs: List["Dataflow"] = Field(
+        default_factory=list, description="Incoming Dataflows"
+    )
+    outputs: List["Dataflow"] = Field(
+        default_factory=list, description="Outgoing Dataflows"
+    )
+    isAdmin: bool = Field(
+        default=False,
+        description="Indicates whether the actor has administrative privileges",
+    )
+
+    @field_validator("data", mode="before")
+    @classmethod
+    def _coerce_dataset(cls, v):
+        """Ensure actor data is stored as a DataSet."""
+        from .data import Data  # Local import to avoid circular dependency
+
+        if isinstance(v, DataSet):
+            return v
+
+        dataset = DataSet()
+
+        if v is None:
+            return dataset
+
+        if isinstance(v, Data):
+            dataset.add(v)
+            return dataset
+
+        if hasattr(v, "__iter__") and not isinstance(v, (str, bytes)):
+            for item in v:
+                if item is None:
+                    continue
+                dataset.add(item)
+            return dataset
+
+        dataset.add(v)
+        return dataset
+
+    def __init__(self, name: str = None, **data):
+        """
+        Initialize an Actor.
+
+        Args:
+            name (str): Name of the actor.
+            **data: Optional actor properties:
+                - port (int): Default TCP port for outgoing data flows
+                - protocol (str): Default network protocol for outgoing data flows
+                - data (DataSet): pytm.Data object(s) in outgoing data flows
+                - inputs (List[Dataflow]): Incoming Dataflows
+                - outputs (List[Dataflow]): Outgoing Dataflows
+                - isAdmin (bool): Indicates whether the actor has administrative privileges
+        """
+        super().__init__(name, **data)
+        # Register with TM actors
+        self._register_with_tm_actors()
+
+    def _register_with_tm_actors(self):
+        """Register this actor with the TM class."""
+        try:
+            from .tm import TM
+
+            TM._actors.append(self)
+        except ImportError:
+            pass