From f7e3355bee107653d6be6dcc3cd9afecd23a2c40 Mon Sep 17 00:00:00 2001
From: Joan Fontanals <joan.martinez@jina.ai>
Date: Fri, 19 Mar 2021 06:46:39 +0100
Subject: [PATCH] fix: use import extensions (#2204)

---
 jina/executors/indexers/vector.py | 10 +++-------
 jina/logging/profile.py           | 19 ++++++++++---------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/jina/executors/indexers/vector.py b/jina/executors/indexers/vector.py
index 21b576ee3960f..c92bd05933d45 100644
--- a/jina/executors/indexers/vector.py
+++ b/jina/executors/indexers/vector.py
@@ -13,6 +13,7 @@
 from . import BaseVectorIndexer
 from ..decorators import batching
 from ...helper import cached_property
+from ...importer import ImportExtensions
 
 
 class BaseNumpyIndexer(BaseVectorIndexer):
@@ -520,11 +521,6 @@ def _cosine(self, cached_A, raw_B):
 
     @batching(merge_over_axis=1, slice_on=2)
     def _cdist(self, *args, **kwargs):
-        try:
+        with ImportExtensions(required=True):
             from scipy.spatial.distance import cdist
-
-            return cdist(*args, **kwargs, metric=self.metric)
-        except ModuleNotFoundError:
-            raise ModuleNotFoundError(
-                f'your metric {self.metric} requires scipy, but scipy is not found'
-            )
+        return cdist(*args, **kwargs, metric=self.metric)
diff --git a/jina/logging/profile.py b/jina/logging/profile.py
index 5047264d1506f..e3e4cab864c7d 100644
--- a/jina/logging/profile.py
+++ b/jina/logging/profile.py
@@ -7,11 +7,12 @@
 from functools import wraps
 from typing import Optional
 
+from ..importer import ImportExtensions
 from ..helper import colored, get_readable_size, get_readable_time
 
 if False:
     # fix type-hint complain for sphinx and flake
-    from . import JinaLogger
+    from ..logging import JinaLogger
 
 
 def used_memory(unit: int = 1024 * 1024 * 1024) -> float:
@@ -21,18 +22,18 @@ def used_memory(unit: int = 1024 * 1024 * 1024) -> float:
     :param unit: Unit of the memory, default in Gigabytes.
     :return: Memory usage of the current process.
     """
-    try:
+    with ImportExtensions(required=False):
         import resource
 
         return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / unit
-    except ModuleNotFoundError:
-        from . import default_logger
 
-        default_logger.error(
-            'module "resource" can not be found and you are likely running it on Windows, '
-            'i will return 0'
-        )
-        return 0
+    from . import default_logger
+
+    default_logger.error(
+        'module "resource" can not be found and you are likely running it on Windows, '
+        'i will return 0'
+    )
+    return 0
 
 
 def used_memory_readable() -> str: