Complete partially typed signatures (monai.utils) (#5891)

Part of #5884.

### Description

Fully type annotate any functions with at least one type annotation in
module `visualize`

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.

Signed-off-by: Felix Schnabel <f.schnabel@tum.de>

Author: Shadow-Devil

monai/apps/auto3dseg/ensemble_builder.py

@@ -15,10 +15,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from copy import deepcopy
-from typing import Any
+from typing import Any, cast
 from warnings import warn
 
 import numpy as np
+import torch
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
 from monai.apps.utils import get_logger
@@ -108,7 +109,7 @@ def ensemble_pred(self, preds, sigmoid=False):
 
         if self.mode == "mean":
             prob = MeanEnsemble()(preds)
-            return prob2class(prob, dim=0, keepdim=True, sigmoid=sigmoid)
+            return prob2class(cast(torch.Tensor, prob), dim=0, keepdim=True, sigmoid=sigmoid)
         elif self.mode == "vote":
             classes = [prob2class(p, dim=0, keepdim=True, sigmoid=sigmoid) for p in preds]
             if sigmoid:

monai/data/dataset.py

@@ -25,7 +25,7 @@
 from multiprocessing.managers import ListProxy
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any
+from typing import IO, TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
@@ -1403,7 +1403,7 @@ def __init__(
         dat = np.load(npzfile)
 
         self.arrays = {storedk: dat[datak] for datak, storedk in self.keys.items()}
-        self.length = self.arrays[first(self.keys.values())].shape[0]
+        self.length = self.arrays[cast(str, first(self.keys.values()))].shape[0]
 
         self.other_keys = {} if other_keys is None else {k: dat[k] for k in other_keys}
 

monai/utils/decorators.py

@@ -15,17 +15,19 @@
 
 __all__ = ["RestartGenerator", "MethodReplacer"]
 
+from typing import Callable, Generator
+
 
 class RestartGenerator:
     """
     Wraps a generator callable which will be called whenever this class is iterated and its result returned. This is
     used to create an iterator which can start iteration over the given generator multiple times.
     """
 
-    def __init__(self, create_gen) -> None:
+    def __init__(self, create_gen: Callable[[], Generator]) -> None:
         self.create_gen = create_gen
 
-    def __iter__(self):
+    def __iter__(self) -> Generator:
         return self.create_gen()
 
 
@@ -36,7 +38,7 @@ class MethodReplacer:
 
     replace_list_name = "__replacemethods__"
 
-    def __init__(self, meth) -> None:
+    def __init__(self, meth: Callable) -> None:
         self.meth = meth
 
     def replace_method(self, meth):

monai/utils/deprecate_utils.py

@@ -43,7 +43,7 @@ def deprecated(
     removed: str | None = None,
     msg_suffix: str = "",
     version_val: str = __version__,
-    warning_category=FutureWarning,
+    warning_category: type[FutureWarning] = FutureWarning,
 ) -> Callable[[T], T]:
     """
     Marks a function or class as deprecated. If `since` is given this should be a version at or earlier than the
@@ -121,13 +121,13 @@ def _wrapper(*args, **kwargs):
 
 
 def deprecated_arg(
-    name,
+    name: str,
     since: str | None = None,
     removed: str | None = None,
     msg_suffix: str = "",
     version_val: str = __version__,
     new_name: str | None = None,
-    warning_category=FutureWarning,
+    warning_category: type[FutureWarning] = FutureWarning,
 ) -> Callable[[T], T]:
     """
     Marks a particular named argument of a callable as deprecated. The same conditions for `since` and `removed` as
@@ -235,7 +235,7 @@ def deprecated_arg_default(
     replaced: str | None = None,
     msg_suffix: str = "",
     version_val: str = __version__,
-    warning_category=FutureWarning,
+    warning_category: type[FutureWarning] = FutureWarning,
 ) -> Callable[[T], T]:
     """
     Marks a particular arguments default of a callable as deprecated. It is changed from `old_default` to `new_default`

monai/utils/dist.py

@@ -11,6 +11,13 @@
 
 from __future__ import annotations
 
+import sys
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+
+from typing import overload
+
 import torch
 import torch.distributed as dist
 
@@ -39,7 +46,22 @@ def get_dist_device():
     return None
 
 
-def evenly_divisible_all_gather(data: torch.Tensor, concat: bool = True):
+@overload
+def evenly_divisible_all_gather(data: torch.Tensor, concat: Literal[True]) -> torch.Tensor:
+    ...
+
+
+@overload
+def evenly_divisible_all_gather(data: torch.Tensor, concat: Literal[False]) -> list[torch.Tensor]:
+    ...
+
+
+@overload
+def evenly_divisible_all_gather(data: torch.Tensor, concat: bool) -> torch.Tensor | list[torch.Tensor]:
+    ...
+
+
+def evenly_divisible_all_gather(data: torch.Tensor, concat: bool = True) -> torch.Tensor | list[torch.Tensor]:
     """
     Utility function for distributed data parallel to pad at first dim to make it evenly divisible and all_gather.
     The input data of every rank should have the same number of dimensions, only the first dim can be different.
@@ -149,6 +171,6 @@ def string_list_all_gather(strings: list[str], delimiter: str = "\t") -> list[st
 
     joined = delimiter.join(strings)
     gathered = evenly_divisible_all_gather(torch.tensor(bytearray(joined, "utf-8"), dtype=torch.long), concat=False)
-    gathered = [bytearray(g.tolist()).decode("utf-8").split(delimiter) for g in gathered]
+    _gathered = [bytearray(g.tolist()).decode("utf-8").split(delimiter) for g in gathered]
 
-    return [i for k in gathered for i in k]
+    return [i for k in _gathered for i in k]

monai/utils/enums.py

@@ -369,19 +369,19 @@ class PostFix(StrEnum):
     """Post-fixes."""
 
     @staticmethod
-    def _get_str(prefix, suffix):
+    def _get_str(prefix: str | None, suffix: str) -> str:
         return suffix if prefix is None else f"{prefix}_{suffix}"
 
     @staticmethod
-    def meta(key: str | None = None):
+    def meta(key: str | None = None) -> str:
         return PostFix._get_str(key, "meta_dict")
 
     @staticmethod
-    def orig_meta(key: str | None = None):
+    def orig_meta(key: str | None = None) -> str:
         return PostFix._get_str(key, "orig_meta_dict")
 
     @staticmethod
-    def transforms(key: str | None = None):
+    def transforms(key: str | None = None) -> str:
         return PostFix._get_str(key, TraceKeys.KEY_SUFFIX[1:])
 
 

monai/utils/jupyter_utils.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from enum import Enum
 from threading import RLock, Thread
 from typing import TYPE_CHECKING, Any
@@ -44,13 +44,13 @@
 
 
 def plot_metric_graph(
-    ax,
+    ax: plt.Axes,
     title: str,
-    graphmap: dict[str, list[float] | tuple[list[float], list[float]]],
+    graphmap: Mapping[str, list[float] | tuple[list[float], list[float]]],
     yscale: str = "log",
     avg_keys: tuple[str] = (LOSS_NAME,),
     window_fraction: int = 20,
-):
+) -> None:
     """
     Plot metrics on a single graph with running averages plotted for selected keys. The values in `graphmap`
     should be lists of (timepoint, value) pairs as stored in MetricLogger objects.
@@ -91,9 +91,9 @@ def plot_metric_graph(
 
 
 def plot_metric_images(
-    fig,
+    fig: plt.Figure,
     title: str,
-    graphmap: dict[str, list[float] | tuple[list[float], list[float]]],
+    graphmap: Mapping[str, list[float] | tuple[list[float], list[float]]],
     imagemap: dict[str, np.ndarray],
     yscale: str = "log",
     avg_keys: tuple[str] = (LOSS_NAME,),
@@ -138,7 +138,7 @@ def plot_metric_images(
     return axes
 
 
-def tensor_to_images(name: str, tensor: torch.Tensor):
+def tensor_to_images(name: str, tensor: torch.Tensor) -> np.ndarray | None:
     """
     Return an tuple of images derived from the given tensor. The `name` value indices which key from the
     output or batch value the tensor was stored as, or is "Batch" or "Output" if these were single tensors
@@ -147,25 +147,25 @@ def tensor_to_images(name: str, tensor: torch.Tensor):
     each channel separately.
     """
     if tensor.ndim == 3 and tensor.shape[1] > 2 and tensor.shape[2] > 2:
-        return tensor.cpu().data.numpy()
+        return tensor.cpu().data.numpy()  # type: ignore[no-any-return]
     if tensor.ndim == 4 and tensor.shape[2] > 2 and tensor.shape[3] > 2:
         dmid = tensor.shape[1] // 2
-        return tensor[:, dmid].cpu().data.numpy()
+        return tensor[:, dmid].cpu().data.numpy()  # type: ignore[no-any-return]
 
     return None
 
 
 def plot_engine_status(
     engine: Engine,
-    logger,
+    logger: Any,
     title: str = "Training Log",
     yscale: str = "log",
     avg_keys: tuple[str] = (LOSS_NAME,),
     window_fraction: int = 20,
-    image_fn: Callable | None = tensor_to_images,
-    fig=None,
+    image_fn: Callable[[str, torch.Tensor], Any] | None = tensor_to_images,
+    fig: plt.Figure = None,
     selected_inst: int = 0,
-) -> tuple:
+) -> tuple[plt.Figure, list]:
     """
     Plot the status of the given Engine with its logger. The plot will consist of a graph of loss values and metrics
     taken from the logger, and images taken from the `output` and `batch` members of `engine.state`. The images are
@@ -191,7 +191,7 @@ def plot_engine_status(
     else:
         fig = plt.Figure(figsize=(20, 10), tight_layout=True, facecolor="white")
 
-    graphmap = {LOSS_NAME: logger.loss}
+    graphmap: dict[str, list[float]] = {LOSS_NAME: logger.loss}
     graphmap.update(logger.metrics)
 
     imagemap: dict = {}
@@ -233,10 +233,12 @@ def plot_engine_status(
     return fig, axes
 
 
-def _get_loss_from_output(output: dict[str, torch.Tensor] | torch.Tensor):
+def _get_loss_from_output(
+    output: list[torch.Tensor | dict[str, torch.Tensor]] | dict[str, torch.Tensor] | torch.Tensor
+) -> torch.Tensor:
     """Returns a single value from the network output, which is a dict or tensor."""
 
-    def _get_loss(data):
+    def _get_loss(data: torch.Tensor | dict[str, torch.Tensor]) -> torch.Tensor:
         if isinstance(data, dict):
             return data["loss"]
         return data
@@ -286,7 +288,7 @@ def __init__(
         self._status_dict: dict[str, Any] = {}
         self.loss_transform = loss_transform
         self.metric_transform = metric_transform
-        self.fig = None
+        self.fig: plt.Figure | None = None
         self.status_format = status_format
 
         self.engine.add_event_handler(Events.ITERATION_COMPLETED, self._update_status)
@@ -357,7 +359,7 @@ def status(self) -> str:
 
         return ", ".join(msgs)
 
-    def plot_status(self, logger, plot_func: Callable = plot_engine_status):
+    def plot_status(self, logger: Any, plot_func: Callable = plot_engine_status) -> plt.Figure:
         """
         Generate a plot of the current status of the contained engine whose loss and metrics were tracked by `logger`.
         The function `plot_func` must accept arguments `title`, `engine`, `logger`, and `fig` which are the plot title,

monai/utils/misc.py

@@ -23,7 +23,7 @@
 from collections.abc import Callable, Iterable, Sequence
 from distutils.util import strtobool
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, TypeVar, cast, overload
 
 import numpy as np
 import torch
@@ -82,7 +82,20 @@ def star_zip_with(op, *vals):
     return zip_with(op, *vals, mapfunc=itertools.starmap)
 
 
-def first(iterable, default=None):
+T = TypeVar("T")
+
+
+@overload
+def first(iterable: Iterable[T], default: T) -> T:
+    ...
+
+
+@overload
+def first(iterable: Iterable[T]) -> T | None:
+    ...
+
+
+def first(iterable: Iterable[T], default: T | None = None) -> T | None:
     """
     Returns the first item in the given iterable or `default` if empty, meaningful mostly with 'for' expressions.
     """
@@ -463,7 +476,7 @@ class ImageMetaKey:
     SPATIAL_SHAPE = "spatial_shape"
 
 
-def has_option(obj, keywords: str | Sequence[str]) -> bool:
+def has_option(obj: Callable, keywords: str | Sequence[str]) -> bool:
     """
     Return a boolean indicating whether the given callable `obj` has the `keywords` in its signature.
     """
@@ -504,7 +517,7 @@ def sample_slices(data: NdarrayOrTensor, dim: int = 1, as_indices: bool = True,
     return data[tuple(slices)]
 
 
-def check_parent_dir(path: PathLike, create_dir: bool = True):
+def check_parent_dir(path: PathLike, create_dir: bool = True) -> None:
     """
     Utility to check whether the parent directory of the `path` exists.
 
@@ -523,7 +536,14 @@ def check_parent_dir(path: PathLike, create_dir: bool = True):
             raise ValueError(f"the directory of specified path does not exist: `{path_dir}`.")
 
 
-def save_obj(obj, path: PathLike, create_dir: bool = True, atomic: bool = True, func: Callable | None = None, **kwargs):
+def save_obj(
+    obj: object,
+    path: PathLike,
+    create_dir: bool = True,
+    atomic: bool = True,
+    func: Callable | None = None,
+    **kwargs: Any,
+) -> None:
     """
     Save an object to file with specified path.
     Support to serialize to a temporary file first, then move to final destination,
@@ -576,7 +596,7 @@ def label_union(x: list) -> list:
     return list(set.union(set(np.array(x).tolist())))
 
 
-def prob2class(x, sigmoid: bool = False, threshold: float = 0.5, **kwargs):
+def prob2class(x: torch.Tensor, sigmoid: bool = False, threshold: float = 0.5, **kwargs: Any) -> torch.Tensor:
     """
     Compute the lab from the probability of predicted feature maps