Enable type-hinted Scenario subclassing and fix Model generic typing#3168
Enable type-hinted Scenario subclassing and fix Model generic typing#3168
Conversation
Enhanced the Scenario class to support both simple instantiation and type-hinted subclassing, allowing class-level defaults to be overridden by instance parameters. Improved documentation and enforced immutability of parameters during model execution. These changes make scenario configuration more flexible and robust for complex models.
The Model class now accepts a second generic type parameter S for the scenario type.
Introduces tests to verify that Scenario can be subclassed with type-hinted attributes, that scenario subclasses integrate correctly with Model, and that each model receives a fresh scenario instance to avoid shared state.
|
I don't have time to review this at the moment. Just one question: why not dataclasses? |
|
Performance benchmarks:
|
Great question. I did consider (and try) that at first, but there are a few problems. The This implementation uses |
| # Collect class-level annotated attributes as defaults | ||
| defaults = {} | ||
| for cls in reversed(self.__class__.__mro__): | ||
| if cls is Scenario or cls is object: | ||
| continue | ||
| # Get type-annotated attributes | ||
| annotations = getattr(cls, "__annotations__", {}) | ||
| for key in annotations: | ||
| if hasattr(cls, key) and not key.startswith("_"): |
There was a problem hiding this comment.
This is a class-level thing, so it's a bit wasteful to do it in the __init__, which is instance-specific. There is probably some metaclass or other trick to doing this only once when creating the class.
From a quick search, I think that a simple scenario meta class with __init_subclass__ magic might do the trick
Below just some snippets to get started.
class ScenarioMeta:
def __new__meta, name, bases, class_dict):
# The __new__ method of metaclasses is run after the class statement’s entire body has been processed.
# the mro stuff might be handled in here
class Scenario(metaclass=ScenarioMeta):
# existing code
# another option, can be done separately or instead of the metaclass approach
def __init_subclass__(cls):
# This method is called whenever the containing class is subclassed.
#cls is then the new subclass. If defined as a normal instance
#method, this method is implicitly converted to a class method.There was a problem hiding this comment.
Good catch. I don’t think we need the metaclass, but can just directly use __init_subclass__ :
class Scenario:
def __init_subclass__(cls):
"""Called once when a subclass is created."""
# Collect defaults once and cache on the class
defaults = {}
for base in reversed(cls.__mro__):
if base is Scenario or base is object:
continue
annotations = getattr(base, '__annotations__', {})
for key in annotations:
if hasattr(base, key) and not key.startswith('_'):
defaults[key] = getattr(base, key)
# Cache on the class itself
cls._scenario_defaults = defaults
def __init__(self, *, rng=None, **kwargs):
# Just use the cached defaults
self.__dict__.update(self.__class__._scenario_defaults)
self.__dict__.update(kwargs)
self.__dict__['rng'] = rngThere was a problem hiding this comment.
It works we a few minor additions. I have pushed the change to this PR.
| def __len__(self): | ||
| """Return number of parameters.""" |
There was a problem hiding this comment.
so the idea is that the __dict__ contains the scenario paramters, while we use __slots__ for the overarching registration stuff (such as scenario_id)?
There was a problem hiding this comment.
Yes, indeed. The __slots__ = ("__dict__", "_scenario_id", "model") part uses slots for Mesa’s fixed internal attributes (_scenario_id for unique identification and model for the model reference), explicitly including __dict__ to allow unlimited dynamic scenario parameters (like citizen_density, cop_vision, etc.).
This gives the memory efficiency of __slots__ for the structural attributes that every scenario needs, while maintaining the flexibility of __dict__ for user-defined parameters that vary between different scenario types.
|
If we fix pre-commit I think we're ready to merge. |
2 participants
Summary
This PR enhances the
Scenarioclass to support type-hinted subclassing and makesModelgeneric over scenario types, providing better IDE support and type safety for scenario parameters.Took some time to get everything working at the same time, but now it does! I did test it on an updated version of the
EpsteinCivilViolencemodel.Motivation
Currently, when using scenarios in Mesa models, there are two issues:
Scenariowith typed attributes, IDEs don't recognize these attributes as instance attributes, causing false warnings and lack of autocomplete.Model, the type is lost andmodel.scenariois typed as baseScenario, losing all the specific parameter type information.modelreference remains set to the first model.Changes
Scenarioclass for type-hinted subclassing__init_subclass__to collect type-annotated class attributes as defaults once per class (not per instance)_scenario_defaultsand applied during__init____slots__ = ("__dict__", "_scenario_id", "model")pattern:__slots__for Mesa's fixed internal attributes (_scenario_id,model)__dict__for unlimited dynamic scenario parametersModelgeneric over scenario typeS: Scenariotype parameter toModel[A: Agent, S: Scenario]scenarioproperty and__init__to use generic typeS__setattr__to checkself.modelexists before accessingmodel.runningNo breaking changes.
Design decisions
Why not dataclasses?
The
@dataclassdecorator would require all fields to be predefined, preventing the flexible**kwargspattern that allows users to pass arbitrary parameters without subclassing (e.g.,Scenario(rng=42, any_param=value)). The current implementation supports both simpleScenario(**kwargs)for quick prototyping and type-hinted subclassing for production code, without the complexities of dataclass inheritance (field ordering requirements,field()specifications).Why
__init_subclass__over__init__?Per reviewer feedback: Collecting class-level annotations in
__init__is wasteful since it happens for every instance. Using__init_subclass__performs MRO traversal once when the subclass is created, caching the results in_scenario_defaultsfor all instances. This is the idiomatic Python 3.6+ pattern for class customization at subclass creation time.Testing
Added three new test cases:
test_scenario_subclassing: Verifies type-hinted subclassing works correctlytest_scenario_subclass_with_model: Verifies scenario subclasses work with Modeltest_scenario_fresh_instance_per_model: Ensures no shared state issuesExample Usage
Follow-up on #3103, and issues found in #3167 regarding scenario type hints and shared instance problems.