Auto step counter

mesa_frames/concrete/model.py

@@ -85,6 +85,20 @@
         self.current_id = 0
         self._agents = AgentsDF(self)
         self._space = None
+        self._steps = 0
+
+        self._user_step = self.step
+        self.step = self._wrapped_step
+
+    def _wrapped_step(self) -> None:
+        """Automatically increments step counter and calls user-defined step()."""
+        self._steps += 1
+        self._user_step()
+
+    @property
+    def steps(self) -> int:
+        """Get the current step count."""
+        return self._steps
 
     def get_agents_of_type(self, agent_type: type) -> "AgentSetDF":
         """Retrieve the AgentSetDF of a specified type.
@@ -133,6 +147,17 @@
         """
         self.agents.step()
 
+    @property
+    def steps(self) -> int:
+        """Get the current step count.
+
+        Returns
+        -------
+        int
+            The current step count of the model.
+        """
+        return self._steps
+
     @property
     def agents(self) -> AgentsDF:
         """Get the AgentsDF object containing all agents in the model.

tests/test_modeldf.py

@@ -0,0 +1,37 @@
+from mesa_frames import ModelDF
+
+
+class CustomModel(ModelDF):
+    def __init__(self):
+        super().__init__()
+        self.custom_step_count = 0
+
+    def step(self):
+        self.custom_step_count += 2
+
+
+class Test_ModelDF:
+    def test_steps(self):
+        model = ModelDF()
+
+        assert model.steps == 0
+
+        model.step()
+        assert model.steps == 1
+
+        model.step()
+        assert model.steps == 2
+
+    def test_user_defined_step(self):
+        model = CustomModel()
+
+        assert model.steps == 0
+        assert model.custom_step_count == 0
+
+        model.step()
+        assert model.steps == 1
+        assert model.custom_step_count == 2
+
+        model.step()
+        assert model.steps == 2
+        assert model.custom_step_count == 4