Add type hints to agent class

projectmesa · Oct 13, 2024 · e488ab1 · e488ab1
1 parent f4b9156
commit e488ab1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/schelling/model.py b/examples/schelling/model.py
@@ -6,7 +6,7 @@ class SchellingAgent(mesa.Agent):
     Schelling segregation agent
     """
 
-    def __init__(self, model, agent_type):
+    def __init__(self, model: mesa.Model, agent_type: int) -> None:
         """
         Create a new Schelling agent.
 
@@ -16,7 +16,7 @@ def __init__(self, model, agent_type):
         super().__init__(model)
         self.type = agent_type
 
-    def step(self):
+    def step(self) -> None:
         neighbors = self.model.grid.iter_neighbors(
             self.pos, moore=True, radius=self.model.radius
         )