Switch to ruff

.github/workflows/main.yml

@@ -39,3 +39,21 @@ jobs:
         run: poetry install --with dev
       - name: Run tests
         run: poetry run pytest
+
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v6
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.13"
+      - name: Install Poetry
+        run: pip install "poetry>=2.1"
+      - name: Install dependencies
+        run: poetry install --with dev
+      - name: Check formatting
+        run: poetry run ruff format --check
+      - name: Lint
+        run: poetry run ruff check --output-format=github

CONTRIBUTING.md

@@ -49,6 +49,16 @@ can install pytm together with its development dependencies:
 Note that the `Makefile` targets drive the tools through Poetry, so reach for `pytest` and the
 `pytm` modules directly in an environment installed this way.
 
+### Linting and formatting
+
+[Ruff](https://docs.astral.sh/ruff/) is both the linter and the formatter (the style is
+black-compatible); the configuration lives in `pyproject.toml`. Before submitting a PR run
+
+    make fmt
+
+to format the code and apply auto-fixable lint findings. CI enforces the read-only
+equivalent, which you can reproduce locally with `make lint`.
+
 ### Dependencies
 
 Dependency changes go into `pyproject.toml` - runtime ones under `[project.dependencies]`,

Makefile

@@ -63,4 +63,10 @@ docs: docs/pytm/index.html docs/threats.md
 
 .PHONY: fmt
 fmt:
-	poetry run black  $(wildcard pytm/*.py) $(wildcard tests/*.py) $(wildcard *.py)
+	poetry run ruff check --fix --exit-zero
+	poetry run ruff format
+
+.PHONY: lint
+lint:
+	poetry run ruff format --check
+	poetry run ruff check

docs/sample_llm.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 """Sample threat model demonstrating LLM element usage."""
 
-from pytm import TM, LLM, Server, Datastore, Boundary, Dataflow, Actor
+from pytm import LLM, TM, Actor, Boundary, Dataflow, Datastore, Server
 
 tm = TM("Sample LLM Threat Model")
-tm.description = "A web app using an LLM API for chat and a self-hosted model for classification"
+tm.description = (
+    "A web app using an LLM API for chat and a self-hosted model for classification"
+)
 
 # Boundaries
 internet = Boundary("Internet")

pyproject.toml

@@ -25,7 +25,6 @@ Homepage = "https://github.com/OWASP/pytm"
 [dependency-groups]
 dev = [
     "pytest>=8.3.5,<10.0.0",
-    "black>=25.9,<27.0",
     "pdoc3>=0.11.6,<0.12.0",
     "ruff>=0.15.11,<0.16.0",
 ]
@@ -42,8 +41,7 @@ select = [
     "UP",  # pyupgrade
 ]
 ignore = [
-    "E501",  # line too long — black handles line length for reformattable code
-    "UP007", # Union[X, Y] -> X | Y syntax requires Python 3.10+ at runtime
+    "E501",  # line too long — the formatter handles line length for reformattable code
 ]
 
 [build-system]

pytm/__init__.py

@@ -29,22 +29,22 @@
 
 import sys
 
-from .json import load, loads
-from .pytm import var
+from .actor import Actor
+from .asset import LLM, Agent, Asset, ExternalEntity, Lambda, Server
+from .base import Assumption, Controls
+from .boundary import Boundary
+from .data import Data
+from .dataflow import Dataflow
+from .datastore import Datastore
+from .element import Element
 
 # 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 Agent, Asset, Lambda, LLM, Server, ExternalEntity
-from .datastore import Datastore
-from .actor import Actor
+from .json import load, loads
 from .process import Process, SetOfProcesses
-from .dataflow import Dataflow
-from .boundary import Boundary
+from .pytm import var
+from .threat import Threat
 from .tm import TM
 
 # Rebuild models to resolve forward references

pytm/actor.py

@@ -1,10 +1,11 @@
 """Actor model - represents entities that initiate actions."""
 
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
+
 from pydantic import Field, field_validator
 
-from .element import Element
 from .base import DataSet
+from .element import Element
 
 if TYPE_CHECKING:
     from .dataflow import Dataflow
@@ -35,10 +36,10 @@ class Actor(Element):
         default_factory=DataSet,
         description="pytm.Data object(s) in outgoing data flows",
     )
-    inputs: List["Dataflow"] = Field(
+    inputs: list["Dataflow"] = Field(
         default_factory=list, description="Incoming Dataflows"
     )
-    outputs: List["Dataflow"] = Field(
+    outputs: list["Dataflow"] = Field(
         default_factory=list, description="Outgoing Dataflows"
     )
     isAdmin: bool = Field(

pytm/asset.py

@@ -1,11 +1,11 @@
 """Asset models - base Asset class and specific asset implementations."""
 
-from typing import List, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from pydantic import Field, field_validator
 
-from .element import Element, sev_to_color
 from .base import DataSet
+from .element import Element, sev_to_color
 
 if TYPE_CHECKING:
     from .dataflow import Dataflow
@@ -36,10 +36,10 @@ class Asset(Element):
         default_factory=DataSet,
         description="pytm.Data object(s) in incoming data flows",
     )
-    inputs: List["Dataflow"] = Field(
+    inputs: list["Dataflow"] = Field(
         default_factory=list, description="incoming Dataflows"
     )
-    outputs: List["Dataflow"] = Field(
+    outputs: list["Dataflow"] = Field(
         default_factory=list, description="outgoing Dataflows"
     )
     onAWS: bool = Field(default=False, description="Is this asset on AWS?")

pytm/base.py

@@ -2,25 +2,26 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable, List, Set, Union, TYPE_CHECKING
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, ConfigDict, Field
 
 if TYPE_CHECKING:
-    from .element import Element
     from .data import Data
-    from .threat import Threat
+    from .element import Element
     from .finding import Finding
+    from .threat import Threat
 
 
 class DataSet(set):
     """Custom set for Data objects with string lookup capability."""
 
     __slots__ = ("_names",)
 
-    def __init__(self, values: Iterable["Data"] | None = None):
+    def __init__(self, values: Iterable[Data] | None = None):
         super().__init__()
-        self._names: Set[str] = set()
+        self._names: set[str] = set()
         if values is not None:
             self.update(values)
 
@@ -74,12 +75,12 @@ def clear(self) -> None:  # type: ignore[override]
         super().clear()
         self._names.clear()
 
-    def _register(self, element: "Data") -> None:
+    def _register(self, element: Data) -> None:
         name = getattr(element, "name", None)
         if isinstance(name, str):
             self._names.add(name)
 
-    def _unregister(self, element: "Data") -> None:
+    def _unregister(self, element: Data) -> None:
         name = getattr(element, "name", None)
         if isinstance(name, str):
             self._names.discard(name)
@@ -177,7 +178,7 @@ class Assumption(BaseModel):
     model_config = ConfigDict(extra="allow")
 
     name: str = Field(description="Name of the assumption")
-    exclude: Set[str] = Field(
+    exclude: set[str] = Field(
         default_factory=set,
         description="A set of threat SIDs to exclude for this assumption. For example: INP01",
     )
@@ -186,7 +187,7 @@ class Assumption(BaseModel):
     )
 
     def __init__(
-        self, name: str = None, exclude: Union[List[str], Set[str]] = None, **kwargs
+        self, name: str = None, exclude: list[str] | set[str] = None, **kwargs
     ):
         """Initialize an Assumption.
 
@@ -208,9 +209,9 @@ def __str__(self):
 
 
 # Type aliases for complex field types that reference forward declarations
-ElementList = List["Element"]
-DataList = List["Data"]
-ThreatList = List["Threat"]
-FindingList = List["Finding"]
+ElementList = list["Element"]
+DataList = list["Data"]
+ThreatList = list["Threat"]
+FindingList = list["Finding"]
 ControlsType = Controls
-AssumptionList = List[Assumption]
+AssumptionList = list[Assumption]

pytm/boundary.py

@@ -1,7 +1,7 @@
 """Boundary model - represents trust boundaries in the threat model."""
 
-from typing import List, TYPE_CHECKING
 from textwrap import indent
+from typing import TYPE_CHECKING
 
 from .element import Element
 
@@ -75,7 +75,7 @@ def _color(self, **kwargs) -> str:
         else:
             return "firebrick2"
 
-    def parents(self) -> List["Boundary"]:
+    def parents(self) -> list["Boundary"]:
         """Get parent boundaries."""
         result = []
         parent = self.inBoundary

pytm/data.py

@@ -1,13 +1,14 @@
 """Data model - represents data that traverses the threat model."""
 
-from typing import List, TYPE_CHECKING
-from pydantic import BaseModel, Field, ConfigDict
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field
 
 from .enums import Classification, Lifetime
 
 if TYPE_CHECKING:
-    from .element import Element
     from .dataflow import Dataflow
+    from .element import Element
 
 
 class Data(BaseModel):
@@ -59,10 +60,10 @@ class Data(BaseModel):
     isSourceEncryptedAtRest: bool = Field(
         default=False, description="Is data encrypted at rest at source?"
     )
-    carriedBy: List["Dataflow"] = Field(
+    carriedBy: list["Dataflow"] = Field(
         default_factory=list, description="Dataflows that carries this piece of data"
     )
-    processedBy: List["Element"] = Field(
+    processedBy: list["Element"] = Field(
         default_factory=list,
         description="Elements that store/process this piece of data",
     )

pytm/dataflow.py

@@ -1,11 +1,12 @@
 """Dataflow model - represents data flows between elements."""
 
 from typing import Optional
+
 from pydantic import Field, field_validator, model_validator
 
+from .base import DataSet
 from .element import Element, sev_to_color
 from .enums import Classification, TLSVersion
-from .base import DataSet
 
 
 class Dataflow(Element):

pytm/datastore.py

@@ -2,6 +2,7 @@
 
 import os
 from typing import TYPE_CHECKING
+
 from pydantic import Field
 
 from .asset import Asset

pytm/element.py

@@ -5,7 +5,7 @@
 import uuid as uuid_module
 from hashlib import sha224
 from textwrap import wrap
-from typing import Any, List, Optional, Set, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
@@ -72,23 +72,23 @@ class Element(BaseModel):
         default=TLSVersion.NONE,
         description="Minimum TLS version required",
     )
-    findings: List["Finding"] = Field(
+    findings: list["Finding"] = Field(
         default_factory=list,
         description="Threats that apply to this element",
     )
-    overrides: List["Finding"] = Field(
+    overrides: list["Finding"] = Field(
         default_factory=list,
         description="Overrides to findings, allowing to set a custom response, CVSS score or override other attributes",
     )
-    assumptions: List[Assumption] = Field(
+    assumptions: list[Assumption] = Field(
         default_factory=list,
         description="Assumptions about the element. These optionally allow to exclude threats with the given SIDs",
     )
-    levels: Set[int] = Field(
+    levels: set[int] = Field(
         default_factory=lambda: {0},
         description="List of levels (0, 1, 2, ...) to be drawn in the model",
     )
-    sourceFiles: List[str] = Field(
+    sourceFiles: list[str] = Field(
         default_factory=list,
         description="Location of the source code that describes this element relative to the directory of the model script",
     )
@@ -119,9 +119,7 @@ def _coerce_levels(cls, value):
             return set(value)
         return {value}
 
-    def __setattr__(
-        self, key: str, value: Any
-    ) -> None:  # noqa: D401 - keep same behaviour
+    def __setattr__(self, key: str, value: Any) -> None:  # noqa: D401 - keep same behaviour
         if (
             key in self._WRITE_ONCE_FIELDS
             and key in self.__dict__
@@ -130,7 +128,7 @@ def __setattr__(
             raise ValueError(f"cannot overwrite {type(self).__name__}.{key} value")
         super().__setattr__(key, value)
 
-    def __init__(self, name: Optional[str] = None, **data: Any):
+    def __init__(self, name: str | None = None, **data: Any):
         """Initialize an Element.
 
         Args:
@@ -300,7 +298,7 @@ def _attr_values(self) -> dict:
         """Return a dictionary of all attribute values."""
         return self.model_dump()
 
-    def checkTLSVersion(self, flows: List["Dataflow"]) -> bool:
+    def checkTLSVersion(self, flows: list["Dataflow"]) -> bool:
         """Check if any flows have insufficient TLS version."""
         return any(f.tlsVersion < self.minTLSVersion for f in flows)