Source code for schrodinger.utils.imputils

"""
Utility functions to import python modules. Always prefer loading from a
location on sys.path.

Copyright Schrodinger, LLC. All rights reserved.

"""
# Contributors: Matvey Adzhigirey

import importlib
import importlib.util
import inspect
import os
import sys

from schrodinger.utils import fileutils
from schrodinger.utils import imputils


def import_module_from_file(filename):
    """
    Import module from a file path.  Returns the imported module object.
    """

    modulename = os.path.splitext(os.path.basename(filename))[0]
    return import_from_file(modulename, filename)


def import_from_file(modulename, filename):
    if modulename in sys.modules:
        return sys.modules[modulename]
    spec = importlib.util.spec_from_file_location(modulename, filename)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    sys.modules[modulename] = module
    return module


def get_path_from_module(module):
    """
    Given a module, return a path string. The path string will either be in '.'
    format if the module can be found in either the schrodinger scripts or
    modules directory or an absolute path if the module is defined in a python
    unittest module.  Inverse function of `get_module_from_path`.

    .. WARNING

        Contains special handling for unit test modules. Unit test modules will
        be functionally equivalent to the original module, but may not be
        identical, since it is being re-imported separately.

    :param module: The module to export a path string for. To get the module
        of any arbitrary object, use `inspect.getmodule`.
    :type module: module
    """
    path = inspect.getfile(module)
    scripts_dir = fileutils.get_mmshare_scripts_dir() + os.path.sep
    if path.startswith(scripts_dir):
        path = path.replace(scripts_dir, '')
        path = path.replace('.py', '')
        return path.replace(os.path.sep, '.')
    elif os.path.join('python', 'test') in path:
        return os.path.abspath(path)
    else:
        return module.__name__


def get_module_from_path(module_path):
    """
    Given a module path string generated by `get_path_from_module`, return the
    corresponding module. Inverse function of `get_path_from_module`.

    :param module_path: The path string describing the module to import
    :type module_path: str
    """
    if (os.path.join('python', 'test') in module_path):
        return imputils.import_module_from_file(module_path)
    else:
        return importlib.import_module(module_path)


def import_script(name, subdir=None, common=False):
    """
    Import the given script residing in mmshare/python/scripts as a module. If
    the script is not found in the scripts directory, an attempt is made to
    import the script from a subdirectory following our standard naming
    convention by replacing '_driver.py' or '_backend.py' with '_gui_dir'.

    :param str name: The name of the script, including the .py extension

    :param str subdir: The name of the subdirectory the script resides in - must
        be a path relative to mmshare/python/scripts. If given, this name will
        be used instead of attempting to derive a subdirectory name from the
        script name.

    :param bool common: Import from python/common rather than python/scripts

    :rtype: module
    :return: The script imported as a module

    :raise FileNotFoundError: If the script can't be found
    """

    if common:
        script_dir = fileutils.get_mmshare_common_dir()
    else:
        script_dir = fileutils.get_mmshare_scripts_dir()
    if subdir:
        path = os.path.join(script_dir, subdir, name)
    else:
        path = os.path.join(script_dir, name)
        if not os.path.exists(path):
            guidir = '_gui_dir'
            subdir = name.replace('_driver.py', guidir)
            subdir = subdir.replace('_backend.py', guidir)
            path = os.path.join(script_dir, subdir, name)
    return import_module_from_file(path)


def lazy_import(name):
    """
    Lazily import a module.  The actual import will not happen until the module
    is first used.  This can help to avoid performance bottlenecks (i.e. if a
    module is rarely used but the import slows down panel launching) or to avoid
    circular imports (although refactoring your code to avoid the circular
    import in the first place may be a better option).

    For example::

        from schrodinger.utils import imputils
        phase_markers = imputils.lazy_import("schrodinger.application.phase.phase_markers")

    `phase_markers` can now be used as if the line had said::

        import schrodinger.application.phase.phase_markers as phase_markers

    but the actual import won't occur until the module is first used (or if the
    module is imported normally somewhere else).  Note that imports of higher
    level packages (i.e. the "schrodinger", "schrodinger.application", and
    "schrodinger.application.phase" packages for the example above) will not
    happen lazily.  If any of those packages contain an `__init__.py`, then the
    `__init__.py` will be executed during the `lazy_import` call.

    :param name: The name of the module to lazily import.  This must be given as
        an absolute name, not a relative one (i.e.
        "schrodinger.application.phase.phase_markers", not ".phase_markers")
    :type name: str

    :return: A placeholder for the module
    :rtype: ModuleType
    """
    # this code is based on
    # https://github.com/python/cpython/blob/master/Doc/library/importlib.rst#implementing-lazy-imports
    if name in sys.modules:
        # the module has already been imported, so just return it
        return sys.modules[name]
    spec = importlib.util.find_spec(name)
    if spec is None:
        raise ModuleNotFoundError(f"No module named {name}")
    loader = importlib.util.LazyLoader(spec.loader)
    spec.loader = loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    loader.exec_module(module)

    # See SHARED-7537. For 'package.module', module must be added as an
    # attribute to package manually.  Note that package is automatically
    # imported when importlib.util.find_spec('package.module') is executed.
    if '.' in spec.name:
        _, module_name = spec.name.rsplit('.', 1)
        if (spec.parent in sys.modules and
                not hasattr(sys.modules[spec.parent], module_name)):
            setattr(sys.modules[spec.parent], module_name, module)
    # End of fix for SHARED-7537 (Which is Python bug 42273)

    return module