lazyloader.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. """module lazyloader."""
  2. import importlib
  3. import sys
  4. import types
  5. class LazyLoader(types.ModuleType):
  6. """Lazily import a module, mainly to avoid pulling in large dependencies.
  7. We use this for tensorflow and other optional libraries primarily at the
  8. top module level.
  9. """
  10. # The lint error here is incorrect.
  11. def __init__(
  12. self,
  13. local_name, # pylint: disable=super-on-old-class
  14. parent_module_globals,
  15. name,
  16. warning=None,
  17. ):
  18. self._local_name = local_name
  19. self._parent_module_globals = parent_module_globals
  20. self._warning = warning
  21. super().__init__(str(name))
  22. def _load(self):
  23. """Load the module and insert it into the parent's globals."""
  24. # Import the target module and insert it into the parent's namespace
  25. module = importlib.import_module(self.__name__)
  26. self._parent_module_globals[self._local_name] = module
  27. # print("import", self.__name__)
  28. # print("Set global", self._local_name)
  29. # print("mod", module)
  30. sys.modules[self._local_name] = module
  31. # Emit a warning if one was specified
  32. if self._warning:
  33. print(self._warning) # noqa: T201
  34. # Make sure to only warn once.
  35. self._warning = None
  36. # Update this object's dict so that if someone keeps a reference to the
  37. # LazyLoader, lookups are efficient (__getattr__ is only called on lookups
  38. # that fail).
  39. self.__dict__.update(module.__dict__)
  40. return module
  41. # def __getattribute__(self, item):
  42. # print("getattribute", item)
  43. def __getattr__(self, item):
  44. # print("getattr", item)
  45. module = self._load()
  46. return getattr(module, item)
  47. def __dir__(self):
  48. # print("dir")
  49. module = self._load()
  50. return dir(module)