implement list methods.

Author: GdoongMathew

src/lightning/pytorch/loggers/utilities.py

@@ -15,14 +15,17 @@
 
 from collections.abc import ItemsView, Iterable, KeysView, Mapping, ValuesView
 from pathlib import Path
-from typing import Any, Optional, SupportsIndex, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, SupportsIndex, TypeVar, Union
 
 from torch import Tensor
 from typing_extensions import Self
 
 import lightning.pytorch as pl
 from lightning.pytorch.callbacks import Checkpoint
 
+if TYPE_CHECKING:
+    from _typeshed import SupportsRichComparison
+
 
 def _version(loggers: list[Any], separator: str = "_") -> Union[int, str]:
     if len(loggers) == 1:
@@ -122,28 +125,96 @@ def __init__(self, loggers: Union[Iterable[_T], Mapping[str, _T]] = None):
             self._dict: dict = {}
 
     def __eq__(self, other: Any) -> bool:
-        self_list = list(self)
+        list_eq = list.__eq__(self, other)
         if isinstance(other, _ListMap):
-            return self_list == list(other) and self._dict == other._dict
-        if isinstance(other, list):
-            return self_list == other
-        return False
+            dict_eq = self._dict == other._dict
+            return list_eq and dict_eq
+        return list_eq
+
+    def copy(self):
+        new_listmap = _ListMap(self)
+        new_listmap._dict = self._dict.copy()
+        return new_listmap
+
+    def extend(self, __iterable: Iterable[_T]) -> None:
+        if isinstance(__iterable, _ListMap):
+            offset = len(self)
+            for key, idx in __iterable._dict.items():
+                self._dict[key] = idx + offset
+        super().extend(__iterable)
+
+    def pop(self, key: Union[SupportsIndex, str] = -1, default: Optional[Any] = None) -> _T:
+        if isinstance(key, int):
+            ret = list.pop(self, key)
+            for str_key, idx in list(self._dict.items()):
+                if idx == key:
+                    self._dict.pop(str_key)
+                elif idx > key:
+                    self._dict[str_key] = idx - 1
+            return ret
+        if isinstance(key, str):
+            if key not in self._dict:
+                return default
+            return self.pop(self._dict[key])
+        raise TypeError("Key must be int or str")
+
+    def insert(self, index: SupportsIndex, __object: _T) -> None:
+        for key, idx in self._dict.items():
+            if idx >= index:
+                self._dict[key] = idx + 1
+        list.insert(self, index, __object)
+
+    def remove(self, __object: _T) -> None:
+        idx = self.index(__object)
+        name = None
+        for key, val in self._dict.items():
+            if val == idx:
+                name = key
+            elif val > idx:
+                self._dict[key] = val - 1
+        if name:
+            self._dict.pop(name, None)
+        list.remove(self, __object)
+
+    def sort(
+        self,
+        *,
+        key: Optional[Callable[[_T], "SupportsRichComparison"]] = None,
+        reverse: bool = False,
+    ) -> None:
+        # Create a mapping from item to its name(s)
+        item_to_names = {}
+        for name, idx in self._dict.items():
+            item = self[idx]
+            item_to_names.setdefault(item, []).append(name)
+        # Sort the list
+        list.sort(self, key=key, reverse=reverse)
+        # Update _dict with new indices
+        new_dict = {}
+        for idx, item in enumerate(self):
+            if item in item_to_names:
+                for name in item_to_names[item]:
+                    new_dict[name] = idx
+        self._dict = new_dict
 
     # --- List-like interface ---
     def __getitem__(self, key: Union[int, slice, str]) -> _T:
-        if isinstance(key, (int, slice)):
-            return list.__getitem__(self, key)
         if isinstance(key, str):
-            return list.__getitem__(self, self._dict[key])
-        raise TypeError("Key must be int / slice (for index) or str (for name).")
+            return self[self._dict[key]]
+        return list.__getitem__(self, key)
 
-    def __add__(self, other: Union[list[_T], Self]) -> list[_T]:
-        # todo
-        return list.__add__(self, other)
+    def __add__(self, other: Union[list[_T], Self]) -> Self:
+        new_listmap = self.copy()
+        new_listmap += other
+        return new_listmap
 
     def __iadd__(self, other: Union[list[_T], Self]) -> Self:
-        # todo
-        return list.__iadd__(self, other)
+        if isinstance(other, _ListMap):
+            offset = len(self)
+            for key, idx in other._dict.items():
+                self._dict[key] = idx + offset
+
+        return super().__iadd__(other)
 
     def __setitem__(self, key: Union[SupportsIndex, slice, str], value: _T) -> None:
         if isinstance(key, (int, slice)):
@@ -192,20 +263,6 @@ def items(self) -> ItemsView[str, _T]:
         return d.items()
 
     # --- List and Dict interface ---
-    def pop(self, key: Union[SupportsIndex, str] = -1, default: Optional[Any] = None) -> _T:
-        if isinstance(key, int):
-            ret = list.pop(self, key)
-            for str_key, idx in list(self._dict.items()):
-                if idx == key:
-                    self._dict.pop(str_key)
-                elif idx > key:
-                    self._dict[str_key] = idx - 1
-            return ret
-        if isinstance(key, str):
-            if key not in self._dict:
-                return default
-            return self.pop(self._dict[key])
-        raise TypeError("Key must be int or str")
 
     def __repr__(self) -> str:
         ret = super().__repr__()

tests/tests_pytorch/loggers/test_utilities.py

@@ -98,20 +98,24 @@ def test_listmap_setitem():
 def test_listmap_add():
     """Test adding two collections together."""
     lm1 = _ListMap([1, 2])
-    lm2 = _ListMap([3, 4])
+    lm2 = _ListMap({"3": 3, "5": 5})
     combined = lm1 + lm2
-    assert isinstance(combined, list)
+    assert isinstance(combined, _ListMap)
     assert len(combined) == 4
-    assert combined[0] == 1
-    assert combined[1] == 2
-    assert combined[2] == 3
-    assert combined[3] == 4
+    assert combined is not lm1
+    assert combined == [1, 2, 3, 5]
+    assert combined["3"] == 3
+    assert combined["5"] == 5
 
-    combined += lm1
-    assert isinstance(combined, list)
-    assert len(combined) == 6
-    for item, expected in zip(combined, [1, 2, 3, 4, 1, 2]):
-        assert item == expected
+    ori_lm1_id = id(lm1)
+
+    lm1 += lm2
+    assert ori_lm1_id == id(lm1)
+    assert isinstance(lm1, _ListMap)
+    assert len(lm1) == 4
+    assert lm1 == [1, 2, 3, 5]
+    assert lm1["3"] == 3
+    assert lm1["5"] == 5
 
 
 def test_listmap_remove():
@@ -190,6 +194,13 @@ def test_listmap_dict_pop():
     assert "b" not in lm
     assert len(lm) == 2
 
+    value = lm.pop(0)
+    assert value == 1
+    assert lm["c"] == 3  # still accessible by key
+    assert len(lm) == 1
+    with pytest.raises(KeyError):
+        lm["a"]  # "a" was removed
+
 
 def test_listmap_dict_setitem():
     lm = _ListMap({
@@ -201,3 +212,23 @@ def test_listmap_dict_setitem():
     lm["c"] = 3
     assert lm["c"] == 3
     assert len(lm) == 3
+
+
+def test_listmap_sort():
+    lm = _ListMap({"b": 1, "c": 3, "a": 2, "z": -7})
+
+    lm.extend([-1, -2, 5])
+    lm.sort(key=lambda x: abs(x))
+    assert lm == [1, -1, 2, -2, 3, 5, -7]
+    assert lm["a"] == 2
+    assert lm["b"] == 1
+    assert lm["c"] == 3
+    assert lm["z"] == -7
+
+    lm = _ListMap({"b": 1, "c": 3, "a": 2, "z": -7})
+    lm.sort(reverse=True)
+    assert lm == [3, 2, 1, -7]
+    assert lm["a"] == 2
+    assert lm["b"] == 1
+    assert lm["c"] == 3
+    assert lm["z"] == -7