fix #2716: missing agent type hints #2885

SiddharthBansal007 · 2025-11-12T16:58:35Z

Add generic typing to Agent, AgentSet, and Model classes

Implement TypeVar-based generics to enable static type checkers to reason about
concrete agent and model types throughout the codebase. This resolves typing gaps
where model.agents returned untyped AgentSet[Agent] and self.model did not
preserve model-specific type information.

Key changes:

Agent[M] now carries model type information through self.model property
AgentSet[A] preserves agent type through iteration, indexing, and selections
Model[A] is parameterized by its agent type with typed agents property
All method signatures updated to propagate type information correctly
Backward compatible: untyped code continues to work without changes

Type checking benefits:

IDEs can now provide accurate autocomplete for model attributes in agents
Static type checkers (mypy, pyright) can verify agent/model interactions
Type information flows through AgentSet operations (select, do, iteration)
Model-specific attributes are properly typed when using Model[ConcreteModel]

Implementation details:

Uses TYPE_CHECKING blocks to avoid circular import issues
Runtime TypeVars and type-checker TypeVars kept separate for compatibility
Generic[T] inheritance pattern used for all three classes
WeakKeyDictionary in AgentSet unchanged; typing layer transparent

Testing:

All 366 existing tests pass without modification
Agent and model test coverage: 99% and 91% respectively
Example file added demonstrating typed agent/model usage
No runtime behavior changes; purely typing improvements

Addresses: Typing improvements for better IDE support and static analysis

EwoutH · 2025-11-12T18:49:05Z

Thanks for tackling this issue, improved type hints are definitely useful.

The implementation looks solid, but I wanted to mention: since we recently dropped Python 3.11 support and now require 3.12+, we can actually use the newer PEP 695 syntax instead of the Generic[T] + TypeVar approach.

Python 3.12 introduced cleaner syntax for generics that eliminates the need for the conditional TYPE_CHECKING blocks and the manual TypeVar declarations. Here's what it would look like:

Instead of:

if TYPE_CHECKING:
    M = TypeVar("M", bound=Model)
    A = TypeVar("A", bound="Agent")
else:
    M = TypeVar("M")
    A = TypeVar("A")

class Agent(Generic[M]):  # noqa: UP046
    def __init__(self, model: M, *args, **kwargs) -> None:
        self.model: M = model

We can now write:

class Agent[M: Model]:
    def __init__(self, model: M, *args, **kwargs) -> None:
        self.model: M = model

Same goes for AgentSet and the method-level generics:

class AgentSet[A: Agent](MutableSet[A], Sequence[A]):
    # ...

@classmethod
def create_agents[T: Agent](cls: type[T], model: M, n: int, *args, **kwargs) -> AgentSet[T]:
    # ...

The PEP 695 syntax is more readable, eliminates the runtime/type-checking split, and is the idiomatic way to do this in Python 3.12+. The # noqa: UP046 comments you added are actually there because Ruff knows this syntax exists but couldn't be used when supporting 3.11 - now we can.

Would you be up for refactoring this to use the new syntax?

SiddharthBansal007 · 2025-11-12T19:49:02Z

@EwoutH Thanks, i have made the following changes

EwoutH

Thanks! Few minor points

EwoutH · 2025-11-13T09:57:03Z

mesa/agent.py


 import numpy as np

 if TYPE_CHECKING:


Could you check if we still need this? (honestly I don't know)

As far as I went through the codebase, I think numpy is needed
Example usecase (Line 146)

def rng(self) -> np.random.Generator: """Return a seeded np.random rng.""" return self.model.rng

for the TYPE_CHECKING, i have used it to prevent circular imports, do let me know if i am wrong in any of the both cases

EwoutH · 2025-11-13T09:57:26Z

mesa/agent.py

+        # Preserve the more specific model type for static type checkers.
+        # At runtime this remains the Model instance passed in.


Comments are not needed

Suggested change

# Preserve the more specific model type for static type checkers.

# At runtime this remains the Model instance passed in.

ok, I will make the neccessary changes

EwoutH · 2025-11-13T09:58:53Z

mesa/model.py

-            type[Agent], AgentSet
-        ] = {}  # a dict with an agentset for each class of agents
-        self._all_agents = AgentSet(
-            [], random=self.random
-        )  # an agenset with all agents


please restore the original comments

Made the changes, please review the new push

SiddharthBansal007 · 2025-11-16T21:53:04Z

@EwoutH incase you missed the changes, could you look upon the changes

EwoutH · 2025-12-03T11:57:24Z

mesa/agent.py

+    def create_agents[T: Agent](
+        cls: type[T], model: Model, n: int, *args, **kwargs
+    ) -> AgentSet[T]:


Why is T used here for Agent? Wouldn't A be more logical?

Please motivate first. But if you agree:

Suggested change

def create_agents[T: Agent](

cls: type[T], model: Model, n: int, *args, **kwargs

) -> AgentSet[T]:

def create_agents[A: Agent](

cls: type[A], model: Model, n: int, *args, **kwargs

) -> AgentSet[A]:

EwoutH · 2025-12-03T11:59:52Z

mesa/agent.py

        inplace: bool = False,
-        agent_type: type[Agent] | None = None,
+        agent_type: type[A] | None = None,
    ) -> AgentSet:


We can also add a type on the return, right?

Suggested change

) -> AgentSet:

) -> AgentSet[A]:

EwoutH · 2025-12-03T12:02:07Z

mesa/model.py

+if TYPE_CHECKING:
+    pass
+


This empty block can be removed, right?

Suggested change

if TYPE_CHECKING:

pass

EwoutH · 2025-12-03T12:02:57Z

mesa/model.py


 # mypy
-from typing import Any
+from typing import TYPE_CHECKING, Any


Can be removed I think.

Suggested change

from typing import TYPE_CHECKING, Any

from typing import Any

EwoutH · 2025-12-03T12:05:42Z

mesa/agent.py

        return (
            AgentSet(sorted_agents, self.random)
            if not inplace
            else self._update(sorted_agents)


In the sort function, [A] can also be added.

def sort( self, key: Callable[[Agent], Any] | str, ascending: bool = False, inplace: bool = False, ) -> AgentSet[A]:

EwoutH · 2025-12-03T12:08:51Z

mesa/model.py

        self.step = self._wrapped_step

        # setup agent registration data structures
        self._agents = {}  # the hard references to all agents in the model


I think this is possible:

Suggested change

self._agents = {} # the hard references to all agents in the model

self._agents = dict[A, None] # the hard references to all agents in the model

for more information, see https://pre-commit.ci

EwoutH · 2025-12-03T12:11:41Z

@SiddharthBansal007 sorry for the late review. I've a few more (minor) comments, can you address them?

Please don't adopt them without thinking, but critically check.

This comment was marked as off-topic.

Sign in to view

EwoutH added the maintenance Release notes label label Nov 12, 2025

EwoutH reviewed Nov 13, 2025

View reviewed changes

EwoutH reviewed Dec 3, 2025

View reviewed changes

SiddharthBansal007 and others added 4 commits December 3, 2025 13:10

fix projectmesa#2716: missing agent type hints

7d0c6ee

Addressed the comment and updated the codebase

01ab92b

Addressed the review and updated the comments

78dee0c

[pre-commit.ci] auto fixes from pre-commit.com hooks

638dc87

for more information, see https://pre-commit.ci

EwoutH force-pushed the fix/#2716-missing-agent-type-hints branch from 129f3c2 to 638dc87 Compare December 3, 2025 12:10

EwoutH added enhancement Release notes label and removed maintenance Release notes label labels Dec 3, 2025

Merge branch 'main' into fix/projectmesa#2716-missing-agent-type-hints

78d02d4

		# Preserve the more specific model type for static type checkers.
		# At runtime this remains the Model instance passed in.

	self._agents = {} # the hard references to all agents in the model
	self._agents = dict[A, None] # the hard references to all agents in the model

Uh oh!

fix #2716: missing agent type hints #2885

Are you sure you want to change the base?

fix #2716: missing agent type hints #2885

Conversation

SiddharthBansal007 commented Nov 12, 2025

Uh oh!

This comment was marked as off-topic.

EwoutH commented Nov 12, 2025

Uh oh!

SiddharthBansal007 commented Nov 12, 2025

Uh oh!

EwoutH left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SiddharthBansal007 Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SiddharthBansal007 commented Nov 16, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EwoutH commented Dec 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SiddharthBansal007 Nov 13, 2025 •

edited

Loading