finish typing

Author: otaj

src/lightning_lite/utilities/data.py

@@ -21,7 +21,7 @@
 from typing import Any, Callable, Dict, Generator, Iterable, Optional, Tuple, Type, Union
 
 from lightning_utilities.core.inheritance import get_all_subclasses
-from torch.utils.data import BatchSampler, DataLoader, IterableDataset, Sampler
+from torch.utils.data import BatchSampler, DataLoader, Dataset, IterableDataset, Sampler
 
 from lightning_lite.utilities.enums import LightningEnum
 from lightning_lite.utilities.exceptions import MisconfigurationException
@@ -34,6 +34,7 @@ class _WrapAttrTag(LightningEnum):
     DEL = "del"
 
     def __call__(self, *args: Any) -> None:
+        fn: Union[Callable[[object, str], None], Callable[[object, str, Any], None]]
         if self == self.SET:
             fn = setattr
         else:
@@ -45,20 +46,20 @@ def has_iterable_dataset(dataloader: DataLoader) -> bool:
     return hasattr(dataloader, "dataset") and isinstance(dataloader.dataset, IterableDataset)
 
 
-def has_len(dataloader: Union[DataLoader, Iterable]) -> bool:
+def has_len(dataloader: Union[DataLoader, Iterable, Dataset]) -> bool:
     """Checks if a given Dataloader has ``__len__`` method implemented i.e. if it is a finite dataloader or
     infinite dataloader."""
     try:
         # try getting the length
-        if len(dataloader) == 0:
+        if len(dataloader) == 0:  # type: ignore [arg-type]
             rank_zero_warn(
                 f"`{dataloader.__class__.__name__}` returned 0 length. Please make sure this was your intention."
             )
         has_len = True
     except (TypeError, NotImplementedError):
         has_len = False
 
-    if has_len and has_iterable_dataset(dataloader):
+    if has_len and isinstance(dataloader, DataLoader) and has_iterable_dataset(dataloader):
         rank_zero_warn(
             "Your `IterableDataset` has `__len__` defined."
             " In combination with multi-process data loading (when num_workers > 1),"
@@ -76,7 +77,7 @@ def _update_dataloader(dataloader: DataLoader, sampler: Union[Sampler, Iterable]
 
 def _get_dataloader_init_args_and_kwargs(
     dataloader: DataLoader,
-    sampler: Optional[Sampler],
+    sampler: Union[Sampler, Iterable],
     disallow_batch_sampler: bool = False,
 ) -> Tuple[Tuple[Any], Dict[str, Any]]:
     if not isinstance(dataloader, DataLoader):
@@ -170,7 +171,7 @@ def _get_dataloader_init_args_and_kwargs(
 
 def _dataloader_init_kwargs_resolve_sampler(
     dataloader: DataLoader,
-    sampler: Optional[Sampler],
+    sampler: Union[Sampler, Iterable],
     disallow_batch_sampler: bool = False,
 ) -> Dict[str, Any]:
     """This function is used to handle the sampler, batch_sampler arguments associated within a DataLoader for its

src/pytorch_lightning/strategies/ipu.py

@@ -245,7 +245,7 @@ def _convert_to_poptorch_loader(
             return dataloader
 
         dl_args, dl_kwargs = _get_dataloader_init_args_and_kwargs(
-            dataloader, sampler, mode, self.replication_factor > 1  # type: ignore[arg-type]
+            dataloader, sampler, mode, self.replication_factor > 1
         )
         opts = self.training_opts if mode == RunningStage.TRAINING else self.inference_opts
         dataloader = _reinstantiate_wrapped_cls(

src/pytorch_lightning/utilities/auto_restart.py

@@ -62,7 +62,7 @@ class FastForwardSampler(Sampler):
     samples seen in the last iterations (for the current worker).
     """
 
-    def __init__(self, sampler: Iterator, attr_name: Optional[str] = None) -> None:
+    def __init__(self, sampler: Union[Sampler, Iterable], attr_name: Optional[str] = None) -> None:
         super().__init__(data_source=None)
         self._sampler = sampler
         self.restarting: bool = False

src/pytorch_lightning/utilities/data.py

@@ -30,7 +30,6 @@
 )
 
 import pytorch_lightning as pl
-from lightning_lite.utilities import LightningEnum
 from lightning_lite.utilities.data import _reinstantiate_wrapped_cls, _replace_value_in_saved_args
 from lightning_lite.utilities.data import has_iterable_dataset as new_has_iterable_dataset
 from lightning_lite.utilities.data import has_len as new_has_len
@@ -41,23 +40,12 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_warn
 
-BType = Union[Tensor, str, Mapping[Any, "BType"], Iterable["BType"]]
+# might be supported in later releases, see https://github.com/python/mypy/pull/13297
+BType = Union[Tensor, str, Mapping[Any, "BType"], Iterable["BType"]]  # type: ignore[misc]
 
 warning_cache = WarningCache()
 
 
-class _WrapAttrTag(LightningEnum):
-    SET = "set"
-    DEL = "del"
-
-    def __call__(self, *args: Any) -> None:
-        if self == self.SET:
-            fn = setattr
-        else:
-            fn = delattr
-        return fn(*args)
-
-
 def _extract_batch_size(batch: BType) -> Generator[Optional[int], None, None]:
     if isinstance(batch, Tensor):
         if batch.ndim == 0:
@@ -109,7 +97,7 @@ def extract_batch_size(batch: BType) -> int:
 
 def has_len_all_ranks(
     dataloader: DataLoader,
-    strategy: "pl.Strategy",
+    strategy: "pl.strategies.Strategy",
     model: Union["pl.LightningModule", "pl.LightningDataModule"],
 ) -> bool:
     """Checks if a given Dataloader has ``__len__`` method implemented i.e. if it is a finite dataloader or
@@ -158,7 +146,7 @@ def get_len(dataloader: Union[DataLoader, Dataset]) -> Union[int, float]:
     """
 
     if new_has_len(dataloader):
-        return len(dataloader)
+        return len(dataloader)  # type: ignore [arg-type]
 
     return float("inf")
 
@@ -173,7 +161,7 @@ def _update_dataloader(
 
 def _get_dataloader_init_args_and_kwargs(
     dataloader: DataLoader,
-    sampler: Optional[Sampler],
+    sampler: Union[Sampler, Iterable],
     mode: Optional[RunningStage] = None,
     disallow_batch_sampler: bool = False,
 ) -> Tuple[Tuple[Any], Dict[str, Any]]:
@@ -273,7 +261,7 @@ def _get_dataloader_init_args_and_kwargs(
 
 def _dataloader_init_kwargs_resolve_sampler(
     dataloader: DataLoader,
-    sampler: Optional[Sampler],
+    sampler: Union[Sampler, Iterable],
     mode: Optional[RunningStage] = None,
     disallow_batch_sampler: bool = False,
 ) -> Dict[str, Any]: