Skip to content

Commit

Permalink
seperate apply and do
Browse files Browse the repository at this point in the history
  • Loading branch information
quaquel committed Aug 21, 2024
1 parent f0dccf6 commit eaf0a5b
Showing 1 changed file with 34 additions and 1 deletion.
35 changes: 34 additions & 1 deletion mesa/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,7 +401,13 @@ def groupby(self, by: Callable | str, result_type: str = "agentset") -> GroupBy:


class GroupBy:
"""Helper class for AgentSet.groupby"""
"""Helper class for AgentSet.groupby
Attributes:
groups (dict): A dictionary with the group_name as key and group as values
"""

def __init__(self, groups: dict[Any, list | AgentSet]):
self.groups: dict[Any, list | AgentSet] = groups
Expand Down Expand Up @@ -436,5 +442,32 @@ def apply(self, method: Callable | str, *args, **kwargs) -> dict[Any, Any]:
else:
return {k: method(v, *args, **kwargs) for k, v in self.groups.items()}

def do(self, method: Callable | str, *args, **kwargs) -> GroupBy:
"""Apply the specified callable to each group
Args:
method (Callable, str): The callable to apply to each group,
* if ``method`` is a callable, it will be called it will be called with the group as first argument
* if ``method`` is a str, it should refer to a method on the group
Additional arguments and keyword arguments will be passed on to the callable.
Returns:
GroupBy
"""
if isinstance(method, str):
for v in self.groups.values():
getattr(v, method)(*args, **kwargs)
else:
for v in self.groups.values():
method(v, *args, **kwargs)

return self

def __iter__(self):
return iter(self.groups.items())

def __len__(self):
return len(self.groups)

0 comments on commit eaf0a5b

Please sign in to comment.