jedi-fork/jedi/evaluate/gradual/stub_context.py

import os

from jedi.evaluate.base_context import ContextWrapper, ContextSet, \
    NO_CONTEXTS
from jedi.evaluate.context.klass import ClassMixin, ClassContext
from jedi.evaluate.context.module import ModuleContext
from jedi.evaluate.filters import ParserTreeFilter, \
    TreeNameDefinition
from jedi.evaluate.utils import to_list
from jedi.evaluate.gradual.typing import TypingModuleFilterWrapper, AnnotatedClass


class _StubContextMixin(object):
    _add_non_stubs_in_filter = False

    def is_stub(self):
        return True

    def _get_stub_only_filters(self, **filter_kwargs):
        return [StubFilter(
            self.evaluator,
            context=self,
            **filter_kwargs
        )]

    def _get_base_filters(self, filters, search_global=False,
                          until_position=None, origin_scope=None):
        next(filters)  # Ignore the first filter and replace it with our own
        stub_only_filters = self._get_stub_only_filters(
            search_global=search_global,
            until_position=until_position,
            origin_scope=origin_scope,
        )
        for f in stub_only_filters:
            yield f

        for f in filters:
            yield f


class StubModuleContext(_StubContextMixin, ModuleContext):
    _add_non_stubs_in_filter = True

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

    def _get_first_non_stub_filters(self):
        for context in self.non_stub_context_set:
            yield next(context.get_filters(search_global=False))

    def _get_stub_only_filters(self, search_global, **filter_kwargs):
        stub_filters = super(StubModuleContext, self)._get_stub_only_filters(
            search_global=search_global, **filter_kwargs
        )
        stub_filters += self.iter_star_filters(search_global=search_global)
        return stub_filters

    def get_filters(self, search_global=False, until_position=None,
                    origin_scope=None, **kwargs):
        filters = super(StubModuleContext, self).get_filters(
            search_global, until_position, origin_scope, **kwargs
        )
        for f in self._get_base_filters(filters, search_global, until_position, origin_scope):
            yield f

    def _iter_modules(self, path):
        dirs = os.listdir(path)
        for name in dirs:
            if os.path.isdir(os.path.join(path, name)):
                yield (None, name, True)
            if name.endswith('.pyi'):
                yield (None, name[:-4], True)
        return []


class StubClass(_StubContextMixin, ClassMixin, ContextWrapper):
    pass


class TypingModuleWrapper(StubModuleContext):
    def get_filters(self, *args, **kwargs):
        filters = super(TypingModuleWrapper, self).get_filters(*args, **kwargs)
        yield TypingModuleFilterWrapper(next(filters))
        for f in filters:
            yield f


def goto_with_stubs_if_possible(name):
    return [name]
    # XXX
    root = name.parent_context.get_root_context()
    stub = root.get_root_context().stub_context
    if stub is None:
        return [name]

    qualified_names = name.parent_context.get_qualified_names()
    if qualified_names is None:
        return [name]

    stub_contexts = ContextSet([stub])
    for n in qualified_names:
        stub_contexts = stub_contexts.py__getattribute__(n)
    names = stub_contexts.py__getattribute__(name.string_name, is_goto=True)
    return [
        n
        for n in names
        if n.start_pos == name.start_pos and n.parent_context == name.parent_context
    ] or [name]


def goto_non_stub(parent_context, tree_name):
    contexts = stub_to_actual_context_set(parent_context)
    return contexts.py__getattribute__(tree_name, is_goto=True)


def stub_to_actual_context_set(stub_context, ignore_compiled=False):
    stub_module = stub_context.get_root_context()
    if not stub_module.is_stub():
        return ContextSet([stub_context])

    qualified_names = stub_context.get_qualified_names()
    if qualified_names is None:
        return NO_CONTEXTS

    assert isinstance(stub_module, StubModuleContext), stub_module
    non_stubs = stub_module.non_stub_context_set
    if ignore_compiled:
        non_stubs = non_stubs.filter(lambda c: not c.is_compiled())
    for name in qualified_names:
        non_stubs = non_stubs.py__getattribute__(name)
    return non_stubs


def try_stubs_to_actual_context_set(stub_contexts, prefer_stub_to_compiled=False):
    return ContextSet.from_sets(
        stub_to_actual_context_set(stub_context, ignore_compiled=prefer_stub_to_compiled)
        or ContextSet([stub_context])
        for stub_context in stub_contexts
    )


@to_list
def try_stub_to_actual_names(names, prefer_stub_to_compiled=False):
    for name in names:
        # Using the tree_name is better, if it's available, becuase no
        # information is lost. If the name given is defineda as `foo: int` we
        # would otherwise land on int, which is not what we want. We want foo
        # from the non-stub module.
        if name.tree_name is None:
            actual_contexts = ContextSet.from_sets(
                stub_to_actual_context_set(c) for c in name.infer()
            )
            actual_contexts = actual_contexts.filter(lambda c: not c.is_compiled())
            if actual_contexts:
                for s in actual_contexts:
                    yield s.name
            else:
                yield name
        else:
            parent_context = name.parent_context
            if not parent_context.is_stub():
                yield name
                continue

            contexts = stub_to_actual_context_set(parent_context)
            if prefer_stub_to_compiled:
                # We don't really care about
                contexts = contexts.filter(lambda c: not c.is_compiled())

            new_names = contexts.py__getattribute__(name.tree_name, is_goto=True)

            if new_names:
                for n in new_names:
                    yield n
            else:
                yield name


def _load_or_get_stub_module(evaluator, names):
    return evaluator.stub_module_cache.get(names)


def load_stubs(context):
    root_context = context.get_root_context()
    stub_module = _load_or_get_stub_module(
        context.evaluator,
        root_context.string_names
    )
    if stub_module is None:
        return NO_CONTEXTS

    qualified_names = context.get_qualified_names()
    if qualified_names is None:
        return NO_CONTEXTS

    stub_contexts = ContextSet([stub_module])
    for name in qualified_names:
        stub_contexts = stub_contexts.py__getattribute__(name)

    if isinstance(context, AnnotatedClass):
        return ContextSet([
            context.annotate_other_class(c) if c.is_class() else c
            for c in stub_contexts
        ])
    return stub_contexts


def _load_stub_module(module):
    if module.is_stub():
        return module
    from jedi.evaluate.gradual.typeshed import _try_to_load_stub
    return _try_to_load_stub(
        module.evaluator,
        ContextSet([module]),
        parent_module_context=None,
        import_names=module.string_names
    )


def name_to_stub(name):
    return ContextSet.from_sets(to_stub(c) for c in name.infer())


def to_stub(context):
    if context.is_stub():
        return ContextSet([context])

    qualified_names = context.get_qualified_names()
    stub_module = _load_stub_module(context.get_root_context())
    if stub_module is None or qualified_names is None:
        return NO_CONTEXTS

    stub_contexts = ContextSet([stub_module])
    for name in qualified_names:
        stub_contexts = stub_contexts.py__getattribute__(name)
    return stub_contexts


class _StubName(TreeNameDefinition):
    def infer(self):
        inferred = super(_StubName, self).infer()
        if self.string_name == 'version_info' and self.get_root_context().py__name__() == 'sys':
            return [VersionInfo(c) for c in inferred]

        return [
            StubClass.create_cached(c.evaluator, c) if isinstance(c, ClassContext) else c
            for c in inferred
        ]


class StubFilter(ParserTreeFilter):
    name_class = _StubName

    def __init__(self, *args, **kwargs):
        self._search_global = kwargs.pop('search_global')  # Python 2 :/
        super(StubFilter, self).__init__(*args, **kwargs)

    def _is_name_reachable(self, name):
        if not super(StubFilter, self)._is_name_reachable(name):
            return False

        if not self._search_global:
            # Imports in stub files are only public if they have an "as"
            # export.
            definition = name.get_definition()
            if definition.type in ('import_from', 'import_name'):
                if name.parent.type not in ('import_as_name', 'dotted_as_name'):
                    return False
            n = name.value
            if n.startswith('_') and not (n.startswith('__') and n.endswith('__')):
                return False
        return True


class VersionInfo(ContextWrapper):
    pass