From 963a6f426b30a8d099690f15f5249948b4a5c0aa Mon Sep 17 00:00:00 2001
From: Thomas BAUDIER <thomas.baudier@creatis.insa-lyon.fr>
Date: Thu, 14 Dec 2023 17:56:45 +0100
Subject: [PATCH] PERF: Do not import torch to reduce itk import time

https://discourse.itk.org/t/torch-import-time/6354
Torch takes time to import. In types.py torch is imported
to check if it's present or not in the environment, but
it's never used.
Now, torch is detected with importlib.metadata and imported
later.

Do the same thing with xarray
---
 .../Generators/Python/itk/support/helpers.py     | 12 +++++++++---
 Wrapping/Generators/Python/itk/support/types.py  | 16 +++++++++-------
 2 files changed, 18 insertions(+), 10 deletions(-)

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]