From 3fec08cc582b208b794c1eb715ce5f77bb15c2da Mon Sep 17 00:00:00 2001
From: rht <rhtbot@protonmail.com>
Date: Sun, 7 Jan 2024 09:17:51 -0500
Subject: [PATCH] model: Implement internal clock

---
 mesa/model.py               | 12 +++++++++++-
 tests/test_batch_run.py     |  1 +
 tests/test_datacollector.py |  1 +
 tests/test_lifespan.py      |  1 +
 tests/test_time.py          |  1 +
 tests/test_visualization.py |  1 +
 6 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mesa/model.py b/mesa/model.py
index ca0c653cd30..cd36f29115c 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -12,11 +12,13 @@
 from collections import defaultdict
 
 # mypy
-from typing import Any
+from typing import Any, Union
 
 from mesa.agent import Agent, AgentSet
 from mesa.datacollection import DataCollector
 
+TimeT = Union[float, int]
+
 
 class Model:
     """Base class for models in the Mesa ABM library.
@@ -67,6 +69,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.current_id = 0
         self._agents: defaultdict[type, dict] = defaultdict(dict)
 
+        self.steps = 0
+        self.time: TimeT = 0  # the model's clock
+
         # Warning flags for current experimental features. These make sure a warning is only printed once per model.
         self.agentset_experimental_warning_given = False
 
@@ -97,6 +102,11 @@ def run_model(self) -> None:
     def step(self) -> None:
         """A single step. Fill in here."""
 
+    def step_and_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
diff --git a/tests/test_batch_run.py b/tests/test_batch_run.py
index 1557666f34f..7156ba02f8b 100644
--- a/tests/test_batch_run.py
+++ b/tests/test_batch_run.py
@@ -85,6 +85,7 @@ def get_local_model_param(self):
     def step(self):
         self.datacollector.collect(self)
         self.schedule.step()
+        self.step_and_advance_time()
 
 
 def test_batch_run():
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index a56ca47745c..63f3d866c87 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -89,6 +89,7 @@ def test_model_calc_comp(self, input1, input2):
     def step(self):
         self.schedule.step()
         self.datacollector.collect(self)
+        self.step_and_advance_time()
 
 
 class TestDataCollector(unittest.TestCase):
diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py
index cfd60cdeb74..cb900bb7469 100644
--- a/tests/test_lifespan.py
+++ b/tests/test_lifespan.py
@@ -43,6 +43,7 @@ def step(self):
                 self.schedule.add(
                     FiniteLifeAgent(self.next_id(), self.agent_lifetime, self)
                 )
+        self.step_and_advance_time()
 
     def run_model(self, step_count=100):
         for _ in range(step_count):
diff --git a/tests/test_time.py b/tests/test_time.py
index a440d502566..cc2f2a37c92 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -96,6 +96,7 @@ def __init__(self, shuffle=False, activation=STAGED, enable_kill_other_agent=Fal
 
     def step(self):
         self.schedule.step()
+        self.step_and_advance_time()
 
     def model_stage(self):
         self.log.append("model_stage")
diff --git a/tests/test_visualization.py b/tests/test_visualization.py
index ca9f0949558..77d9aecff9c 100644
--- a/tests/test_visualization.py
+++ b/tests/test_visualization.py
@@ -49,6 +49,7 @@ def __init__(self, width, height, key1=103, key2=104):
 
     def step(self):
         self.schedule.step()
+        self.step_and_advance_time()
 
 
 class TestModularServer(TestCase):