Merge remote-tracking branch 'origin/feat/logger_dict' into feat/logger_dict

Author: GdoongMathew

src/lightning/pytorch/loggers/utilities.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 """Utilities for loggers."""
 
-from collections.abc import Mapping
+from collections.abc import ItemsView, Iterable, KeysView, Mapping, ValuesView
 from pathlib import Path
-from typing import Any, Optional, Self, TypeVar, Union
+from typing import Any, Optional, SupportsIndex, TypeVar, Union
 
 from torch import Tensor
+from typing_extensions import Self
 
 import lightning.pytorch as pl
 from lightning.pytorch.callbacks import Checkpoint
@@ -109,7 +110,7 @@ def _log_hyperparams(trainer: "pl.Trainer") -> None:
 class _ListMap(list[_T]):
     """A hybrid container for loggers allowing both index and name access."""
 
-    def __init__(self, loggers: Union[list[_T], Mapping[str, _T]] = None):
+    def __init__(self, loggers: Union[Iterable[_T], Mapping[str, _T]] = None):
         if isinstance(loggers, Mapping):
             # super inits list with values
             if any(not isinstance(x, str) for x in loggers):
@@ -120,7 +121,7 @@ def __init__(self, loggers: Union[list[_T], Mapping[str, _T]] = None):
             super().__init__(() if loggers is None else loggers)
             self._dict: dict = {}
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         self_list = list(self)
         if isinstance(other, _ListMap):
             return self_list == list(other) and self._dict == other._dict
@@ -144,7 +145,7 @@ def __iadd__(self, other: Union[list[_T], Self]) -> Self:
         # todo
         return list.__iadd__(self, other)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: Union[SupportsIndex, slice, str], value: _T) -> None:
         if isinstance(key, (int, slice)):
             # replace element by index
             return list.__setitem__(self, key, value)
@@ -158,14 +159,14 @@ def __setitem__(self, key, value):
             return None
         raise TypeError("Key must be int or str")
 
-    def __contains__(self, item):
+    def __contains__(self, item: Union[_T, str]) -> bool:
         if isinstance(item, str):
             return item in self._dict
         return list.__contains__(self, item)
 
     # --- Dict-like interface ---
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: Union[int, slice, str]) -> None:
         if isinstance(key, (int, slice)):
             loggers = list.__getitem__(self, key)
             super(list, self).__delitem__(key)
@@ -179,19 +180,19 @@ def __delitem__(self, key):
         else:
             raise TypeError("Key must be int or str")
 
-    def keys(self):
+    def keys(self) -> KeysView[str]:
         return self._dict.keys()
 
-    def values(self):
+    def values(self) -> ValuesView[_T]:
         d = {k: self[v] for k, v in self._dict.items()}
         return d.values()
 
-    def items(self):
+    def items(self) -> ItemsView[str, _T]:
         d = {k: self[v] for k, v in self._dict.items()}
         return d.items()
 
     # --- List and Dict interface ---
-    def pop(self, key: Union[int, str] = -1, default: Optional[Any] = None) -> _T:
+    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()):
@@ -206,6 +207,18 @@ def pop(self, key: Union[int, str] = -1, default: Optional[Any] = None) -> _T:
             return self.pop(self._dict[key])
         raise TypeError("Key must be int or str")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         ret = super().__repr__()
         return f"_ListMap({ret}, keys={list(self._dict.keys())})"
+
+    def __reversed__(self) -> Iterable[_T]:
+        return reversed(list(self))
+
+    def reverse(self) -> None:
+        for key, idx in self._dict.items():
+            self._dict[key] = len(self) - 1 - idx
+        list.reverse(self)
+
+    def clear(self):
+        self._dict.clear()
+        list.clear(self)

src/lightning/pytorch/trainer/trainer.py

@@ -494,7 +494,7 @@ def __init__(
         setup._init_profiler(self, profiler)
 
         # init logger flags
-        self._loggers: list[Logger]
+        self._loggers: _ListMap[Logger]
         self._logger_connector.on_trainer_init(logger, log_every_n_steps)
 
         # init debugging flags

tests/tests_pytorch/loggers/test_utilities.py

@@ -122,6 +122,30 @@ def test_listmap_remove():
     assert 2 not in lm
 
 
+def test_listmap_reverse():
+    """Test reversing the collection."""
+    lm = _ListMap({"1": 1, "2": 2, "3": 3})
+    lm.reverse()
+    assert lm == [3, 2, 1]
+    for (key, value), expected in zip(lm.items(), [("1", 1), ("2", 2), ("3", 3)]):
+        assert (key, value) == expected
+
+
+def test_listmap_reversed():
+    """Test reversed iterator of the collection."""
+    lm = _ListMap({"1": 1, "2": 2, "3": 3})
+    rev_lm = list(reversed(lm))
+    assert rev_lm == [3, 2, 1]
+
+
+def test_listmap_clear():
+    """Test clearing the collection."""
+    lm = _ListMap({"1": 1, "2": 2, "3": 3})
+    lm.clear()
+    assert len(lm) == 0
+    assert len(lm.keys()) == 0
+
+
 # Dict type properties tests
 def test_listmap_keys():
     lm = _ListMap({

tests/tests_pytorch/trainer/properties/test_loggers.py

@@ -23,24 +23,21 @@ def test_trainer_loggers_property():
     """Test for correct initialization of loggers in Trainer."""
     logger1 = CustomLogger()
     logger2 = CustomLogger()
-    logger3 = CustomLogger()
+    CustomLogger()
 
     # trainer.loggers should be a copy of the input list
     trainer = Trainer(logger=[logger1, logger2])
 
     assert trainer.loggers == [logger1, logger2]
-    assert trainer.logger_map == {}
 
     # trainer.loggers should create a list of size 1
     trainer = Trainer(logger=logger1)
 
     assert trainer.logger == logger1
     assert trainer.loggers == [logger1]
-    assert trainer.logger_map == {}
 
     trainer.loggers.append(logger2)
     assert trainer.loggers == [logger1, logger2]
-    assert trainer.logger_map == {}
 
     # trainer.loggers should be a list of size 1 holding the default logger
     trainer = Trainer(logger=True)
@@ -51,19 +48,15 @@ def test_trainer_loggers_property():
     trainer = Trainer(logger={"log1": logger1, "log2": logger2})
     assert trainer.logger == logger1
     assert trainer.loggers == [logger1, logger2]
-    assert isinstance(trainer.logger_map, dict)
-    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
-
-    trainer.loggers.append(logger3)
-    assert trainer.loggers == [logger1, logger2, logger3]
-    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
+    assert trainer.loggers["log1"] is logger1
+    assert trainer.loggers["log2"] is logger2
 
 
 def test_trainer_loggers_setters():
     """Test the behavior of setters for trainer.logger and trainer.loggers."""
     logger1 = CustomLogger()
     logger2 = CustomLogger()
-    logger3 = CustomLogger()
+    CustomLogger()
 
     trainer = Trainer()
     assert type(trainer.logger) is TensorBoardLogger
@@ -78,46 +71,45 @@ def test_trainer_loggers_setters():
     assert trainer.logger is None
     assert trainer.loggers == []
     assert isinstance(trainer.loggers, list)
-    assert trainer.logger_map == {}
 
     # Test setters for trainer.loggers
     trainer.loggers = [logger1, logger2]
     assert trainer.loggers == [logger1, logger2]
     assert isinstance(trainer.loggers, list)
-    assert trainer.logger_map == {}
 
     trainer.loggers = [logger1]
     assert trainer.loggers == [logger1]
     assert trainer.logger == logger1
-    assert trainer.logger_map == {}
 
     trainer.loggers = []
     assert trainer.loggers == []
     assert trainer.logger is None
     assert isinstance(trainer.loggers, list)
-    assert trainer.logger_map == {}
 
     trainer.loggers = None
     assert trainer.loggers == []
     assert trainer.logger is None
     assert isinstance(trainer.loggers, list)
-    assert trainer.logger_map == {}
+
+    trainer.loggers = {}
+    assert trainer.loggers == []
+    assert trainer.logger is None
+    assert isinstance(trainer.loggers, list)
 
     trainer.loggers = {"log1": logger1, "log2": logger2}
     assert trainer.loggers == [logger1, logger2]
     assert isinstance(trainer.loggers, list)
-    assert isinstance(trainer.logger_map, dict)
-    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
 
-    trainer.loggers.append(logger3)
-    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
+    assert trainer.loggers["log1"] is logger1
+    assert trainer.loggers["log2"] is logger2
 
 
 @pytest.mark.parametrize(
     "logger_value",
     [
         False,
         [],
+        {},
     ],
 )
 def test_no_logger(tmp_path, logger_value):
@@ -130,4 +122,3 @@ def test_no_logger(tmp_path, logger_value):
     assert trainer.logger is None
     assert trainer.loggers == []
     assert trainer.log_dir == str(tmp_path)
-    assert trainer.logger_map == {}