Add generic type parameters to Agent, AgentSet, and Model

[SiddharthBansal007]

Summary

Implements PEP 695 generic typing for Agent, AgentSet, and Model classes to enable static type checkers to preserve concrete agent and model types throughout the codebase.

Motive

Previously, type checkers could not infer specific agent or model types:

• model.agents returned untyped AgentSet[Agent] instead of the specific agent type
• agent.model returned generic Model instead of the concrete model type
• Method chaining on AgentSet lost type information

This made it difficult to catch type errors at development time and reduced IDE autocomplete effectiveness.

Implementation

• Added generic type parameter M to Agent[M: Model] to track the specific model type
• Added generic type parameter A to Model[A: Agent] to track the specific agent type
• Added generic type parameter A to AgentSet[A: Agent] for type-safe collections
• Updated all method signatures to use generic type variables (filter_func: Callable[[A], bool], key: Callable[[A], Any])
• Added @overload decorators to AgentSet.__getitem__ to distinguish between int and slice indexing
• Updated internal storage type hints (_agents: dict[A, None], _agents_by_type: dict[type[A], AgentSet[A]])

Usage Examples

class MyAgent(Agent["MyModel"]):
    def custom_method(self) -> None:
        pass

class MyModel(Model[MyAgent]):
    pass

model = MyModel()
agent = MyAgent(model)

# Type checker now knows:
reveal_type(agent.model)  # MyModel (not just Model)
reveal_type(model.agents)  # AgentSet[MyAgent] (not AgentSet[Agent])
reveal_type(model.agents[0])  # MyAgent (not Agent)

# Method chaining preserves types
filtered: AgentSet[MyAgent] = model.agents.select(lambda a: a.pos is not None)
filtered[0].custom_method()  # Type checker knows custom_method exists

Additional Notes

• Requires Python 3.12+ for PEP 695 syntax
• No runtime behavior changes; purely improves static type checking
• Known limitation: Agent[M: Model] and Model[A: Agent] create a circular type bound dependency, which is acceptable for this use case
• create_agents uses independent type variable T to support subclass creation

[EwoutH]

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]

@EwoutH Thanks, i have made the following changes

[EwoutH]

Thanks! Few minor points

[SiddharthBansal007]

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

[EwoutH]

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]:

[quaquel]

In my understanding of typing, T is just the common default to use. But I am by no means an expert yet on type hinting in python.

[EwoutH]

@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.

[quaquel]

@EwoutH, I picked this up and made a bunch of changes. Can you check to see if it's good to go at this point? I'd like to get this merged.

[EwoutH]

Thanks for picking this up @quaquel. I agree it would be very useful to have this in Mesa sooner than later.

A few suggestions and questions. Please check them carefully, this is not my core area of expertise.

[quaquel]

@EwoutH, I addressed your comments. I think this is ready for now.

I want to guard against scope creep, so as long as these changes are fine, I suggest we move forward. From my own first impressions, some type warnings seem to be gone when using this PR.

[EwoutH]

Looks good to me! @SiddharthBansal007 thanks for the initial PR and @quaquel for landing it.