Make infer public on classes

jedi/api/classes.py

@@ -320,8 +320,7 @@ class BaseDefinition(object):
         names = self._evaluator.goto(self._name.parent_context, self._name.tree_name)
         return [Definition(self._evaluator, n) for n in names]
 
-    def _goto_definitions(self):
-        # TODO make this function public.
+    def infer(self):
         return [Definition(self._evaluator, d.name) for d in self._name.infer()]
 
     @property

jedi/evaluate/filters.py

@@ -72,7 +72,7 @@ class ContextNameMixin(object):
         return ContextSet([self._context])
 
     def get_root_context(self):
-        if self.parent_context is None:
+        if self.parent_context is None:  # A module
             return self._context
         return super(ContextNameMixin, self).get_root_context()
 

jedi/evaluate/gradual/stub_context.py

@@ -214,7 +214,8 @@ class StubName(NameWrapper):
             return self._wrapped_name.infer()
 
         typ = self._wrapped_name.tree_name.parent.type
-        # TODO is this if a performance optimization?
+        # Only for these two we want to merge, the function doesn't support
+        # anything else.
         if typ in ('classdef', 'funcdef'):
             actual_context, = self._wrapped_name.infer()
             return _add_stub_if_possible(self.parent_context, actual_context, stub_contexts)
@@ -244,6 +245,7 @@ def _add_stub_if_possible(parent_context, actual_context, stub_contexts):
 
 
 def with_stub_context_if_possible(actual_context):
+    assert actual_context.tree_node.type in ('classdef', 'funcdef')
     names = actual_context.get_qualified_names()
     stub_module = actual_context.get_root_context().stub_context
     if stub_module is None: