Skip to content

Commit ead9cd1

Browse files
make typing behavior of AgentSet.get explicit (#2293)
* make typing behavior of AgentSet.get explicit * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent eaf7457 commit ead9cd1

File tree

1 file changed

+19
-3
lines changed

1 file changed

+19
-3
lines changed

mesa/agent.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from random import Random
2121

2222
# mypy
23-
from typing import TYPE_CHECKING, Any, Literal
23+
from typing import TYPE_CHECKING, Any, Literal, overload
2424

2525
if TYPE_CHECKING:
2626
# We ensure that these are not imported during runtime to prevent cyclic
@@ -348,12 +348,28 @@ def agg(self, attribute: str, func: Callable) -> Any:
348348
values = self.get(attribute)
349349
return func(values)
350350

351+
@overload
351352
def get(
352353
self,
353-
attr_names: str | list[str],
354+
attr_names: str,
354355
handle_missing: Literal["error", "default"] = "error",
355356
default_value: Any = None,
356-
) -> list[Any] | list[list[Any]]:
357+
) -> list[Any]: ...
358+
359+
@overload
360+
def get(
361+
self,
362+
attr_names: list[str],
363+
handle_missing: Literal["error", "default"] = "error",
364+
default_value: Any = None,
365+
) -> list[list[Any]]: ...
366+
367+
def get(
368+
self,
369+
attr_names,
370+
handle_missing="error",
371+
default_value=None,
372+
):
357373
"""
358374
Retrieve the specified attribute(s) from each agent in the AgentSet.
359375

0 commit comments

Comments
 (0)