Factor out dynamic class finder/loader to a dedicated file (#2954)

It could be reused to do the same for MPP modules/classes.
Ref: #2940 and #2950

patroni/dcs/__init__.py

@@ -1,26 +1,21 @@
 """Abstract classes for Distributed Configuration Store."""
 import abc
 import datetime
-import importlib
-import inspect
 import json
 import logging
-import os
-import pkgutil
 import re
-import sys
 import time
 from collections import defaultdict
 from copy import deepcopy
 from random import randint
 from threading import Event, Lock
-from types import ModuleType
+from typing import Any, Callable, Collection, Dict, Iterator, List, \
-from typing import Any, Callable, Collection, Dict, List, NamedTuple, Optional, Set, Tuple, Union, TYPE_CHECKING, \
+    NamedTuple, Optional, Tuple, Type, TYPE_CHECKING, Union
-    Type, Iterator
 from urllib.parse import urlparse, urlunparse, parse_qsl
 
 import dateutil.parser
 
+from ..dynamic_loader import iter_classes, iter_modules
 from ..exceptions import PatroniFatalException
 from ..utils import deep_compare, uri
 from ..tags import Tags
@@ -87,28 +82,9 @@ def parse_connection_string(value: str) -> Tuple[str, Union[str, None]]:
 def dcs_modules() -> List[str]:
     """Get names of DCS modules, depending on execution environment.
 
-    .. note::
-        If being packaged with PyInstaller, modules aren't discoverable dynamically by scanning source directory because
-        :class:`importlib.machinery.FrozenImporter` doesn't implement :func:`iter_modules`. But it is still possible to
-        find all potential DCS modules by iterating through ``toc``, which contains list of all "frozen" resources.
-
     :returns: list of known module names with absolute python module path namespace, e.g. ``patroni.dcs.etcd``.
     """
-    dcs_dirname = os.path.dirname(__file__)
+    return iter_modules(__package__)
-    module_prefix = __package__ + '.'
-
-    if getattr(sys, 'frozen', False):
-        toc: Set[str] = set()
-        # dcs_dirname may contain a dot, which causes pkgutil.iter_importers()
-        # to misinterpret the path as a package name. This can be avoided
-        # altogether by not passing a path at all, because PyInstaller's
-        # FrozenImporter is a singleton and registered as top-level finder.
-        for importer in pkgutil.iter_importers():
-            if hasattr(importer, 'toc'):
-                toc |= getattr(importer, 'toc')
-        return [module for module in toc if module.startswith(module_prefix) and module.count('.') == 2]
-
-    return [module_prefix + name for _, name, is_pkg in pkgutil.iter_modules([dcs_dirname]) if not is_pkg]
 
 
 def iter_dcs_classes(
@@ -122,44 +98,16 @@ def iter_dcs_classes(
     :param config: configuration information with possible DCS names as keys. If given, only attempt to import DCS
                    modules defined in the configuration. Else, if ``None``, attempt to import any supported DCS module.
 
-    :yields: a tuple containing the module ``name`` and the imported DCS class object.
+    :returns: an iterator of tuples, each containing the module ``name`` and the imported DCS class object.
     """
-    for mod_name in dcs_modules():
+    return iter_classes(__package__, AbstractDCS, config)
-        name = mod_name.rpartition('.')[2]
-        if config is None or name in config:
-
-            try:
-                module = importlib.import_module(mod_name)
-                dcs_module = find_dcs_class_in_module(module)
-                if dcs_module:
-                    yield name, dcs_module
-
-            except ImportError:
-                logger.log(logging.DEBUG if config is not None else logging.INFO,
-                           'Failed to import %s', mod_name)
-
-
-def find_dcs_class_in_module(module: ModuleType) -> Optional[Type['AbstractDCS']]:
-    """Try to find the implementation of :class:`AbstractDCS` interface in *module* matching the *module* name.
-
-    :param module: Imported DCS module.
-
-    :returns: class with a name matching the name of *module* that implements :class:`AbstractDCS` or ``None`` if not
-              found.
-    """
-    module_name = module.__name__.rpartition('.')[2]
-    return next(
-        (obj for obj_name, obj in module.__dict__.items()
-         if (obj_name.lower() == module_name
-             and inspect.isclass(obj) and issubclass(obj, AbstractDCS))),
-        None)
 
 
 def get_dcs(config: Union['Config', Dict[str, Any]]) -> 'AbstractDCS':
     """Attempt to load a Distributed Configuration Store from known available implementations.
 
     .. note::
-        Using the list of available DCS classes returned by :func:`iter_dcs_classes` attempt to dynamically
+        Using the list of available DCS classes returned by :func:`iter_classes` attempt to dynamically
         instantiate the class that implements a DCS using the abstract class :class:`AbstractDCS`.
 
         Basic top-level configuration parameters retrieved from *config* are propagated to the DCS specific config
@@ -185,9 +133,9 @@ def get_dcs(config: Union['Config', Dict[str, Any]]) -> 'AbstractDCS':
             config[name].update(config['citus'])
         return dcs_class(config[name])
 
-    raise PatroniFatalException(
+    available_implementations = ', '.join(sorted([n for n, _ in iter_dcs_classes()]))
-        f"Can not find suitable configuration of distributed configuration store\n"
+    raise PatroniFatalException("Can not find suitable configuration of distributed configuration store\n"
-        f"Available implementations: {', '.join(sorted([n for n, _ in iter_dcs_classes()]))}")
+                                f"Available implementations: {available_implementations}")
 
 
 _Version = Union[int, str]

patroni/dynamic_loader.py

@@ -0,0 +1,96 @@
+"""Helper functions to search for implementations of specific abstract interface in a package."""
+import importlib
+import inspect
+import logging
+import os
+import pkgutil
+import sys
+from types import ModuleType
+
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, TYPE_CHECKING, Type, TypeVar, Union
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .config import Config
+
+logger = logging.getLogger(__name__)
+
+
+def iter_modules(package: str) -> List[str]:
+    """Get names of modules from *package*, depending on execution environment.
+
+    .. note::
+        If being packaged with PyInstaller, modules aren't discoverable dynamically by scanning source directory because
+        :class:`importlib.machinery.FrozenImporter` doesn't implement :func:`iter_modules`. But it is still possible to
+        find all potential modules by iterating through ``toc``, which contains list of all "frozen" resources.
+
+    :param package: a package name to search modules in, e.g. ``patroni.dcs``.
+
+    :returns: list of known module names with absolute python module path namespace, e.g. ``patroni.dcs.etcd``.
+    """
+    module_prefix = package + '.'
+
+    if getattr(sys, 'frozen', False):
+        toc: Set[str] = set()
+        # dirname may contain a few dots, which causes pkgutil.iter_importers()
+        # to misinterpret the path as a package name. This can be avoided
+        # altogether by not passing a path at all, because PyInstaller's
+        # FrozenImporter is a singleton and registered as top-level finder.
+        for importer in pkgutil.iter_importers():
+            if hasattr(importer, 'toc'):
+                toc |= getattr(importer, 'toc')
+        dots = module_prefix.count('.')  # search for modules only on the same level
+        return [module for module in toc if module.startswith(module_prefix) and module.count('.') == dots]
+
+    # here we are making an assumption that the package which is calling this function is already imported
+    pkg_file = sys.modules[package].__file__
+    if TYPE_CHECKING:  # pragma: no cover
+        assert isinstance(pkg_file, str)
+    return [name for _, name, is_pkg in pkgutil.iter_modules([os.path.dirname(pkg_file)], module_prefix) if not is_pkg]
+
+
+ClassType = TypeVar("ClassType")
+
+
+def find_class_in_module(module: ModuleType, cls_type: Type[ClassType]) -> Optional[Type[ClassType]]:
+    """Try to find the implementation of *cls_type* class interface in *module* matching the *module* name.
+
+    :param module: imported module.
+    :param cls_type: a class type we are looking for.
+
+    :returns: class with a name matching the name of *module* that implements *cls_type* or ``None`` if not found.
+    """
+    module_name = module.__name__.rpartition('.')[2]
+    return next(
+        (obj for obj_name, obj in module.__dict__.items()
+         if (obj_name.lower() == module_name
+             and inspect.isclass(obj) and issubclass(obj, cls_type))),
+        None)
+
+
+def iter_classes(
+        package: str, cls_type: Type[ClassType],
+        config: Optional[Union['Config', Dict[str, Any]]] = None
+) -> Iterator[Tuple[str, Type[ClassType]]]:
+    """Attempt to import modules and find implementations of *cls_type* that are present in the given configuration.
+
+    .. note::
+            If a module successfully imports we can assume that all its requirements are installed.
+
+    :param package: a package name to search modules in, e.g. ``patroni.dcs``.
+    :param cls_type: a class type we are looking for.
+    :param config: configuration information with possible module names as keys. If given, only attempt to import
+                   modules defined in the configuration. Else, if ``None``, attempt to import any supported module.
+
+    :yields: a tuple containing the module ``name`` and the imported class object.
+    """
+    for mod_name in iter_modules(package):
+        name = mod_name.rpartition('.')[2]
+        if config is None or name in config:
+            try:
+                module = importlib.import_module(mod_name)
+                module_cls = find_class_in_module(module, cls_type)
+                if module_cls:
+                    yield name, module_cls
+            except ImportError:
+                logger.log(logging.DEBUG if config is not None else logging.INFO,
+                           'Failed to import %s', mod_name)

tests/test_ctl.py

@@ -238,7 +238,7 @@ class TestCtl(unittest.TestCase):
         result = self.runner.invoke(ctl, ['failover', 'dummy', '--group', '0', '--candidate', 'async'], input='N')
         self.assertEqual(result.exit_code, 1)
 
-    @patch('patroni.dcs.dcs_modules', Mock(return_value=['patroni.dcs.dummy', 'patroni.dcs.etcd']))
+    @patch('patroni.dynamic_loader.iter_modules', Mock(return_value=['patroni.dcs.dummy', 'patroni.dcs.etcd']))
     def test_get_dcs(self):
         with click.Context(click.Command('list')) as ctx:
             ctx.obj = {'__config': {'dummy': {}}}