From e2fea0a5de9da5a743faefb11fba9812eb3569f8 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sat, 16 Mar 2019 01:06:13 +0100
Subject: [PATCH] Fix some tests because of stub_context changes

---
 jedi/evaluate/context/function.py             |  5 ++++-
 jedi/evaluate/context/instance.py             |  7 ++-----
 jedi/evaluate/gradual/typing.py               | 21 ++++++++-----------
 .../test_gradual/test_typeshed.py             |  2 +-
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/jedi/evaluate/context/function.py b/jedi/evaluate/context/function.py
index 8154d599..91affa41 100644
--- a/jedi/evaluate/context/function.py
+++ b/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:
diff --git a/jedi/evaluate/context/instance.py b/jedi/evaluate/context/instance.py
index c9d57359..d135facc 100644
--- a/jedi/evaluate/context/instance.py
+++ b/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:
+        stub_context = self._wrapped_context.stub_context
+        if stub_context is not None:
             return stub_context.py__call__(arguments)
 
         function_execution = self.get_function_execution(arguments)
diff --git a/jedi/evaluate/gradual/typing.py b/jedi/evaluate/gradual/typing.py
index 69456a67..5b1c0187 100644
--- a/jedi/evaluate/gradual/typing.py
+++ b/jedi/evaluate/gradual/typing.py
@@ -513,19 +513,16 @@ 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(
-                # Take the first filter, which is here to filter module contents
-                # and wrap it.
-                self.parent_context,
-                [filter_],
-                search_global=False,
-                origin_scope=origin_scope,
-            )
+        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,
+            [filter_],
+            search_global=False,
+            origin_scope=origin_scope,
+        )
 
     def get_filters(self, search_global=False, *args, **kwargs):
         filters = super(AbstractAnnotatedClass, self).get_filters(
diff --git a/test/test_evaluate/test_gradual/test_typeshed.py b/test/test_evaluate/test_gradual/test_typeshed.py
index 08c74d21..d9019d06 100644
--- a/test/test_evaluate/test_gradual/test_typeshed.py
+++ b/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.