Fix some tests because of stub_context changes

jedi/evaluate/context/function.py

@@ -336,7 +336,10 @@ class FunctionExecutionContext(TreeContext):
                 # Only the first generic is relevant.
                 generics = (return_contexts.py__class__(), NO_CONTEXTS, NO_CONTEXTS)
                 return ContextSet(
-                    AnnotatedSubClass(getattr(c, 'stub_context', c), generics) for c in async_classes
+                    AnnotatedSubClass(
+                        c if c.stub_context is None else c.stub_context,
+                        generics
+                    ) for c in async_classes
                 ).execute_annotation()
         else:
             if is_generator:

jedi/evaluate/context/instance.py

@@ -395,11 +395,8 @@ class BoundMethod(FunctionMixin, ContextWrapper):
 
         # This might not be the most beautiful way, but prefer stub_contexts
         # and execute those if possible.
-        try:
         stub_context = self._wrapped_context.stub_context
-        except AttributeError:
-            pass
-        else:
+        if stub_context is not None:
             return stub_context.py__call__(arguments)
 
         function_execution = self.get_function_execution(arguments)

jedi/evaluate/gradual/typing.py

@@ -513,12 +513,9 @@ class AbstractAnnotatedClass(ClassMixin, ContextWrapper):
         filter_ = super(AbstractAnnotatedClass, self)._create_class_filter(
             cls, origin_scope, is_instance
         )
-        try:
-            stub_context = cls.stub_context
-        except AttributeError:
+        if cls.stub_context is None:
             return filter_
-        else:
-            return stub_context.get_stub_only_filter(
+        return cls.stub_context.get_stub_only_filter(
             # Take the first filter, which is here to filter module contents
             # and wrap it.
             self.parent_context,

test/test_evaluate/test_gradual/test_typeshed.py

@@ -190,7 +190,7 @@ def test_goto_stubs(Script):
         'import os; os.walk'
         'from collections import Counter; Counter'
     ])
-def test_goto_stubs_on_itself(Script):
+def test_goto_stubs_on_itself(Script, code):
     """
     If goto_stubs is used on an identifier in e.g. the stdlib, we should goto
     the stub of it.