Implement rename functionality for AgentSet and AgentSetRegistry with conflict handling

Author: adamamer20

mesa_frames/abstract/agentset.py

@@ -468,6 +468,50 @@ def random(self) -> Generator:
     def space(self) -> mesa_frames.abstract.space.Space | None:
         return self.model.space
 
+    def rename(self, new_name: str, inplace: bool = True) -> Self:
+        """Rename this AgentSet.
+
+        If this set is contained in the model's AgentSetRegistry, delegate to
+        the registry's rename implementation so that name uniqueness and
+        conflicts are handled consistently. If the set is not yet part of a
+        registry, update the local name directly.
+
+        Parameters
+        ----------
+        new_name : str
+            Desired new name for this AgentSet.
+
+        Returns
+        -------
+        Self
+            The updated AgentSet (or a renamed copy when ``inplace=False``).
+        """
+        obj = self._get_obj(inplace)
+        try:
+            # If contained in registry, delegate to it so conflicts are handled
+            if self in self.model.sets:  # type: ignore[operator]
+                # Preserve index to retrieve copy when not inplace
+                idx = None
+                try:
+                    idx = list(self.model.sets).index(self)  # type: ignore[arg-type]
+                except Exception:
+                    idx = None
+                reg = self.model.sets.rename(self, new_name, inplace=inplace)
+                if inplace:
+                    return self
+                # Non-inplace: return the corresponding set from the copied registry
+                if idx is not None:
+                    return reg[idx]  # type: ignore[index]
+                # Fallback: look up by name (may be canonicalized)
+                return reg.get(new_name)  # type: ignore[return-value]
+        except Exception:
+            # If delegation cannot be resolved, fall back to local rename
+            obj._name = new_name
+            return obj
+        # Not in a registry: local rename
+        obj._name = new_name
+        return obj
+
     def __setitem__(
         self,
         key: str

mesa_frames/abstract/agentsetregistry.py

@@ -93,6 +93,44 @@ def discard(
             return self.remove(sets, inplace=inplace)
         return self._get_obj(inplace)
 
+    @abstractmethod
+    def rename(
+        self,
+        target: (
+            mesa_frames.abstract.agentset.AbstractAgentSet
+            | str
+            | dict[mesa_frames.abstract.agentset.AbstractAgentSet | str, str]
+            | list[tuple[mesa_frames.abstract.agentset.AbstractAgentSet | str, str]]
+        ),
+        new_name: str | None = None,
+        *,
+        on_conflict: Literal["canonicalize", "raise"] = "canonicalize",
+        mode: Literal["atomic", "best_effort"] = "atomic",
+        inplace: bool = True,
+    ) -> Self:
+        """Rename AgentSets in this registry, handling conflicts.
+
+        Parameters
+        ----------
+        target : AgentSet | str | dict | list[tuple]
+            Single target (instance or existing name) with ``new_name`` provided,
+            or a mapping/sequence of (target, new_name) pairs for batch rename.
+        new_name : str | None
+            New name for single-target rename.
+        on_conflict : {"canonicalize", "raise"}
+            When a desired name collides, either canonicalize by appending a
+            numeric suffix (default) or raise ``ValueError``.
+        mode : {"atomic", "best_effort"}
+            In "atomic" mode, validate all renames before applying any. In
+            "best_effort" mode, apply what can be applied and skip failures.
+
+        Returns
+        -------
+        Self
+            Updated registry (or a renamed copy when ``inplace=False``).
+        """
+        ...
+
     @abstractmethod
     def add(
         self,

mesa_frames/concrete/agentset.py

@@ -103,7 +103,7 @@ def __init__(
         self._df = pl.DataFrame()
         self._mask = pl.repeat(True, len(self._df), dtype=pl.Boolean, eager=True)
 
-    def rename(self, new_name: str) -> str:
+    def rename(self, new_name: str, inplace: bool = True) -> Self:
         """Rename this agent set. If attached to AgentSetRegistry, delegate for uniqueness enforcement.
 
         Parameters
@@ -113,22 +113,40 @@ def rename(self, new_name: str) -> str:
 
         Returns
         -------
-        str
-            The final name used (may be canonicalized if duplicates exist).
+        Self
+            The updated AgentSet (or a renamed copy when ``inplace=False``).
 
         Raises
         ------
         ValueError
             If name conflicts occur and delegate encounters errors.
         """
+        # Respect inplace semantics consistently with other mutators
+        obj = self._get_obj(inplace)
+
         # Always delegate to the container's accessor if available through the model's sets
         # Check if we have a model and can find the AgentSetRegistry that contains this set
-        if self in self.model.sets:
-            return self.model.sets.rename(self._name, new_name)
+        try:
+            if self in self.model.sets:
+                # Save index to locate the copy on non-inplace path
+                try:
+                    idx = list(self.model.sets).index(self)  # type: ignore[arg-type]
+                except Exception:
+                    idx = None
+                reg = self.model.sets.rename(self, new_name, inplace=inplace)
+                if inplace:
+                    return self
+                if idx is not None:
+                    return reg[idx]
+                return reg.get(new_name)  # type: ignore[return-value]
+        except Exception:
+            # Fall back to local rename if delegation fails
+            obj._name = new_name
+            return obj
 
         # Set name locally if no container found
-        self._name = new_name
-        return new_name
+        obj._name = new_name
+        return obj
 
     def add(
         self,

mesa_frames/concrete/agentsetregistry.py

@@ -113,6 +113,111 @@ def add(
         obj._ids = new_ids
         return obj
 
+    def rename(
+        self,
+        target: (
+            AgentSet
+            | str
+            | dict[AgentSet | str, str]
+            | list[tuple[AgentSet | str, str]]
+        ),
+        new_name: str | None = None,
+        *,
+        on_conflict: Literal["canonicalize", "raise"] = "canonicalize",
+        mode: Literal["atomic", "best_effort"] = "atomic",
+        inplace: bool = True,
+    ) -> Self:
+        """Rename AgentSets with conflict handling.
+
+        Supports single-target ``(set | old_name, new_name)`` and batch rename via
+        dict or list of pairs. Names remain unique across the registry.
+        """
+
+        # Normalize to list of (index_in_self, desired_name) using the original registry
+        def _resolve_one(x: AgentSet | str) -> int:
+            if isinstance(x, AgentSet):
+                for i, s in enumerate(self._agentsets):
+                    if s is x:
+                        return i
+                raise KeyError("AgentSet not found in registry")
+            # name lookup on original registry
+            for i, s in enumerate(self._agentsets):
+                if s.name == x:
+                    return i
+            raise KeyError(f"Agent set '{x}' not found")
+
+        if isinstance(target, (AgentSet, str)):
+            if new_name is None:
+                raise TypeError("new_name must be provided for single rename")
+            pairs_idx: list[tuple[int, str]] = [(_resolve_one(target), new_name)]
+            single = True
+        elif isinstance(target, dict):
+            pairs_idx = [(_resolve_one(k), v) for k, v in target.items()]
+            single = False
+        else:
+            pairs_idx = [(_resolve_one(k), v) for k, v in target]
+            single = False
+
+        # Choose object to mutate
+        obj = self._get_obj(inplace)
+        # Translate indices to object AgentSets in the selected registry object
+        target_sets = [obj._agentsets[i] for i, _ in pairs_idx]
+
+        # Build the set of names that remain fixed (exclude targets' current names)
+        targets_set = set(target_sets)
+        fixed_names: set[str] = {
+            s.name
+            for s in obj._agentsets
+            if s.name is not None and s not in targets_set
+        }  # type: ignore[comparison-overlap]
+
+        # Plan final names
+        final: list[tuple[AgentSet, str]] = []
+        used = set(fixed_names)
+
+        def _canonicalize(base: str) -> str:
+            if base not in used:
+                used.add(base)
+                return base
+            counter = 1
+            cand = f"{base}_{counter}"
+            while cand in used:
+                counter += 1
+                cand = f"{base}_{counter}"
+            used.add(cand)
+            return cand
+
+        errors: list[Exception] = []
+        for aset, (_idx, desired) in zip(target_sets, pairs_idx):
+            if on_conflict == "canonicalize":
+                final_name = _canonicalize(desired)
+                final.append((aset, final_name))
+            else:  # on_conflict == 'raise'
+                if desired in used:
+                    err = ValueError(
+                        f"Duplicate agent set name disallowed: '{desired}'"
+                    )
+                    if mode == "atomic":
+                        errors.append(err)
+                    else:
+                        # best_effort: skip this rename
+                        continue
+                else:
+                    used.add(desired)
+                    final.append((aset, desired))
+
+        if errors and mode == "atomic":
+            # Surface first meaningful error
+            raise errors[0]
+
+        # Apply renames
+        for aset, newn in final:
+            # Set the private name directly to avoid external uniqueness hooks
+            if hasattr(aset, "_name"):
+                aset._name = newn  # type: ignore[attr-defined]
+
+        return obj
+
     def replace(
         self,
         mapping: (dict[int | str, AgentSet] | list[tuple[int | str, AgentSet]]),

tests/test_agentset.py

@@ -260,6 +260,36 @@ def test_select(self, fix1_AgentSet: ExampleAgentSet):
             selected.active_agents["wealth"].to_list() == agents.df["wealth"].to_list()
         )
 
+    def test_rename(self, fix1_AgentSet: ExampleAgentSet) -> None:
+        agents = fix1_AgentSet
+        reg = agents.model.sets
+        # Inplace rename returns self and updates registry
+        old_name = agents.name
+        result = agents.rename("alpha", inplace=True)
+        assert result is agents
+        assert agents.name == "alpha"
+        assert reg.get("alpha") is agents
+        assert reg.get(old_name) is None
+
+        # Add a second set and claim the same name via registry first
+        other = ExampleAgentSet(agents.model)
+        other["wealth"] = other.starting_wealth
+        other["age"] = [1, 2, 3, 4]
+        reg.add(other)
+        reg.rename(other, "omega")
+        # Now rename the first to an existing name; should canonicalize to omega_1
+        agents.rename("omega", inplace=True)
+        assert agents.name != "omega"
+        assert agents.name.startswith("omega_")
+        assert reg.get(agents.name) is agents
+
+        # Non-inplace: returns a renamed copy of the set
+        copy_set = agents.rename("beta", inplace=False)
+        assert copy_set is not agents
+        assert copy_set.name in ("beta", "beta_1")
+        # Original remains unchanged
+        assert agents.name not in ("beta", "beta_1")
+
         # Test with a pl.Series[bool]
         mask = pl.Series("mask", [True, False, True, True], dtype=pl.Boolean)
         selected = agents.select(mask, inplace=False)

tests/test_agentsetregistry.py

@@ -186,6 +186,42 @@ def test_sort(self, fix_registry_with_two: AgentSetRegistry) -> None:
             reg[0]["wealth"].to_list(), reverse=True
         )
 
+    # Public: rename
+    def test_rename(self, fix_registry_with_two: AgentSetRegistry) -> None:
+        reg = fix_registry_with_two
+        # Single rename by instance, inplace
+        a0 = reg[0]
+        reg.rename(a0, "X")
+        assert a0.name == "X"
+        assert reg.get("X") is a0
+
+        # Rename second to same name should canonicalize
+        a1 = reg[1]
+        reg.rename(a1, "X")
+        assert a1.name != "X" and a1.name.startswith("X_")
+        assert reg.get(a1.name) is a1
+
+        # Non-inplace copy
+        reg2 = reg.rename(a0, "Y", inplace=False)
+        assert reg2 is not reg
+        assert reg.get("Y") is None
+        assert reg2.get("Y") is not None
+
+        # Atomic conflict raise: attempt to rename to existing name
+        with pytest.raises(ValueError):
+            reg.rename({a0: a1.name}, on_conflict="raise", mode="atomic")
+        # Names unchanged
+        assert reg.get(a1.name) is a1
+
+        # Best-effort: one ok, one conflicting → only ok applied
+        unique_name = "Z_unique"
+        reg.rename(
+            {a0: unique_name, a1: unique_name}, on_conflict="raise", mode="best_effort"
+        )
+        assert a0.name == unique_name
+        # a1 stays with its previous (non-unique_name) value
+        assert a1.name != unique_name
+
     # Dunder: __getattr__
     def test__getattr__(self, fix_registry_with_two: AgentSetRegistry) -> None:
         reg = fix_registry_with_two