Skip to content

Conversation

@SiddharthBansal007
Copy link

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

@github-actions

This comment was marked as off-topic.

@EwoutH
Copy link
Member

EwoutH commented Nov 12, 2025

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?

@EwoutH EwoutH added the maintenance Release notes label label Nov 12, 2025
@SiddharthBansal007
Copy link
Author

@EwoutH Thanks, i have made the following changes

Copy link
Member

@EwoutH EwoutH left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Few minor points


import numpy as np

if TYPE_CHECKING:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Author

@SiddharthBansal007 SiddharthBansal007 Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

mesa/agent.py Outdated
Comment on lines 63 to 64
# Preserve the more specific model type for static type checkers.
# At runtime this remains the Model instance passed in.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, I will make the neccessary changes

Comment on lines 110 to 114
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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please restore the original comments

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Made the changes, please review the new push

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

maintenance Release notes label

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants