diff --git a/Wrapping/Generators/Python/itk/support/helpers.py b/Wrapping/Generators/Python/itk/support/helpers.py index 07da3f94e28..5ff75eb8270 100644 --- a/Wrapping/Generators/Python/itk/support/helpers.py +++ b/Wrapping/Generators/Python/itk/support/helpers.py @@ -16,6 +16,8 @@ # # ==========================================================================*/ +import importlib +from importlib.metadata import metadata import os import re import functools @@ -24,17 +26,17 @@ _HAVE_XARRAY = False try: - import xarray as xr + metadata('xarray') _HAVE_XARRAY = True except ImportError: pass _HAVE_TORCH = False try: - import torch + metadata('torch') _HAVE_TORCH = True -except ImportError: +except importlib.metadata.PackageNotFoundError: pass @@ -84,6 +86,10 @@ def accept_array_like_xarray_torch(image_filter): If a xarray DataArray is passed as an input, output itk.Image's are converted to xarray.DataArray's.""" import numpy as np import itk + if _HAVE_XARRAY: + import xarray as xr + if _HAVE_TORCH: + import torch @functools.wraps(image_filter) def image_filter_wrapper(*args, **kwargs): diff --git a/Wrapping/Generators/Python/itk/support/types.py b/Wrapping/Generators/Python/itk/support/types.py index 643f407e727..524c3a5868e 100644 --- a/Wrapping/Generators/Python/itk/support/types.py +++ b/Wrapping/Generators/Python/itk/support/types.py @@ -16,6 +16,8 @@ # # ==========================================================================*/ +import importlib +from importlib.metadata import metadata from typing import Union, Optional, Tuple, TYPE_CHECKING import os @@ -26,17 +28,17 @@ _HAVE_XARRAY = False try: - import xarray as xr + metadata('xarray') _HAVE_XARRAY = True -except ImportError: +except importlib.metadata.PackageNotFoundError: pass _HAVE_TORCH = False try: - import torch + metadata('torch') _HAVE_TORCH = True -except ImportError: +except importlib.metadata.PackageNotFoundError: pass # noinspection PyPep8Naming @@ -218,11 +220,11 @@ def initialize_c_types_once() -> ( ImageOrImageSource = Union[ImageBase, ImageSource] # Can be coerced into an itk.ImageBase if _HAVE_XARRAY and _HAVE_TORCH: - ImageLike = Union[ImageBase, ArrayLike, xr.DataArray, torch.Tensor] + ImageLike = Union[ImageBase, ArrayLike, "xr.DataArray", "torch.Tensor"] elif _HAVE_XARRAY: - ImageLike = Union[ImageBase, ArrayLike, xr.DataArray] + ImageLike = Union[ImageBase, ArrayLike, "xr.DataArray"] elif _HAVE_TORCH: - ImageLike = Union[ImageBase, ArrayLike, torch.Tensor] + ImageLike = Union[ImageBase, ArrayLike, "torch.Tensor"] else: ImageLike = Union[ImageBase, ArrayLike]