make dataflow.imgaug and dataflow.dataset lazy

tensorpack/dataflow/__init__.py

@@ -5,6 +5,7 @@
 from pkgutil import iter_modules
 import os
 import os.path
+from ..utils.develop import LazyLoader
 
 __all__ = []
 
@@ -30,14 +31,9 @@ for _, module_name, __ in iter_modules(
         _global_import(module_name)
 
 
-class _LazyModule(object):
-    def __init__(self, modname):
-        self._modname = modname
+dataset = LazyLoader('dataset', globals(), 'tensorpack.dataflow.dataset')
+imgaug = LazyLoader('imgaug', globals(), 'tensorpack.dataflow.imgaug')
 
-    def __getattr__(self, name):
-        dataset = __import__(self._modname, globals(), locals(), [name], 1)
-        return getattr(dataset, name)
+del LazyLoader
 
-
-dataset = _LazyModule('dataset')
 __all__.extend(['imgaug', 'dftools', 'dataset'])

tensorpack/utils/develop.py

@@ -10,6 +10,9 @@ appeared in docs."""
 import os
 import functools
 from datetime import datetime
+import importlib
+import types
+
 from . import logger
 
 
@@ -117,3 +120,31 @@ def deprecated(text="", eos=""):
             return func(*args, **kwargs)
         return new_func
     return deprecated_inner
+
+
+# Copied from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/lazy_loader.py
+class LazyLoader(types.ModuleType):
+    def __init__(self, local_name, parent_module_globals, name):
+        self._local_name = local_name
+        self._parent_module_globals = parent_module_globals
+        super(LazyLoader, self).__init__(name)
+
+    def _load(self):
+        # Import the target module and insert it into the parent's namespace
+        module = importlib.import_module(self.__name__)
+        self._parent_module_globals[self._local_name] = module
+
+        # Update this object's dict so that if someone keeps a reference to the
+        #   LazyLoader, lookups are efficient (__getattr__ is only called on lookups
+        #   that fail).
+        self.__dict__.update(module.__dict__)
+
+        return module
+
+    def __getattr__(self, item):
+        module = self._load()
+        return getattr(module, item)
+
+    def __dir__(self):
+        module = self._load()
+        return dir(module)