temp commit

util/misc.py

@@ -1,4 +1,6 @@
+import importlib
 import logging
+import pkgutil
 from pathlib import Path
 from typing import Optional
 
@@ -12,6 +14,24 @@ def add_spectral_norm(module):
         return module
 
 
+def import_submodules(package, recursive=True):
+    """ Import all submodules of a module, recursively, including subpackages
+
+    :param package: package (name or actual module)
+    :type package: str | module
+    :rtype: dict[str, types.ModuleType]
+    """
+    if isinstance(package, str):
+        package = importlib.import_module(package)
+    results = {}
+    for loader, name, is_pkg in pkgutil.walk_packages(package.__path__):
+        full_name = package.__name__ + '.' + name
+        results[name] = importlib.import_module(full_name)
+        if recursive and is_pkg:
+            results.update(import_submodules(full_name))
+    return results
+
+
 def setup_logger(
         name: Optional[str] = None,
         level: int = logging.INFO,