refactor: streamline key management methods in AgentSetRegistry and enforce model consistency

Author: adamamer20

mesa_frames/abstract/agentsetregistry.py

@@ -538,6 +538,13 @@ def __setitem__(
         - For name keys, the key is authoritative for the assigned set's name
         - For index keys, collisions on a different entry's name must raise
         """
+        if value.model is not self.model:
+            raise TypeError("Assigned AgentSet must belong to the same model")
+        if isinstance(key, (int, str)):
+            # Delegate to replace() so subclasses centralize invariant handling.
+            self.replace({key: value}, inplace=True, atomic=True)
+            return
+        raise TypeError("Key must be int index or str name")
 
     @abstractmethod
     def __getattr__(self, name: str) -> Any | dict[str, Any]:
@@ -568,14 +575,23 @@ def __str__(self) -> str:
         """Get a string representation of the AgentSets in the registry."""
         ...
 
-    @abstractmethod
     def keys(
         self, *, key_by: KeyBy = "name"
     ) -> Iterable[str | int | type[mesa_frames.abstract.agentset.AbstractAgentSet]]:
         """Iterate keys for contained AgentSets (by name|index|type)."""
-        ...
+        if key_by == "index":
+            yield from range(len(self))
+            return
+        if key_by == "type":
+            for agentset in self:
+                yield type(agentset)
+            return
+        if key_by != "name":
+            raise ValueError("key_by must be 'name'|'index'|'type'")
+        for agentset in self:
+            if agentset.name is not None:
+                yield agentset.name
 
-    @abstractmethod
     def items(
         self, *, key_by: KeyBy = "name"
     ) -> Iterable[
@@ -585,12 +601,23 @@ def items(
         ]
     ]:
         """Iterate (key, AgentSet) pairs for contained sets."""
-        ...
+        if key_by == "index":
+            for idx, agentset in enumerate(self):
+                yield idx, agentset
+            return
+        if key_by == "type":
+            for agentset in self:
+                yield type(agentset), agentset
+            return
+        if key_by != "name":
+            raise ValueError("key_by must be 'name'|'index'|'type'")
+        for agentset in self:
+            if agentset.name is not None:
+                yield agentset.name, agentset
 
-    @abstractmethod
     def values(self) -> Iterable[mesa_frames.abstract.agentset.AbstractAgentSet]:
         """Iterate contained AgentSets (values view)."""
-        ...
+        yield from self
 
     @property
     def model(self) -> mesa_frames.concrete.model.Model:

mesa_frames/concrete/agentsetregistry.py

@@ -669,78 +669,9 @@ def __repr__(self) -> str:
     def __reversed__(self) -> Iterator[AgentSet]:
         return reversed(self._agentsets)
 
-    def __setitem__(self, key: int | str, value: AgentSet) -> None:
-        """Assign/replace a single AgentSet at an index or name.
-
-        Enforces name uniqueness and model consistency.
-        """
-        if value.model is not self.model:
-            raise TypeError("Assigned AgentSet must belong to the same model")
-        if isinstance(key, int):
-            if value.name is not None:
-                for i, s in enumerate(self._agentsets):
-                    if i != key and s.name == value.name:
-                        raise ValueError(
-                            f"Duplicate agent set name disallowed: {value.name}"
-                        )
-            self._agentsets[key] = value
-        elif isinstance(key, str):
-            try:
-                value.rename(key)
-            except Exception:
-                if hasattr(value, "_name"):
-                    setattr(value, "_name", key)
-            idx = None
-            for i, s in enumerate(self._agentsets):
-                if s.name == key:
-                    idx = i
-                    break
-            if idx is None:
-                self._agentsets.append(value)
-            else:
-                self._agentsets[idx] = value
-        else:
-            raise TypeError("Key must be int index or str name")
-        # Recompute ids cache
-        self._recompute_ids()
-
     def __str__(self) -> str:
         return "\n".join([str(agentset) for agentset in self._agentsets])
 
-    def keys(self, *, key_by: KeyBy = "name") -> Iterable[Any]:
-        if key_by not in ("name", "index", "type"):
-            raise ValueError("key_by must be 'name'|'index'|'type'")
-        if key_by == "index":
-            yield from range(len(self._agentsets))
-            return
-        if key_by == "type":
-            for s in self._agentsets:
-                yield type(s)
-            return
-        # name
-        for s in self._agentsets:
-            if s.name is not None:
-                yield s.name
-
-    def items(self, *, key_by: KeyBy = "name") -> Iterable[tuple[Any, AgentSet]]:
-        if key_by not in ("name", "index", "type"):
-            raise ValueError("key_by must be 'name'|'index'|'type'")
-        if key_by == "index":
-            for i, s in enumerate(self._agentsets):
-                yield i, s
-            return
-        if key_by == "type":
-            for s in self._agentsets:
-                yield type(s), s
-            return
-        # name
-        for s in self._agentsets:
-            if s.name is not None:
-                yield s.name, s
-
-    def values(self) -> Iterable[AgentSet]:
-        return iter(self._agentsets)
-
     @property
     def ids(self) -> pl.Series:
         """Public view of all agent unique_id values across contained sets."""