Skip to content
Back to top

# Source code for xonsh.completers.imports

    """
    Import statement completions.
    Contains modified code from the IPython project (at core/completerlib.py).

    # Copyright (c) IPython Development Team.
    # Distributed under the terms of the Modified BSD License.
    """

    import inspect
    import os
    import re
    import sys
    from importlib import import_module
    from importlib.machinery import all_suffixes
    from time import time
    from zipimport import zipimporter

    from xonsh.built_ins import XSH
    from xonsh.completers.tools import (
        RichCompletion,
        contextual_completer,
        get_filter_function,
    )
    from xonsh.lib.lazyasd import lazyobject
    from xonsh.parsers.completion_context import CompletionContext

    _suffixes = all_suffixes()

    # Time in seconds after which we give up
    TIMEOUT_GIVEUP = 2

    @lazyobject
    def IMPORT_RE():
        # Regular expression for the python import statement
        suffixes = r"|".join(re.escape(s) for s in _suffixes)
        return re.compile(
            r"(?P<name>[^\W\d]\w*?)"
            r"(?P<package>[/\\]__init__)?"
            rf"(?P<suffix>{suffixes})$"
        )

    def module_list(path):
        """
        Return the list containing the names of the modules available in the given
        folder.
        """
        # sys.path has the cwd as an empty string, but isdir/listdir need it as '.'
        if path == "":
            path = "."

        # A few local constants to be used in loops below
        pjoin = os.path.join

        if os.path.isdir(path):
            # Build a list of all files in the directory and all files
            # in its subdirectories. For performance reasons, do not
            # recurse more than one level into subdirectories.
            files = []
            for root, dirs, nondirs in os.walk(path, followlinks=True):
                subdir = root[len(path) + 1 :]
                if subdir:
                    files.extend(pjoin(subdir, f) for f in nondirs)
                    dirs[:] = []  # Do not recurse into additional subdirectories.
                else:
                    files.extend(nondirs)

        else:
            try:
                files = list(zipimporter(path)._files.keys())
            except (ImportError, AttributeError, TypeError, OSError):
                # Not bare except: let KeyboardInterrupt propagate
                files = []

        # Build a list of modules which match the import_re regex.
        modules = []
        for f in files:
            m = IMPORT_RE.match(f)
            if m:
                modules.append(m.group("name"))
        return list(set(modules))

    def get_root_modules():
        """
        Returns a list containing the names of all the modules available in the
        folders of the pythonpath.
        """
        rootmodules_cache = XSH.modules_cache
        rootmodules = list(sys.builtin_module_names)
        start_time = time()
        for path in sys.path:
            try:
                modules = rootmodules_cache[path]
            except KeyError:
                modules = module_list(path)
                try:
                    modules.remove("__init__")
                except ValueError:
                    pass
                if path not in ("", "."):  # cwd modules should not be cached
                    rootmodules_cache[path] = modules
                if time() - start_time > TIMEOUT_GIVEUP:
                    print("\nwarning: Getting root modules is taking too long, we give up")
                    return []
            rootmodules.extend(modules)
        rootmodules = list(set(rootmodules))
        return rootmodules

    def is_importable(module, attr, only_modules):
        if only_modules:
            return inspect.ismodule(getattr(module, attr))
        else:
            return not (attr[:2] == "__" and attr[-2:] == "__")

    def is_possible_submodule(module, attr):
        try:
            obj = getattr(module, attr)
        except AttributeError:
            # Is possilby an unimported submodule
            return True
        except TypeError:
            # https://github.com/ipython/ipython/issues/9678
            return False
        return inspect.ismodule(obj)

    def try_import(mod: str, only_modules=False) -> list[str]:
        """
        Try to import given module and return list of potential completions.
        """
        mod = mod.rstrip(".")
        try:
            m = import_module(mod)
        except Exception:
            return []

        m_is_init = "__init__" in (getattr(m, "__file__", "") or "")

        completions = []
        if (not hasattr(m, "__file__")) or (not only_modules) or m_is_init:
            completions.extend(
                [attr for attr in dir(m) if is_importable(m, attr, only_modules)]
            )

        m_all = getattr(m, "__all__", [])
        if only_modules:
            completions.extend(attr for attr in m_all if is_possible_submodule(m, attr))
        else:
            completions.extend(m_all)

        if m_is_init:
            if m.__file__:
                completions.extend(module_list(os.path.dirname(m.__file__)))
        completions_set = {c for c in completions if isinstance(c, str)}
        completions_set.discard("__init__")
        return list(completions_set)

    ###############
    # Xonsh code: #
    ###############

    def filter_completions(prefix, completions):
        filt = get_filter_function()
        for comp in completions:
            if filt(comp, prefix):
                yield comp

    @contextual_completer
    def complete_import(context: CompletionContext):
        """
        Completes module names and objects for "import ..." and "from ... import
        ...".
        """
        if not (context.command and context.python):
            # Imports are only possible in independent lines (not in `$()` or `@()`).
            # This means it's python code, but also can be a command as far as the parser is concerned.
            return None

        command = context.command

        if command.opening_quote:
            # can't have a quoted import
            return None

        arg_index = command.arg_index
        prefix = command.prefix
        args = command.args

        if arg_index == 1 and args[0].value == "from":
            # completing module to import
            return complete_module(prefix)
        if arg_index >= 1 and args[0].value == "import":
            # completing module to import, might be multiple modules
            prefix = prefix.rsplit(",", 1)[-1]
            return complete_module(prefix), len(prefix)
        if arg_index == 2 and args[0].value == "from":
            return {RichCompletion("import", append_space=True)}
        if arg_index > 2 and args[0].value == "from" and args[2].value == "import":
            # complete thing inside a module, might be multiple objects
            module = args[1].value
            prefix = prefix.rsplit(",", 1)[-1]
            return filter_completions(prefix, try_import(module)), len(prefix)
        return set()

    def complete_module(prefix):
        if not prefix:
            modules = get_root_modules()
        else:
            mod = prefix.split(".")
            if len(mod) < 2:
                modules = get_root_modules()
            else:
                completion_list = try_import(".".join(mod[:-1]), only_modules=True)
                modules = (".".join(mod[:-1] + [el]) for el in completion_list)

        yield from filter_completions(prefix, modules)