Python 包懒加载

在Python中，模块导入通常是在程序开始时进行的，这意味着一旦模块被导入，其相关的代码和资源就会被加载到内存中，直到程序结束。尽管这在大多数情况下都能正常工作，但在处理大型项目或具有复杂依赖关系的代码时，这种“早期导入”的方式可能会导致不必要的内存使用和启动延迟。本文记录 Python 在 import 时常用的懒加载实现方案。

简介

为了解决这个问题，Python提供了惰性导入的概念，即在真正需要时才导入模块。这可以通过几种方式实现，包括使用importlib模块、使用__import__函数、以及使用延迟执行的技术（如lambda函数或装饰器）。

_LazyModule

https://github.com/huggingface/diffusers

huggingface 包的组织方式中出现了_LazyModule，这个模块实际上对应了python中的一种lazy import的思想。也就是在整个包很大的情况下不再将所有的包都import，而仅仅在使用的时候进行真正的import。这么做可以极大地缩短整个的import时间。

核心实现

源码链接： https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/import_utils.py

# Inspired by diffusers repo
# https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/import_utils.py
import os
import importlib.util
from itertools import chain
from types import ModuleType
from typing import Any


class _LazyModule(ModuleType):
    """
    Module class that surfaces all objects but only performs associated imports when the objects are requested.
    """

    # Very heavily inspired by optuna.integration._IntegrationModule
    # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
    def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):
        super().__init__(name)
        self._modules = set(import_structure.keys())
        self._class_to_module = {}
        for key, values in import_structure.items():
            for value in values:
                self._class_to_module[value] = key
        # Needed for autocompletion in an IDE
        self.__all__ = list(import_structure.keys()) + \
            list(chain(*import_structure.values()))
        self.__file__ = module_file
        self.__spec__ = module_spec
        self.__path__ = [os.path.dirname(module_file)]
        self._objects = {} if extra_objects is None else extra_objects
        self._name = name
        self._import_structure = import_structure

    # Needed for autocompletion in an IDE
    def __dir__(self):
        result = super().__dir__()
        # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether
        # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir.
        for attr in self.__all__:
            if attr not in result:
                result.append(attr)
        return result

    def __getattr__(self, name: str) -> Any:
        if name in self._objects:
            return self._objects[name]
        if name in self._modules:
            value = self._get_module(name)
        elif name in self._class_to_module.keys():
            module = self._get_module(self._class_to_module[name])
            value = getattr(module, name)
        else:
            raise AttributeError(
                f"module {self.__name__} has no attribute {name}")

        setattr(self, name, value)
        return value

    def _get_module(self, module_name: str):
        try:
            return importlib.import_module("." + module_name, self.__name__)
        except Exception as e:
            raise RuntimeError(
                f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
                f" traceback):\n{e}"
            ) from e

    def __reduce__(self):
        return (self.__class__, (self._name, self.__file__, self._import_structure))

使用方法

首先看一下最高层的包是如何调用这个 lazy import 类的，也就是package_name下的 __init__.py：

# Only support lazy import for now.

# TODO: support slow import

import sys
__version__ = "0.1"

from .utils import (
    _LazyModule
)

_import_structure = {
    "pipelines": []
}

_import_structure["pipelines"].extend(
    [
        "a",
        "LayoutDMPipeline"
    ]
)


sys.modules[__name__] = _LazyModule(
    __name__,
    globals()["__file__"],
    _import_structure,
    module_spec=__spec__,
    extra_objects={"__version__": __version__},
)

我们首先用了一个字典将我们想要导入的东西包起来，然后一起喂给_LazyModule，最后由_LazyModule传给sys.modules

这个包起来的东西包含了根目录的下级目录，extend的部分是我们最终想要导入的东西。流程是这样：

我们想要一个名为a的东西，那么packge就会去找下级目录找a，如果下级目录能够找到a，那么显然可以直接 from XX import a。但是问题出在下级目录显然也没有 a，下级目录又要到下下级目录中去找，直至找到。我们不妨看看叶子的__init__.py

# Only support lazy import for now.
# TODO: support slow import
 
import sys
from ...utils import (
    _LazyModule
)
 
_import_structure = {}
_import_structure["bar"] = ["a"]
 
 
sys.modules[__name__] = _LazyModule(
    __name__,
    globals()["__file__"],
    _import_structure,
    module_spec=__spec__,
)

总结：

我们通过from XX import a，python有如下操作

从 XX 中要a，XX再找pipelines要a，pipelines再找foo要a，foo再找bar要a，最终要到了a。

这么一来，from XX import a, from XX.pipelines import a, from XX.pipelines.foo import a, from XX.pipelines.foo.bar import a，全部都是可用的。

TensorFlow

Tensorflow 的 LazyLoad 方案：

核心实现

# Code copied from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/lazy_loader.py
"""A LazyLoader class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import importlib
import types


class LazyLoader(types.ModuleType):
  """Lazily import a module, mainly to avoid pulling in large dependencies.

  `contrib`, and `ffmpeg` are examples of modules that are large and not always
  needed, and this allows them to only be loaded when they are used.
  """

  # The lint error here is incorrect.
  def __init__(self, local_name, parent_module_globals, name):  # pylint: disable=super-on-old-class
    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)

使用方法

• 代码说明：

类 LazyLoader 继承自 types.ModuleType，初始化函数确保惰性模块将像真正的模块一样正确添加到全局变量中，只要真正用到模块的时候，也就是执行 getattr 或 dir 时，才会真正的 import 实际模块，更新全局变量以指向实际模块，并且将其所有状态（dict）更新为实际模块的状态，以便对延迟加载的引用，加载模块不需要每次访问都经过加载过程。

• 代码使用：

正常情况下我们这样导入模块：

import tensorflow.contrib as contrib1.

其对应的惰性导入版本如下：

contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib')

Python 原生实现

importlib

importlib是Python标准库中的一个模块，提供了动态导入模块的功能。通过importlib.import_module函数，可以在运行时按需导入模块。

import importlib
def lazy_import(module_name):
    return importlib.import_module(module_name)
# 当需要使用时才导入模块
numpy = lazy_import('numpy')
array = numpy.array([1, 2, 3])

__import__

__import__是一个内置函数，其行为类似于import语句。它接受模块名作为字符串，并返回导入的模块对象。

def lazy_import(module_name):
    return __import__(module_name)
# 当需要使用时才导入模块
numpy = lazy_import('numpy')
array = numpy.array([1, 2, 3])

lambda函数

另一种实现惰性导入的方法是使用延迟执行的技术，如lambda函数。

numpy = lambda: __import__('numpy')
# 当需要使用时才执行导入
array = numpy().array([1, 2, 3])

装饰器

装饰器是另一种延迟执行技术的实现

def lazy_import(module_name):
    def wrapper():
        globals()[module_name] = __import__(module_name)
        return globals()[module_name]
    return wrapper
@lazy_import('numpy')
def numpy_function():
    return numpy.array([1, 2, 3])
# 当需要使用时才执行导入和函数调用
result = numpy_function()

LazyLoader

importlib.util 中的 LazyLoader 可以懒加载。

目录结构

.
├── a.py
└── impt.py

impt.py

import sys
import importlib
from importlib.util import LazyLoader


for i, mp in enumerate(sys.meta_path):
    if str(mp) == "<class '_frozen_importlib_external.PathFinder'>":
        index = i
        path_finder = mp
        continue


class LazyPathFinder(object):
    
    def find_spec(self, fullname, path, target=None):
        spec = path_finder.find_spec(fullname, path, target=target)
        if spec is not None:
            spec.loader = LazyLoader(spec.loader)
        return spec


sys.meta_path[index] = LazyPathFinder()

import os
import a

print('a lazy imported')
print(a.b)

a.py

print('exec a.py')

b = 2

执行并查看输出结果

$ python impt.py 
a lazy imported
exec a.py
2

参考资料

• https://cloud.baidu.com/article/3296561
• https://blog.csdn.net/weixin_43590796/article/details/134945144
• https://www.51cto.com/article/720765.html
• https://www.cnblogs.com/lyg-blog/p/14806886.html

文章链接：
https://www.zywvvd.com/notes/coding/python/python-lazy-load/python-lazy-load/