refactor: remove methods

Author: GdoongMathew

src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py

@@ -11,12 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from collections.abc import ItemsView, Iterable, KeysView, Mapping, ValuesView
-from typing import TYPE_CHECKING, Any, Callable, Optional, SupportsIndex, TypeVar, Union
+from collections.abc import Iterable, KeysView, Mapping
+from typing import TYPE_CHECKING, Any, Optional, SupportsIndex, TypeVar, Union
 
 from lightning_utilities.core.apply_func import apply_to_collection
 from torch import Tensor
-from typing_extensions import Self, overload
+from typing_extensions import overload
 
 import lightning.pytorch as pl
 from lightning.fabric.loggers.tensorboard import _TENSORBOARD_AVAILABLE, _TENSORBOARDX_AVAILABLE
@@ -29,9 +29,8 @@
 
 warning_cache = WarningCache()
 
-
 if TYPE_CHECKING:
-    from _typeshed import SupportsRichComparison
+    pass
 
 
 class _LoggerConnector:
@@ -309,207 +308,73 @@ class _ListMap(list[_T]):
 
     """
 
-    _dict: dict[str, int]
+    _dict_map: dict[str, int]
 
-    def __init__(self, __iterable: Optional[Union[Mapping[str, _T], Iterable[_T]]] = None):
+    def __init__(self, __iterable: Optional[Union[dict[str, _T], Iterable[_T]]] = None):
         if isinstance(__iterable, Mapping):
             # super inits list with values
             if any(not isinstance(x, str) for x in __iterable):
                 raise TypeError("When providing a Mapping, all keys must be of type str.")
             super().__init__(__iterable.values())
-            _dict = dict(zip(__iterable.keys(), range(len(__iterable))))
+            self._dict_map = {key: idx for idx, key in enumerate(__iterable)}
         else:
-            default_dict: dict[str, int] = {}
-            if isinstance(__iterable, _ListMap):
-                default_dict = __iterable._dict.copy()
-            super().__init__(() if __iterable is None else __iterable)
-            _dict = default_dict
-        self._dict = _dict
+            super().__init__(__iterable or ())
+            self._dict_map = {}
 
     def __eq__(self, other: Any) -> bool:
         list_eq = super().__eq__(other)
         if isinstance(other, _ListMap):
-            return list_eq and self._dict == other._dict
+            list_eq &= other._dict_map == self._dict_map
         return list_eq
 
-    def copy(self) -> "_ListMap":
-        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, index: SupportsIndex = -1, /):
+        if self._dict_map:
+            index_int = index.__index__()
+            if index_int < 0:
+                index_int += len(self)
+
+            for key, idx in list(self._dict_map.items()):
+                if idx == index_int:
+                    self._dict_map.pop(key)
+                elif idx > index_int:
+                    self._dict_map[key] -= 1
+        return super().pop(index)
+
+    def insert(self, index: SupportsIndex, __object: _T, /) -> None:
+        if self._dict_map:
+            index_int = index.__index__()
+            if index_int < 0:
+                index_int += len(self)
+
+            for key, idx in list(self._dict_map.items()):
+                if idx >= index_int:
+                    self._dict_map[key] += 1
 
-    @overload
-    def pop(self, key: SupportsIndex = -1, /) -> _T: ...
-
-    @overload
-    def pop(self, key: Union[str, SupportsIndex], default: _T, /) -> _T: ...
-
-    @overload
-    def pop(self, key: str, default: _PT, /) -> Union[_T, _PT]: ...
-
-    def pop(self, key: Union[SupportsIndex, str] = -1, default: Any = None) -> _T:
-        if isinstance(key, int):
-            ret = super().pop(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:
-        idx_int = int(index)
-        # Check for negative indices
-        if idx_int < 0:
-            idx_int += len(self)
-        for key, idx in self._dict.items():
-            if idx >= idx_int:
-                self._dict[key] = idx + 1
         return super().insert(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)
-        super().remove(__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: dict[_T, list[str]] = {}
-        for name, idx in self._dict.items():
-            item = self[idx]
-            item_to_names.setdefault(item, []).append(name)
-        # Sort the list
-        super().sort(key=key, reverse=reverse)
-        # Update _dict with new indices
-        new_dict: dict[str, int] = {}
-        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
-
     @overload
     def __getitem__(self, key: Union[SupportsIndex, str], /) -> _T: ...
 
     @overload
     def __getitem__(self, key: slice, /) -> list[_T]: ...
 
     def __getitem__(self, key: Union[SupportsIndex, str, slice], /) -> Union[_T, list[_T]]:
-        if isinstance(key, str):
-            return self[self._dict[key]]
+        if self._dict_map and isinstance(key, str):
+            return self[self._dict_map[key]]
         return super().__getitem__(key)
 
-    def __add__(self, other: Union[list[_T], "_ListMap[_T]"]) -> "_ListMap[_T]":  # type: ignore[override]
-        new_listmap = self.copy()
-        new_listmap += other
-        return new_listmap
-
-    def __iadd__(self, other: Iterable[_T]) -> Self:  # type: ignore[override]
-        if isinstance(other, _ListMap):
-            offset = len(self)
-            for key, idx in other._dict.items():
-                # notes: if there are duplicate keys, the ones from other will overwrite self
-                self._dict[key] = idx + offset
-
-        return super().__iadd__(other)
-
-    @overload
-    def __setitem__(self, key: Union[SupportsIndex, str], value: _T, /) -> None: ...
-
-    @overload
-    def __setitem__(self, key: slice, value: Iterable[_T], /) -> None: ...
-
-    def __setitem__(self, key: Union[SupportsIndex, str, slice], value: Any, /) -> None:
-        if isinstance(key, str):
-            # replace or insert by name
-            if key in self._dict:
-                self[self._dict[key]] = value
-            else:
-                self.append(value)
-                self._dict[key] = len(self) - 1
-            return None
-        return super().__setitem__(key, value)
-
     def __contains__(self, item: Union[object, str]) -> bool:
         if isinstance(item, str):
-            return item in self._dict
+            return item in self._dict_map
         return super().__contains__(item)
 
-    # --- Dict-like interface ---
-
-    def __delitem__(self, key: Union[SupportsIndex, slice, str]) -> None:
-        index: Union[SupportsIndex, slice]
-        if isinstance(key, str):
-            if key not in self._dict:
-                raise KeyError(f"Key '{key}' not found.")
-            index = self._dict[key]
-        else:
-            index = key
-
-        if isinstance(index, (int, slice)):
-            super().__delitem__(index)
-            for _key in index.indices(len(self)) if isinstance(index, slice) else [index]:
-                # update indices in the dict
-                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
-        else:
-            raise TypeError("Key must be int or str")
-
-    def keys(self) -> KeysView[str]:
-        return self._dict.keys()
-
-    def values(self) -> ValuesView[_T]:
-        return {k: self[v] for k, v in self._dict.items()}.values()
-
-    def items(self) -> ItemsView[str, _T]:
-        return {k: self[v] for k, v in self._dict.items()}.items()
-
-    @overload
-    def get(self, __key: str) -> Optional[_T]: ...
-
-    @overload
-    def get(self, __key: str, default: _PT) -> Union[_T, _PT]: ...
-
-    def get(self, __key: str, default: Optional[_PT] = None) -> Optional[Union[_T, _PT]]:
-        if __key in self._dict:
-            return self[self._dict[__key]]
-        return default
-
     def __repr__(self) -> str:
         ret = super().__repr__()
-        return f"{type(self).__name__}({ret}, keys={list(self._dict.keys())})"
-
-    def reverse(self) -> None:
-        for key, idx in self._dict.items():
-            self._dict[key] = len(self) - 1 - idx
-        return super().reverse()
+        return f"{type(self).__name__}({ret}, keys={list(self._dict_map.keys())})"
 
     def clear(self) -> None:
-        self._dict.clear()
+        self._dict_map.clear()
         return super().clear()
+
+    def keys(self) -> KeysView[str]:
+        return self._dict_map.keys()