From 9014308d569002ebb3deab9d0803477137c223c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=90=D0=BD=D1=82=D0=BE=D0=BD?=
 <88902421+BouFFaNTCRiB@users.noreply.github.com>
Date: Wed, 18 Oct 2023 10:06:56 +0300
Subject: [PATCH] refactor: simplify code in module importer (#6077)

---
 jina/importer.py | 125 +++++++++++++++++++++++++++--------------------
 1 file changed, 73 insertions(+), 52 deletions(-)

diff --git a/jina/importer.py b/jina/importer.py
index 2caa887968209..60402365335c6 100644
--- a/jina/importer.py
+++ b/jina/importer.py
@@ -39,61 +39,82 @@ def __init__(
 
     def __enter__(self):
         return self
+    
+    def _check_v(self, v, missing_module):
+        if (
+                v.strip()
+                and not v.startswith('#')
+                and v.startswith(missing_module)
+                and ':' in v
+            ):
+            return True
+
+    def _find_missing_module_in_extra_req(self, missing_module):
+        with open(os.path.join(__resources_path__, 'extra-requirements.txt'), encoding='utf-8') as fp:
+            for v in fp:
+                if self._check_v(v, missing_module):
+                    missing_module, install_tags = v.split(':')
+                    self._tags.append(missing_module)
+                    self._tags.extend(vv.strip() for vv in install_tags.split(','))
+                    break
+
+    def _find_missing_module(self, exc_val):
+        missing_module = self._pkg_name or exc_val.name
+        missing_module = self._find_missing_module_in_extra_req(missing_module)
+        return missing_module
+
+
+    def _err_msg(self, exc_val, missing_module):
+        if self._tags:
+            from jina.helper import colored
+
+            req_msg = colored('fallback to default behavior', color='yellow')
+            if self._required:
+                req_msg = colored('and it is required', color='red')
+            err_msg = f'''Python package "{colored(missing_module, attrs='bold')}" is not installed, {req_msg}.
+            You are trying to use a feature not enabled by your current Jina installation.'''
 
-    def __exit__(self, exc_type, exc_val, traceback):
-        if exc_type == ModuleNotFoundError:
-            missing_module = self._pkg_name or exc_val.name
-            with open(os.path.join(__resources_path__, 'extra-requirements.txt'), encoding='utf-8') as fp:
-                for v in fp:
-                    if (
-                        v.strip()
-                        and not v.startswith('#')
-                        and v.startswith(missing_module)
-                        and ':' in v
-                    ):
-                        missing_module, install_tags = v.split(':')
-                        self._tags.append(missing_module)
-                        self._tags.extend(vv.strip() for vv in install_tags.split(','))
-                        break
-
-            if self._tags:
-                from jina.helper import colored
-
-                req_msg = colored('fallback to default behavior', color='yellow')
-                if self._required:
-                    req_msg = colored('and it is required', color='red')
-                err_msg = f'''Python package "{colored(missing_module, attrs='bold')}" is not installed, {req_msg}.
-                    You are trying to use a feature not enabled by your current Jina installation.'''
-
-                avail_tags = ' '.join(
-                    colored(f'[{tag}]', attrs='bold') for tag in self._tags
-                )
-                err_msg += (
-                    f'\n\nTo enable this feature, use {colored("pip install jina[TAG]", attrs="bold")}, '
-                    f'where {colored("[TAG]", attrs="bold")} is one of {avail_tags}.\n'
-                )
+            avail_tags = ' '.join(
+                colored(f'[{tag}]', attrs='bold') for tag in self._tags
+            )
+            err_msg += (
+                f'\n\nTo enable this feature, use {colored("pip install jina[TAG]", attrs="bold")}, '
+                f'where {colored("[TAG]", attrs="bold")} is one of {avail_tags}.\n'
+            )
+        else:
+            err_msg = f'{exc_val.msg}'
+        return err_msg
+    
+    def _log_critical(self, err_msg):
+        if self._verbose and self._logger:
+            self._logger.critical(err_msg)
+            if self._help_text:
+                self._logger.error(self._help_text)
+
+    def _log_warning(self, err_msg):
+        if self._verbose and self._logger:
+            self._logger.warning(err_msg)
+            if self._help_text:
+                self._logger.info(self._help_text)   
+
+    def _raise_or_supress(self, err_msg, exc_val):
+        if self._verbose and not self._logger:
+            warnings.warn(err_msg, RuntimeWarning, stacklevel=2)
+        if self._required:
+            self._log_critical(err_msg)
+            raise exc_val
+        else:
+            self._log_warning(err_msg)
+            return True  # suppress the error
 
-            else:
-                err_msg = f'{exc_val.msg}'
 
-            if self._required:
-                if self._verbose:
-                    if self._logger:
-                        self._logger.critical(err_msg)
-                        if self._help_text:
-                            self._logger.error(self._help_text)
-                    else:
-                        warnings.warn(err_msg, RuntimeWarning, stacklevel=2)
-                raise exc_val
-            else:
-                if self._verbose:
-                    if self._logger:
-                        self._logger.warning(err_msg)
-                        if self._help_text:
-                            self._logger.info(self._help_text)
-                    else:
-                        warnings.warn(err_msg, RuntimeWarning, stacklevel=2)
-                return True  # suppress the error
+    def __exit__(self, exc_type, exc_val, traceback):
+        if exc_type != ModuleNotFoundError:
+            return
+        missing_module = self._find_missing_module(exc_val)
+        err_msg = self._err_msg(exc_val, missing_module)
+        return self._raise_or_supress(err_msg, exc_val)
+        
 
 
 def _path_import(absolute_path: str):