Make it possible to gather annotation classes for Union and Optional

jedi/evaluate/base_context.py

@@ -32,6 +32,9 @@ class HelperContextMixin:
     def execute_annotation(self):
         return self.execute_evaluated()
 
+    def gather_annotation_classes(self):
+        return ContextSet([self])
+
     def merge_types_of_iterate(self, contextualized_node=None, is_async=False):
         return ContextSet.from_sets(
             lazy_context.infer()
@@ -353,6 +356,9 @@ class ContextSet(BaseContextSet):
                 context_set |= method()
         return context_set
 
+    def gather_annotation_classes(self):
+        return ContextSet.from_sets([c.gather_annotation_classes() for c in self._set])
+
     def get_signatures(self):
         return [sig for c in self._set for sig in c.get_signatures()]
 

jedi/evaluate/context/typing.py

@@ -180,6 +180,11 @@ class TypingContextWithIndex(_WithIndexBase):
             self._context_of_index
         )])
 
+    def gather_annotation_classes(self):
+        return ContextSet.from_sets(
+            _iter_over_arguments(self._index_context, self._context_of_index)
+        )
+
 
 class TypingContext(_BaseTypingContext):
     index_class = TypingContextWithIndex

jedi/evaluate/param.py

@@ -52,7 +52,7 @@ class ExecutedParam(object):
             return True
         matches = any(c1.is_sub_class_of(c2)
                       for c1 in argument_contexts
-                      for c2 in annotations)
+                      for c2 in annotations.gather_annotation_classes())
         debug.dbg("signature compare %s: %s <=> %s",
                   matches, argument_contexts, annotations, color='BLUE')
         return matches