model: Automatically increase time and step count

This pull request adds automatically increasing of time and step counters in the Mesa model, even if the users overwrites the Model.step() method.

Author: EwoutH

mesa/model.py

@@ -83,6 +83,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._steps: int = 0
         self._time: TimeT = 0  # the model's clock
 
+        # Wrap the user-defined step method
+        if hasattr(self, "step"):
+            self._user_step = self.step
+            self.step = self._wrapped_step
+
+    def _wrapped_step(self, *args: Any, **kwargs: Any) -> None:
+        """Automatically increments time and steps after calling the user's step method."""
+        # Automatically increment time and step counters
+        self._time += kwargs.get("time", 1)
+        self._steps += kwargs.get("step", 1)
+        # Call the original user-defined step method
+        if self._user_step:
+            self._user_step(*args, **kwargs)
+
     @property
     def agents(self) -> AgentSet:
         """Provides an AgentSet of all agents in the model, combining agents from all types."""
@@ -180,11 +194,6 @@ def run_model(self) -> None:
     def step(self) -> None:
         """A single step. Fill in here."""
 
-    def _advance_time(self, deltat: TimeT = 1):
-        """Increment the model's steps counter and clock."""
-        self._steps += 1
-        self._time += deltat
-
     def next_id(self) -> int:
         """Return the next unique ID for agents, increment current_id"""
         self.current_id += 1

mesa/time.py

@@ -69,8 +69,6 @@ def __init__(self, model: Model, agents: Iterable[Agent] | None = None) -> None:
         self.model = model
         self.steps = 0
         self.time: TimeT = 0
-        self._original_step = self.step
-        self.step = self._wrapped_step
 
         if agents is None:
             agents = []
@@ -113,11 +111,6 @@ def step(self) -> None:
         self.steps += 1
         self.time += 1
 
-    def _wrapped_step(self):
-        """Wrapper for the step method to include time and step updating."""
-        self._original_step()
-        self.model._advance_time()
-
     def get_agent_count(self) -> int:
         """Returns the current number of agents in the queue."""
         return len(self._agents)

tests/test_batch_run.py

@@ -83,8 +83,8 @@ def get_local_model_param(self):
         return 42
 
     def step(self):
-        self.datacollector.collect(self)
         self.schedule.step()
+        self.datacollector.collect(self)
 
 
 def test_batch_run():