diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py index 8e9bfa23..3bcc4d76 100644 --- a/mesa_frames/concrete/model.py +++ b/mesa_frames/concrete/model.py @@ -85,6 +85,20 @@ def __init__(self, seed: int | Sequence[int] | None = None) -> None: 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 @@ def step(self) -> None: """ 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. diff --git a/tests/test_modeldf.py b/tests/test_modeldf.py new file mode 100644 index 00000000..afc45405 --- /dev/null +++ b/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