diff --git a/jedi/evaluate/gradual/annotation.py b/jedi/evaluate/gradual/annotation.py
index c084f16c..27f6c00f 100644
--- a/jedi/evaluate/gradual/annotation.py
+++ b/jedi/evaluate/gradual/annotation.py
@@ -205,7 +205,7 @@ def infer_return_types(function_execution_context):
 
     return ContextSet.from_sets(
         ann.define_generics(type_var_dict)
-        if isinstance(ann, (AbstractAnnotatedClass, TypeVar)) else ann
+        if isinstance(ann, (AbstractAnnotatedClass, TypeVar)) else ContextSet({ann})
         for ann in annotation_contexts
     ).execute_annotation()
 
diff --git a/test/test_evaluate/test_stdlib.py b/test/test_evaluate/test_stdlib.py
index 58898ead..fb98d8a4 100644
--- a/test/test_evaluate/test_stdlib.py
+++ b/test/test_evaluate/test_stdlib.py
@@ -91,9 +91,9 @@ def test_re_sub(Script, environment):
 
     names = run("import re; re.sub('a', 'a', 'f')")
     if environment.version_info.major == 2:
-        assert names == {'str', 'unicode'}
+        assert names == {'str'}
     else:
-        assert names == {'str', 'bytes'}
+        assert names == {'str'}
 
     # This param is missing because of overloading.
     names = run("import re; re.sub('a', 'a')")