Skip to content

Commit 5e6dd9a

Browse files
adamamer20adamamer20
committed
Enhance parameter documentation for agent handling and rename functionality across multiple classes
1 parent 03575eb commit 5e6dd9a

File tree

4 files changed

+45
-24
lines changed

4 files changed

+45
-24
lines changed

‎mesa_frames/abstract/agentset.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ def contains(self, agents: IdsLike) -> bool | BoolSeries:
9696
9797
Parameters
9898
----------
99-
agents : mesa_frames.concrete.agents.AgentSetDF | IdsLike
99+
agents : IdsLike
100100
The ID(s) to check for.
101101
102102
Returns
@@ -480,6 +480,9 @@ def rename(self, new_name: str, inplace: bool = True) -> Self:
480480
----------
481481
new_name : str
482482
Desired new name for this AgentSet.
483+
inplace : bool, optional
484+
Whether to perform the rename in place. If False, a renamed copy is
485+
returned, by default True.
483486
484487
Returns
485488
-------

‎mesa_frames/abstract/agentsetregistry.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,18 +43,20 @@ def __init__(self, model):
4343
from __future__ import annotations # PEP 563: postponed evaluation of type annotations
4444

4545
from abc import abstractmethod
46-
from collections.abc import Callable, Collection, Iterator, Sequence, Iterable
46+
from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
4747
from contextlib import suppress
4848
from typing import Any, Literal, Self, overload
4949

5050
from numpy.random import Generator
5151

5252
from mesa_frames.abstract.mixin import CopyMixin
53+
from mesa_frames.types_ import (
54+
AbstractAgentSetSelector as AgentSetSelector,
55+
)
5356
from mesa_frames.types_ import (
5457
BoolSeries,
5558
Index,
5659
KeyBy,
57-
AbstractAgentSetSelector as AgentSetSelector,
5860
Series,
5961
)
6062

@@ -112,22 +114,28 @@ def rename(
112114
113115
Parameters
114116
----------
115-
target : AgentSet | str | dict | list[tuple]
117+
target : mesa_frames.abstract.agentset.AbstractAgentSet | str | dict[mesa_frames.abstract.agentset.AbstractAgentSet | str, str] | list[tuple[mesa_frames.abstract.agentset.AbstractAgentSet | str, str]]
116118
Single target (instance or existing name) with ``new_name`` provided,
117119
or a mapping/sequence of (target, new_name) pairs for batch rename.
118120
new_name : str | None
119121
New name for single-target rename.
120-
on_conflict : {"canonicalize", "raise"}
122+
on_conflict : Literal["canonicalize", "raise"]
121123
When a desired name collides, either canonicalize by appending a
122124
numeric suffix (default) or raise ``ValueError``.
123-
mode : {"atomic", "best_effort"}
125+
mode : Literal["atomic", "best_effort"]
124126
In "atomic" mode, validate all renames before applying any. In
125127
"best_effort" mode, apply what can be applied and skip failures.
126128
127129
Returns
128130
-------
129131
Self
130132
Updated registry (or a renamed copy when ``inplace=False``).
133+
134+
Parameters
135+
----------
136+
inplace : bool, optional
137+
Whether to perform the rename in place. If False, a renamed copy is
138+
returned, by default True.
131139
"""
132140
...
133141

‎mesa_frames/abstract/space.py‎

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ def move_agents(
121121
122122
Parameters
123123
----------
124-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
124+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
125125
The agents to move
126126
pos : SpaceCoordinate | SpaceCoordinates
127127
The coordinates for each agents. The length of the coordinates must match the number of agents.
@@ -157,7 +157,7 @@ def place_agents(
157157
158158
Parameters
159159
----------
160-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
160+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
161161
The agents to place in the space
162162
pos : SpaceCoordinate | SpaceCoordinates
163163
The coordinates for each agents. The length of the coordinates must match the number of agents.
@@ -218,9 +218,9 @@ def swap_agents(
218218
219219
Parameters
220220
----------
221-
agents0 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
221+
agents0 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
222222
The first set of agents to swap
223-
agents1 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
223+
agents1 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
224224
The second set of agents to swap
225225
inplace : bool, optional
226226
Whether to perform the operation inplace, by default True
@@ -290,9 +290,9 @@ def get_directions(
290290
The starting positions
291291
pos1 : SpaceCoordinate | SpaceCoordinates | None, optional
292292
The ending positions
293-
agents0 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
293+
agents0 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
294294
The starting agents
295-
agents1 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
295+
agents1 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
296296
The ending agents
297297
normalize : bool, optional
298298
Whether to normalize the vectors to unit norm. By default False
@@ -334,9 +334,9 @@ def get_distances(
334334
The starting positions
335335
pos1 : SpaceCoordinate | SpaceCoordinates | None, optional
336336
The ending positions
337-
agents0 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
337+
agents0 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
338338
The starting agents
339-
agents1 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
339+
agents1 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
340340
The ending agents
341341
342342
Returns
@@ -369,7 +369,7 @@ def get_neighbors(
369369
The radius(es) of the neighborhood
370370
pos : SpaceCoordinate | SpaceCoordinates | None, optional
371371
The coordinates of the cell to get the neighborhood from, by default None
372-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
372+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
373373
The id of the agents to get the neighborhood from, by default None
374374
include_center : bool, optional
375375
If the center cells or agents should be included in the result, by default False
@@ -391,15 +391,17 @@ def get_neighbors(
391391
def move_to_empty(
392392
self,
393393
agents: IdsLike
394+
| AbstractAgentSet
394395
| AbstractAgentSetRegistry
396+
| Collection[AbstractAgentSet]
395397
| Collection[AbstractAgentSetRegistry],
396398
inplace: bool = True,
397399
) -> Self:
398400
"""Move agents to empty cells/positions in the space (cells/positions where there isn't any single agent).
399401
400402
Parameters
401403
----------
402-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
404+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
403405
The agents to move to empty cells/positions
404406
inplace : bool, optional
405407
Whether to perform the operation inplace, by default True
@@ -414,15 +416,17 @@ def move_to_empty(
414416
def place_to_empty(
415417
self,
416418
agents: IdsLike
419+
| AbstractAgentSet
417420
| AbstractAgentSetRegistry
421+
| Collection[AbstractAgentSet]
418422
| Collection[AbstractAgentSetRegistry],
419423
inplace: bool = True,
420424
) -> Self:
421425
"""Place agents in empty cells/positions in the space (cells/positions where there isn't any single agent).
422426
423427
Parameters
424428
----------
425-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
429+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
426430
The agents to place in empty cells/positions
427431
inplace : bool, optional
428432
Whether to perform the operation inplace, by default True
@@ -468,7 +472,7 @@ def remove_agents(
468472
469473
Parameters
470474
----------
471-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
475+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
472476
The agents to remove from the space
473477
inplace : bool, optional
474478
Whether to perform the operation inplace, by default True
@@ -703,7 +707,7 @@ def move_to_available(
703707
704708
Parameters
705709
----------
706-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry]
710+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry]
707711
The agents to move to available cells/positions
708712
inplace : bool, optional
709713
Whether to perform the operation inplace, by default True
@@ -856,7 +860,9 @@ def get_neighborhood(
856860
radius: int | float | Sequence[int] | Sequence[float] | ArrayLike,
857861
pos: DiscreteCoordinate | DiscreteCoordinates | None = None,
858862
agents: IdsLike
863+
| AbstractAgentSet
859864
| AbstractAgentSetRegistry
865+
| Collection[AbstractAgentSet]
860866
| Collection[AbstractAgentSetRegistry] = None,
861867
include_center: bool = False,
862868
) -> DataFrame:
@@ -870,7 +876,7 @@ def get_neighborhood(
870876
The radius(es) of the neighborhoods
871877
pos : DiscreteCoordinate | DiscreteCoordinates | None, optional
872878
The coordinates of the cell(s) to get the neighborhood from
873-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry], optional
879+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry], optional
874880
The agent(s) to get the neighborhood from
875881
include_center : bool, optional
876882
If the cell in the center of the neighborhood should be included in the result, by default False
@@ -1040,7 +1046,7 @@ def _sample_cells(
10401046
The number of cells to sample. If None, samples the maximum available.
10411047
with_replacement : bool
10421048
If the sampling should be with replacement
1043-
condition : Callable[[DiscreteSpaceCapacity], BoolSeries]
1049+
condition : Callable[[DiscreteSpaceCapacity], BoolSeries | np.ndarray]
10441050
The condition to apply on the capacity
10451051
respect_capacity : bool, optional
10461052
If the capacity should be respected in the sampling.
@@ -1659,9 +1665,9 @@ def _calculate_differences(
16591665
The starting positions
16601666
pos1 : GridCoordinate | GridCoordinates | None
16611667
The ending positions
1662-
agents0 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None
1668+
agents0 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None
16631669
The starting agents
1664-
agents1 : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None
1670+
agents1 : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None
16651671
The ending agents
16661672
16671673
Returns
@@ -1756,7 +1762,7 @@ def _get_df_coords(
17561762
----------
17571763
pos : GridCoordinate | GridCoordinates | None, optional
17581764
The positions to get the DataFrame from, by default None
1759-
agents : IdsLike | AbstractAgentSetRegistry | Collection[AbstractAgentSetRegistry] | None, optional
1765+
agents : IdsLike | AbstractAgentSet | AbstractAgentSetRegistry | Collection[AbstractAgentSet] | Collection[AbstractAgentSetRegistry] | None, optional
17601766
The agents to get the DataFrame from, by default None
17611767
check_bounds: bool, optional
17621768
If the positions should be checked for out-of-bounds in non-toroidal grids, by default True

‎mesa_frames/concrete/agentset.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,10 @@ def rename(self, new_name: str, inplace: bool = True) -> Self:
111111
new_name : str
112112
Desired new name.
113113
114+
inplace : bool, optional
115+
Whether to perform the rename in place. If False, a renamed copy is
116+
returned, by default True.
117+
114118
Returns
115119
-------
116120
Self

0 commit comments

Comments
 (0)