bug fixes, added tests for AgentsDF

mesa_frames/abstract/agents.py

@@ -3,12 +3,11 @@
 from abc import abstractmethod
 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from contextlib import suppress
+from typing import TYPE_CHECKING, Literal
 
 from numpy.random import Generator
 from typing_extensions import Any, Self, overload
 
-from typing import TYPE_CHECKING, Literal
-
 from mesa_frames.abstract.mixin import CopyMixin
 from mesa_frames.types import BoolSeries, DataFrame, IdsLike, Index, MaskLike, Series
 
@@ -86,7 +85,7 @@ def discard(self, agents, inplace: bool = True) -> Self:
         ----------
         Self
         """
-        with suppress(KeyError):
+        with suppress(KeyError, ValueError):
             return self.remove(agents, inplace=inplace)
         return self._get_obj(inplace)
 
@@ -363,22 +362,20 @@ def sort(
     def __add__(self, other) -> Self:
         return self.add(agents=other, inplace=False)
 
-    def __contains__(self, id: int) -> bool:
+    def __contains__(self, agents: int | IdsLike | AgentSetDF) -> bool:
         """Check if an agent is in the AgentContainer.
 
         Parameters
         ----------
-        id : Hashable
-            The ID(s) to check for.
+        id : int | IdsLike | AgentSetDF
+            The ID(s) or AgentSetDF to check for.
 
         Returns
         -------
         bool
             True if the agent is in the AgentContainer, False otherwise.
         """
-        if not isinstance(id, int):
-            raise TypeError("id must be an integer")
-        return self.contains(ids=id)
+        return self.contains(agents=agents)
 
     def __getitem__(
         self,
@@ -511,7 +508,7 @@ def __getattr__(self, name: str) -> Any | dict[str, Any]:
         """
 
     @abstractmethod
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> Iterator[dict[str, Any]]:
         """Iterate over the agents in the AgentContainer.
 
         Returns
@@ -973,15 +970,12 @@ def __getitem__(
         ),
     ) -> Series | DataFrame:
         attr = super().__getitem__(key)
-        assert isinstance(attr, (Series, DataFrame))
+        assert isinstance(attr, (Series, DataFrame, Index))
         return attr
 
     def __len__(self) -> int:
         return len(self._agents)
 
-    def __iter__(self) -> Iterator:
-        return iter(self._agents)
-
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}\n {str(self._agents)}"
 

mesa_frames/abstract/mixin.py

@@ -26,7 +26,7 @@ class CopyMixin(ABC):
         _description_
     """
 
-    _copy_with_method: dict[str, tuple[str, list[str]]]
+    _copy_with_method: dict[str, tuple[str, list[str]]] = {}
     _copy_only_reference: list[str] = [
         "_model",
     ]
@@ -38,6 +38,7 @@ def copy(
         self,
         deep: bool = False,
         memo: dict | None = None,
+        skip: list[str] | None = None,
     ) -> Self:
         """Create a copy of the Class.
 
@@ -48,10 +49,12 @@ def copy(
             If True, all attributes of the AgentContainer will be recursively copied (except attributes in self._copy_reference_only).
             If False, only the top-level attributes will be copied.
             Defaults to False.
-
         memo : dict | None, optional
             A dictionary used to track already copied objects during deep copy.
             Defaults to None.
+        skip : list[str] | None, optional
+            A list of attribute names to skip during the copy process.
+            Defaults to None.
 
         Returns
         -------
@@ -61,6 +64,9 @@ def copy(
         cls = self.__class__
         obj = cls.__new__(cls)
 
+        if skip is None:
+            skip = []
+
         if deep:
             if not memo:
                 memo = {}
@@ -71,13 +77,15 @@ def copy(
                 for k, v in attributes.items()
                 if k not in self._copy_with_method
                 and k not in self._copy_only_reference
+                and k not in skip
             ]
         else:
             [
                 setattr(obj, k, copy(v))
                 for k, v in self.__dict__.items()
                 if k not in self._copy_with_method
                 and k not in self._copy_only_reference
+                and k not in skip
             ]
 
         # Copy attributes with a reference only