fix mypy error3

Author: Jungwon-Lee

src/pytorch_lightning/trainer/connectors/data_connector.py

@@ -439,7 +439,7 @@ def _request_dataloader(self, stage: RunningStage) -> Union[DataLoader, List[Dat
         return dataloader
 
     @staticmethod
-    def _resolve_overfit_batches(dataloaders: Collection[DataLoader], mode: RunningStage) -> Collection[DataLoader]:
+    def _resolve_overfit_batches(dataloaders: Union[Collection[DataLoader], DataLoader], mode: RunningStage) -> Collection[DataLoader]:
         all_have_sequential_sampler = True
 
         def resolve_has_no_sequential_sampler(dataloader: DataLoader):

src/pytorch_lightning/trainer/trainer.py

@@ -25,7 +25,7 @@
 from datetime import timedelta
 from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Type, Union
 from weakref import proxy
 
 import torch
@@ -77,7 +77,7 @@
 from pytorch_lightning.trainer.connectors.checkpoint_connector import CheckpointConnector
 from pytorch_lightning.trainer.connectors.data_connector import DataConnector
 from pytorch_lightning.trainer.connectors.logger_connector import LoggerConnector
-from pytorch_lightning.trainer.connectors.logger_connector.result import _ResultCollection
+from pytorch_lightning.trainer.connectors.logger_connector.result import _OUT_DICT, _ResultCollection
 from pytorch_lightning.trainer.connectors.signal_connector import SignalConnector
 from pytorch_lightning.trainer.data_loading import TrainerDataLoadingMixin
 from pytorch_lightning.trainer.optimizers import TrainerOptimizersMixin
@@ -545,6 +545,7 @@ def __init__(
         self._logger_connector.on_trainer_init(logger, log_every_n_steps, move_metrics_to_cpu)
 
         # init debugging flags
+        self.val_check_batch: Union[int, float]
         self.val_check_interval: Union[int, float]
         self.num_sanity_val_steps: Union[int, float]
         self.limit_train_batches: Union[int, float]
@@ -741,7 +742,7 @@ def _fit_impl(
         # TODO: ckpt_path only in v2.0
         ckpt_path = ckpt_path or self.resume_from_checkpoint
         self._ckpt_path = self.__set_ckpt_path(
-            ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
+            ckpt_path, model_provided=True, model_connected=self.lightning_module is not None  # type: ignore
         )
         results = self._run(model, ckpt_path=self.ckpt_path)
 
@@ -985,7 +986,7 @@ def _predict_impl(
         self.state.status = TrainerStatus.RUNNING
         self.predicting = True
 
-        self.predict_loop.return_predictions = return_predictions
+        self.predict_loop.return_predictions = return_predictions # type: ignore
 
         # if a datamodule comes in as the second arg, then fix it for the user
         if isinstance(dataloaders, LightningDataModule):
@@ -1395,7 +1396,7 @@ def __set_ckpt_path(self, ckpt_path: Optional[str], model_provided: bool, model_
 
         if model_provided and ckpt_path is None:
             # use passed model to function without loading weights
-            return
+            return None
 
         if model_connected and ckpt_path is None:
             ckpt_path = "best"
@@ -1449,8 +1450,8 @@ def __set_ckpt_path(self, ckpt_path: Optional[str], model_provided: bool, model_
                     f'.{fn}(ckpt_path="last") is set, but there is no fault tolerant'
                     " or last checkpoint available. No checkpoint will be loaded."
                 )
-                return
-            ckpt_path = max(candidates_ts.keys(), key=partial(operator.getitem, candidates_ts))
+                return None
+            ckpt_path = max(candidates_ts.keys(), key=partial(operator.getitem, candidates_ts)) # type: ignore
 
         if not ckpt_path:
             raise MisconfigurationException(
@@ -1664,7 +1665,7 @@ def _call_callbacks_on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None
             prev_fx_name = pl_module._current_fx_name
             pl_module._current_fx_name = "on_load_checkpoint"
 
-        callback_states: Dict[Union[Type, str], Dict] = checkpoint.get("callbacks")
+        callback_states: Optional[Dict[Union[Type, str], Dict]] = checkpoint.get("callbacks")
 
         if callback_states is None:
             return
@@ -1692,7 +1693,7 @@ def _call_callbacks_on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None
 
     def _call_callbacks_load_state_dict(self, checkpoint: Dict[str, Any]) -> None:
         """Called when loading a model checkpoint, calls every callback's `load_state_dict`."""
-        callback_states: Dict[Union[Type, str], Dict] = checkpoint.get("callbacks")
+        callback_states: Optional[Dict[Union[Type, str], Dict]] = checkpoint.get("callbacks")
 
         if callback_states is None:
             return
@@ -1745,6 +1746,7 @@ def __init_profiler(self, profiler: Optional[Union[Profiler, str]]) -> None:
                 )
             profiler_class = PROFILERS[profiler]
             profiler = profiler_class()
+        assert isinstance(profiler, Profiler)
         self.profiler: Profiler = profiler or PassThroughProfiler()
 
     def __setup_profiler(self) -> None:
@@ -2126,8 +2128,9 @@ def data_parallel_device_ids(self) -> Optional[List[int]]:
         return self.device_ids if isinstance(self.accelerator, CUDAAccelerator) else None
 
     @property
-    def lightning_module(self) -> "pl.LightningModule":
+    def lightning_module(self) -> "pl.LightningModule": # type: ignore
         # TODO: this is actually an optional return
+        assert self.strategy.lightning_module is not None
         return self.strategy.lightning_module
 
     @property
@@ -2219,12 +2222,12 @@ def model(self, model: torch.nn.Module) -> None:
 
     @property
     def log_dir(self) -> Optional[str]:
-        assert self.logger is not None
         if len(self.loggers) == 1:
-            if isinstance(self.logger, TensorBoardLogger):
-                dirpath = self.logger.log_dir
-            else:
+            assert self.logger is not None
+            if not isinstance(self.logger, TensorBoardLogger):
                 dirpath = self.logger.save_dir
+            else:
+                dirpath = self.logger.log_dir
         else:
             dirpath = self.default_root_dir
 
@@ -2709,7 +2712,7 @@ def logger(self, logger: Optional[Logger]) -> None:
         if not logger:
             self.loggers = []
         elif isinstance(logger, LoggerCollection):
-            self.loggers = list(logger)
+            self.loggers = [x for x in logger]
         else:
             self.loggers = [logger]
 
@@ -2722,17 +2725,17 @@ def loggers(self, loggers: Optional[List[Logger]]) -> None:
         self._loggers = loggers if loggers else []
 
     @property
-    def callback_metrics(self) -> Dict[str, Tensor]:
+    def callback_metrics(self) -> Dict:
         # TODO: the true typing return can include dictionaries as defined in
         # `pytorch_lightning.trainer.connectors.logger_connector.result._OUT_DICT`
         return self._logger_connector.callback_metrics
 
     @property
-    def logged_metrics(self) -> dict:
+    def logged_metrics(self) -> _OUT_DICT:
         return self._logger_connector.logged_metrics
 
     @property
-    def progress_bar_metrics(self) -> dict:
+    def progress_bar_metrics(self) -> Dict:
         return self._logger_connector.progress_bar_metrics
 
     @property
@@ -2748,7 +2751,7 @@ def _exit_gracefully_on_signal(self) -> None:
 
     def _should_terminate_gracefully(self) -> bool:
         value = torch.tensor(int(self._terminate_gracefully), device=self.strategy.root_device)
-        return self.strategy.reduce(value, reduce_op="sum") > 0
+        return bool(self.strategy.reduce(value, reduce_op="sum") > 0)
 
     """
     Other

src/pytorch_lightning/utilities/parsing.py

@@ -17,9 +17,8 @@
 import inspect
 import pickle
 import types
-from argparse import Namespace
 from dataclasses import fields, is_dataclass
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
+from typing import Any, Dict, List, MutableMapping, Optional, Sequence, Tuple, Type, Union
 
 from torch import nn
 from typing_extensions import Literal
@@ -94,18 +93,13 @@ def is_picklable(obj: object) -> bool:
         return False
 
 
-def clean_namespace(hparams: Union[Dict[str, Any], Namespace]) -> None:
+def clean_namespace(hparams: MutableMapping) -> None:
     """Removes all unpicklable entries from hparams."""
-
-    hparams_dict = hparams
-    if isinstance(hparams, Namespace):
-        hparams_dict = hparams.__dict__
-
-    del_attrs = [k for k, v in hparams_dict.items() if not is_picklable(v)]
+    del_attrs = [k for k, v in hparams.items() if not is_picklable(v)]
 
     for k in del_attrs:
         rank_zero_warn(f"attribute '{k}' removed from hparams because it cannot be pickled")
-        del hparams_dict[k]
+        del hparams[k]
 
 
 def parse_class_init_keys(