Repository: jina-ai/serve
Pull request: fix(importer): fix jina check
resolved #1526 
- [x] Add unit tests for mocking the loading errors in `import_module`
- [x] Refactor the `import_classes` to improve the readability
- [x] Remove the unused `targets` argument from `import_classes`.  The only usage of `targets`argument is removed in #1759

Problem statement:
`jina check` is broken
**Describe the bug**

jina check can not detect any broken hub executors even though they are warned during *first* imports.

**Describe how you solve it**

Maybe something related to the import orders in `jina/__init__.py`

---

<!-- Optional, but really help us locate the problem faster -->

**Environment**
<!-- Run `jina --version-full` and copy paste the output here -->

**Screenshots**
<!-- If applicable, add screenshots to help explain your problem. -->

Hints from issue discussion:
This issue is stale because it has been open 20 days with no activity. Remove stale label or comment or this will be closed in 4 days
need someone to address this important ticket, @nan-wang can you make sure at least we have a test for it?

Commit message: refactor: clean up import_class

Patch to review:
diff --git a/jina/importer.py b/jina/importer.py
index e6c317a6717c1..e9931e9cb6c8a 100644
--- a/jina/importer.py
+++ b/jina/importer.py
@@ -11,6 +11,72 @@
 IMPORTED.hub = False
 
 
+def _get_targets(targets):
+    ret = None
+    if isinstance(targets, str):
+        ret = {targets}
+    elif isinstance(targets, list):
+        ret = set(targets)
+    elif targets is None:
+        ret = {}
+    return ret
+
+
+def _raise_bad_imports_warnings(bad_imports, namespace):
+    if not bad_imports:
+        return
+    if namespace != 'jina.hub':
+        warnings.warn(
+            f'theses modules or classes can not be imported {bad_imports}. '
+            f'You can use `jina check` to list all executors and drivers')
+    else:
+        warnings.warn(
+            f'due to the missing dependencies or bad implementations, '
+            f'{bad_imports} can not be imported '
+            f'if you are using these executors/drivers, they wont work. '
+            f'You can use `jina check` to list all executors and drivers')
+
+
+def _get_modules(namespace):
+    from setuptools import find_packages
+    from pkgutil import get_loader
+
+    try:
+        path = os.path.dirname(get_loader(namespace).path)
+    except AttributeError:
+        if namespace == 'jina.hub':
+            warnings.warn(f'hub submodule is not initialized. Please try "git submodule update --init"', ImportWarning)
+        return {}
+
+    modules = _get_submodules(path, namespace)
+
+    for pkg in find_packages(path):
+        modules.add('.'.join([namespace, pkg]))
+        pkgpath = os.path.join(path, pkg.replace('.', '/'))
+        modules.union(_get_submodules(pkgpath, namespace, prefix=pkg))
+
+    return _filter_modules(modules)
+
+
+def _get_submodules(path, namespace, prefix=None):
+    from pkgutil import iter_modules
+    _prefix = '.'.join([namespace, prefix]) if prefix else namespace
+    modules = set()
+    for info in iter_modules([path]):
+        is_hub_module = namespace == 'jina.hub' and info.ispkg
+        is_nonhub_module = namespace != 'jina.hub' and not info.ispkg
+        module_name = '.'.join([_prefix, info.name])
+        if is_hub_module or is_nonhub_module:
+            modules.add(module_name)
+    return modules
+
+
+def _filter_modules(modules):
+    import re
+    ignored_module_pattern = re.compile(r'\.tests|\.api|\.bump_version')
+    return {m for m in modules if not ignored_module_pattern.findall(m)}
+
+
 def import_classes(namespace: str, targets=None,
                    show_import_table: bool = False,
                    import_once: bool = False):
@@ -26,62 +92,32 @@ def import_classes(namespace: str, targets=None,
 
     import os, re
 
-    if namespace == 'jina.executors':
-        import_type = 'ExecutorType'
-        if import_once and IMPORTED.executors:
-            return
-    elif namespace == 'jina.drivers':
-        import_type = 'DriverType'
-        if import_once and IMPORTED.drivers:
-            return
-    elif namespace == 'jina.hub':
-        import_type = 'ExecutorType'
-        if import_once and IMPORTED.hub:
-            return
-    else:
+    _namespace2type = {
+        'jina.executors': 'ExecutorType',
+        'jina.drivers': 'DriverType',
+        'jina.hub': 'ExecutorType'
+    }
+    _import_type = _namespace2type.get(namespace)
+    if _import_type is None:
         raise TypeError(f'namespace: {namespace} is unrecognized')
 
-    from setuptools import find_packages
-    import pkgutil
-    from pkgutil import iter_modules
+    _targets = _get_targets(targets)
+    if _targets is None:
+        raise TypeError(f'target must be a set, but received {targets!r}')
 
-    try:
-        path = os.path.dirname(pkgutil.get_loader(namespace).path)
-    except AttributeError:
-        if namespace == 'jina.hub':
-            warnings.warn(f'hub submodule is not initialized. Please try "git submodule update --init"', ImportWarning)
+    _imported_property = namespace.split('.')[-1]
+    _is_imported = getattr(IMPORTED, _imported_property)
+    if import_once and _is_imported:
         return {}
 
-    modules = set()
-
-    for info in iter_modules([path]):
-        if (namespace != 'jina.hub' and not info.ispkg) or (namespace == 'jina.hub' and info.ispkg):
-            modules.add('.'.join([namespace, info.name]))
-
-    for pkg in find_packages(path):
-        modules.add('.'.join([namespace, pkg]))
-        pkgpath = path + '/' + pkg.replace('.', '/')
-        for info in iter_modules([pkgpath]):
-            if (namespace != 'jina.hub' and not info.ispkg) or (namespace == 'jina.hub' and info.ispkg):
-                modules.add('.'.join([namespace, pkg, info.name]))
-
-    # filter
-    ignored_module_pattern = r'\.tests|\.api|\.bump_version'
-    modules = {m for m in modules if not re.findall(ignored_module_pattern, m)}
-
+    modules = _get_modules(namespace)
+    if not modules:
+        return {}
+    
     from collections import defaultdict
     load_stat = defaultdict(list)
     bad_imports = []
 
-    if isinstance(targets, str):
-        targets = {targets}
-    elif isinstance(targets, list):
-        targets = set(targets)
-    elif targets is None:
-        targets = {}
-    else:
-        raise TypeError(f'target must be a set, but received {targets!r}')
-
     depend_tree = {}
     import importlib
     from .helper import colored
@@ -90,7 +126,7 @@ def import_classes(namespace: str, targets=None,
             mod = importlib.import_module(m)
             for k in dir(mod):
                 # import the class
-                if (getattr(mod, k).__class__.__name__ == import_type) and (not targets or k in targets):
+                if (getattr(mod, k).__class__.__name__ == _import_type) and (not _targets or k in _targets):
                     try:
                         _c = getattr(mod, k)
 
@@ -100,9 +136,9 @@ def import_classes(namespace: str, targets=None,
                                 d[vvv.__name__] = {}
                             d = d[vvv.__name__]
                         d['module'] = m
-                        if k in targets:
-                            targets.remove(k)
-                            if not targets:
+                        if k in _targets:
+                            _targets.remove(k)
+                            if not _targets:
                                 return getattr(mod, k)  # target execs are all found and loaded, return
                         try:
                             # load the default request for this executor if possible
@@ -115,37 +151,21 @@ def import_classes(namespace: str, targets=None,
                     except Exception as ex:
                         load_stat[m].append((k, False, ex))
                         bad_imports.append('.'.join([m, k]))
-                        if k in targets:
+                        if k in _targets:
                             raise ex  # target class is found but not loaded, raise return
         except Exception as ex:
             load_stat[m].append(('', False, ex))
             bad_imports.append(m)
 
-    if targets:
-        raise ImportError(f'{targets} can not be found/load')
+    if _targets:
+        raise ImportError(f'{_targets} can not be found/load')
 
     if show_import_table:
         _print_load_table(load_stat)
-
     else:
-        if bad_imports:
-            if namespace != 'jina.hub':
-                warnings.warn(
-                    f'theses modules or classes can not be imported {bad_imports}. '
-                    f'You can use `jina check` to list all executors and drivers')
-            else:
-                warnings.warn(
-                    f'due to the missing dependencies or bad implementations, '
-                    f'{bad_imports} can not be imported '
-                    f'if you are using these executors/drivers, they wont work. '
-                    f'You can use `jina check` to list all executors and drivers')
-
-    if namespace == 'jina.executors':
-        IMPORTED.executors = True
-    elif namespace == 'jina.drivers':
-        IMPORTED.drivers = True
-    elif namespace == 'jina.hub':
-        IMPORTED.hub = True
+        _raise_bad_imports_warnings(bad_imports, namespace)
+
+    setattr(IMPORTED, _imported_property, True)
 
     return depend_tree
 
diff --git a/tests/unit/test_importer.py b/tests/unit/test_importer.py
index 4d9de482c8131..4148b6bb1daed 100644
--- a/tests/unit/test_importer.py
+++ b/tests/unit/test_importer.py
@@ -1,6 +1,6 @@
 import pytest
 
-from jina.importer import ImportExtensions
+from jina.importer import ImportExtensions, import_classes
 from jina.logging import default_logger
 
 
@@ -36,7 +36,6 @@ def test_bad_import():
     assert not ie._tags
 
 
-
 def test_no_suppress_other_exception():
     with pytest.raises(Exception):
         with ImportExtensions(required=False, logger=default_logger):
@@ -44,4 +43,49 @@ def test_no_suppress_other_exception():
 
     with pytest.raises(Exception):
         with ImportExtensions(required=True, logger=default_logger):
-            raise Exception
\ No newline at end of file
+            raise Exception
+
+
+@pytest.mark.parametrize('import_once', [True, False])
+@pytest.mark.parametrize('ns', ['jina.executors', 'jina.hub', 'jina.drivers'])
+def test_import_classes_import_once(ns, import_once):
+    depend_tree = import_classes(namespace=ns, import_once=import_once)
+    print(depend_tree)
+    assert (not depend_tree) == import_once
+
+
+@pytest.mark.parametrize('import_once', [True, False])
+def test_import_classes_import_once_exception(import_once):
+    with pytest.raises(TypeError):
+        _ = import_classes(namespace='dummy', import_once=import_once)
+
+
+@pytest.mark.parametrize('ns', ['jina.executors', 'jina.hub', 'jina.drivers'])
+def test_import_classes_failed_find_package(ns, mocker):
+    mocker.patch('pkgutil.get_loader', return_value=None)
+    depend_tree = import_classes(namespace=ns)
+    assert len(depend_tree) == 0
+
+
+@pytest.mark.parametrize('ns', ['jina.executors', 'jina.hub', 'jina.drivers'])
+def test_import_classes_failed_import_module(ns, mocker, recwarn):
+    import importlib
+    mocker.patch.object(importlib, 'import_module', side_effect=Exception('mocked error'))
+    depend_tree = import_classes(namespace=ns)
+    assert len(depend_tree) == 0
+    assert len(recwarn) == 1
+    assert 'You can use `jina check` to list all executors and drivers' in recwarn[0].message.args[0]
+
+
+@pytest.mark.parametrize('print_table', [True, False])
+@pytest.mark.parametrize('ns', ['jina.executors', 'jina.hub', 'jina.drivers'])
+def test_import_classes_failed_get_default_reqs(ns, print_table, mocker, recwarn, capsys):
+    mocker.patch('pkg_resources.resource_stream', side_effect=Exception('mocked error'))
+    _ = import_classes(namespace=ns, show_import_table=print_table)
+    if print_table:
+        captured = capsys.readouterr()
+        assert '✗' in captured.out
+    else:
+        assert len(recwarn) == 1
+        assert 'You can use `jina check` to list all executors and drivers' in recwarn[0].message.args[0]
+


Code review comments:
diff_hunk: @@ -26,62 +92,32 @@ def import_classes(namespace: str, targets=None,
 
     import os, re
 
-    if namespace == 'jina.executors':
-        import_type = 'ExecutorType'
-        if import_once and IMPORTED.executors:
-            return
-    elif namespace == 'jina.drivers':
-        import_type = 'DriverType'
-        if import_once and IMPORTED.drivers:
-            return
-    elif namespace == 'jina.hub':
-        import_type = 'ExecutorType'
-        if import_once and IMPORTED.hub:
-            return
-    else:
+    _namespace2type = {
+        'jina.executors': 'ExecutorType',
+        'jina.drivers': 'DriverType',
+        'jina.hub': 'ExecutorType'
+    }
+    _import_type = _namespace2type.get(namespace)
+    if _import_type is None:
         raise TypeError(f'namespace: {namespace} is unrecognized')
 
-    from setuptools import find_packages
-    import pkgutil
-    from pkgutil import iter_modules
+    _targets = _get_targets(targets)
+    if _targets is None:
+        raise TypeError(f'target must be a set, but received {targets!r}')
path: jina/importer.py
text: @user1:
This can never happen, since `_get_targets` will never return `None`. Regardless, I'd rename `_get_targets` into `_parse_targets` and do any error handling there.

@author:
As `targets` is no longer used after #1759, I've removed this part.