Add mypy type checking

mkdocs.yml

@@ -13,3 +13,8 @@ extra:
   version:
     provider: mike
     default: latest
+
+plugins:
+- search
+- mkdocs-autoapi
+- mkdocstrings

pyproject.toml

@@ -35,8 +35,8 @@ dependencies =[
 ]
 [project.optional-dependencies]
 testing =["pytest", "pytest-cov"]
-docs = ["mkdocs-material", "mike"]
-dev = ["ruff", "pre-commit"]
+docs = ["mkdocs-material", "mike", "mkdocstrings[python]", "mkdocs-autoapi"]
+dev = ["ruff", "pre-commit", "mypy"]
 all = ["funtracks[testing,docs,dev]"]
 
 [project.urls]
@@ -102,3 +102,7 @@ testing = { features = ["testing"], solve-group = "default" }
 
 [tool.pixi.feature.testing.tasks]
 test = "pytest --cov=funtracks tests/"
+
+
+[tool.pixi.feature.docs.tasks]
+localdocs = "mike build"

src/funtracks/data_model/action_history.py

@@ -14,7 +14,7 @@ class ActionHistory:
     the undone actions)
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.undo_stack: list[TracksAction] = []  # list of actions that can be undone
         self.redo_stack: list[TracksAction] = []  # list of actions that can be redone
 

src/funtracks/data_model/actions.py

@@ -16,10 +16,15 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from .graph_attributes import NodeAttr
 from .solution_tracks import SolutionTracks
 from .tracks import Attrs, Edge, Node, SegMask, Tracks
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 class TracksAction:
     def __init__(self, tracks: Tracks):
@@ -81,9 +86,9 @@ class AddNodes(TracksAction):
     def __init__(
         self,
         tracks: Tracks,
-        nodes: list[Node],
+        nodes: Iterable[Node],
         attributes: Attrs,
-        pixels: list[SegMask] | None = None,
+        pixels: Iterable[SegMask] | None = None,
     ):
         """Create an action to add new nodes, with optional segmentation
 
@@ -100,9 +105,9 @@ def __init__(
         self.times = attributes.get(NodeAttr.TIME.value, None)
         if NodeAttr.TIME.value in attributes:
             del user_attrs[NodeAttr.TIME.value]
-        self.positions = attributes.get(tracks.pos_attr, None)
-        if tracks.pos_attr in attributes:
-            del user_attrs[tracks.pos_attr]
+        self.positions = attributes.get(NodeAttr.POS.value, None)
+        if NodeAttr.POS.value in attributes:
+            del user_attrs[NodeAttr.POS.value]
         self.pixels = pixels
         self.attributes = user_attrs
         self._apply()
@@ -127,7 +132,10 @@ class DeleteNodes(TracksAction):
     """
 
     def __init__(
-        self, tracks: Tracks, nodes: list[Node], pixels: list[SegMask] | None = None
+        self,
+        tracks: Tracks,
+        nodes: Iterable[Node],
+        pixels: Iterable[SegMask] | None = None,
     ):
         super().__init__(tracks)
         self.nodes = nodes
@@ -170,8 +178,8 @@ class UpdateNodeSegs(TracksAction):
     def __init__(
         self,
         tracks: Tracks,
-        nodes: list[Node],
-        pixels: list[SegMask],
+        nodes: Iterable[Node],
+        pixels: Iterable[SegMask],
         added: bool = True,
     ):
         """
@@ -211,13 +219,13 @@ class UpdateNodeAttrs(TracksAction):
     def __init__(
         self,
         tracks: Tracks,
-        nodes: list[Node],
+        nodes: Iterable[Node],
         attrs: Attrs,
     ):
         """
         Args:
             tracks (Tracks): The tracks to update the node attributes for
-            nodes (list[Node]): The nodes to update the attributes for
+            nodes (Iterable[Node]): The nodes to update the attributes for
             attrs (Attrs): A mapping from attribute name to list of new attribute values
                 for the given nodes.
 
@@ -257,7 +265,7 @@ def _apply(self):
 class AddEdges(TracksAction):
     """Action for adding new edges"""
 
-    def __init__(self, tracks: Tracks, edges: list[Edge]):
+    def __init__(self, tracks: Tracks, edges: Iterable[Edge]):
         super().__init__(tracks)
         self.edges = edges
         self._apply()
@@ -277,7 +285,7 @@ def _apply(self):
 class DeleteEdges(TracksAction):
     """Action for deleting edges"""
 
-    def __init__(self, tracks: Tracks, edges: list[Edge]):
+    def __init__(self, tracks: Tracks, edges: Iterable[Edge]):
         super().__init__(tracks)
         self.edges = edges
         self._apply()
@@ -303,6 +311,7 @@ def __init__(self, tracks: SolutionTracks, start_node: Node, track_id: int):
             track_id (int): The new track id to assign.
         """
         super().__init__(tracks)
+        self.tracks: SolutionTracks = tracks
         self.start_node = start_node
         self.old_track_id = self.tracks.get_track_id(start_node)
         self.new_track_id = track_id

src/funtracks/data_model/graph_attributes.py

@@ -7,12 +7,12 @@ class NodeAttr(Enum):
     implementations of commonly used ones, listed here.
     """
 
-    POS: str = "pos"
-    TIME: str = "time"
-    SEG_ID: str = "seg_id"
-    SEG_HYPO: str = "seg_hypo"
-    AREA: str = "area"
-    TRACK_ID: str = "track_id"
+    POS = "pos"
+    TIME = "time"
+    SEG_ID = "seg_id"
+    SEG_HYPO = "seg_hypo"
+    AREA = "area"
+    TRACK_ID = "track_id"
 
 
 class EdgeAttr(Enum):
@@ -21,7 +21,7 @@ class EdgeAttr(Enum):
     implementations of commonly used ones, listed here.
     """
 
-    IOU: str = "iou"
+    IOU = "iou"
 
 
 class NodeType(Enum):

src/funtracks/data_model/solution_tracks.py

@@ -150,6 +150,10 @@ def add_nodes(
         positions: np.ndarray | None = None,
         attrs: Attrs | None = None,
     ):
+        if attrs is None:
+            raise ValueError(
+                f"Node attributes cannot be None, must contain {NodeAttr.TRACK_ID.value}"
+            )
         # overriding add_nodes to add new nodes to the track_id_to_node mapping
         super().add_nodes(nodes, times, positions, attrs)
         for node, track_id in zip(nodes, attrs[NodeAttr.TRACK_ID.value], strict=True):