Gently import a python module, failing only at runtime if a downstream module object invokes __call__

gentleimport.py

"""
gentleimport is a context manager that patches
__import__ to allow for softly importing mocked modules. To use gentleimport,
with gentleimport():
    import typing
    import asdf

typing.Any  # this is a real object as typing was a found module
asdf  # this is a mock module as asdf was not a real module

asdf() # (or asdf.<anything>....())
# will raise a ModuleNotFoundError to catch late usage of the library at runtime
"""
import builtins
import contextlib
import sys
import unittest.mock as mock


__all__ = ['gentleimport']


class ModuleMock(mock.MagicMock):
    """ module mock makes .<attribute> access have the same
    side effect as the parent (which in this module is an error
    for the invalid module)
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __getattr__(self, name):
        m = super().__getattr__(name)
        m.side_effect = self.side_effect
        return m


def _module_mock(name: str) -> mock.Mock:
    side_effect = ModuleNotFoundError(
        f'{name} was gently imported and could not be found - please install it'
    )
    return ModuleMock(name=name, side_effect=side_effect)


@contextlib.contextmanager
def gentleimport():
    """gentleimport is a context manager that allows for flexible
    importing of modules that can not be found so that import exceptions
    can be raised upon *module use* and not on import for soft dependencies

    Example::

        with gentleimport():
            import numpy as np
            import asdf

        np.ndarray  # (valid if numpy is installed)
        asdf.test # raises ModuleNotFoundError
    """
    orig_import = builtins.__import__
    def import_mock(name, *args):
        try:
            return orig_import(name, *args)
        except ModuleNotFoundError:
            module_mock = _module_mock(name)
            sys.modules[name] = module_mock
            return module_mock

    with mock.patch('builtins.__import__', side_effect=import_mock):
        yield
    builtins.__import__ = orig_import