From 4f34712858b24f4153d512a988827a78340d3e33 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Fri, 1 Jan 2021 03:51:13 +0100
Subject: [PATCH] Fix signatures for TypeVar and cast, fixes #1709

---
 jedi/inference/gradual/base.py        |  3 +++
 jedi/inference/gradual/type_var.py    |  8 ++++----
 jedi/inference/gradual/typing.py      | 12 ++++++------
 test/test_inference/test_signature.py | 14 ++++++++++++--
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/jedi/inference/gradual/base.py b/jedi/inference/gradual/base.py
index fb4740d3..cd9533f3 100644
--- a/jedi/inference/gradual/base.py
+++ b/jedi/inference/gradual/base.py
@@ -380,6 +380,9 @@ class BaseTypingValue(LazyValueWrapper):
     def _get_wrapped_value(self):
         return _PseudoTreeNameClass(self.parent_context, self._tree_name)
 
+    def get_signatures(self):
+        return self._wrapped_value.get_signatures()
+
     def __repr__(self):
         return '%s(%s)' % (self.__class__.__name__, self._tree_name.value)
 
diff --git a/jedi/inference/gradual/type_var.py b/jedi/inference/gradual/type_var.py
index acc6f46d..c09773f1 100644
--- a/jedi/inference/gradual/type_var.py
+++ b/jedi/inference/gradual/type_var.py
@@ -3,7 +3,7 @@ from jedi.inference.base_value import ValueSet, NO_VALUES, ValueWrapper
 from jedi.inference.gradual.base import BaseTypingValue
 
 
-class TypeVarClass(BaseTypingValue):
+class TypeVarClass(ValueWrapper):
     def py__call__(self, arguments):
         unpacked = arguments.unpack()
 
@@ -17,9 +17,9 @@ class TypeVarClass(BaseTypingValue):
         return ValueSet([TypeVar.create_cached(
             self.inference_state,
             self.parent_context,
-            self._tree_name,
-            var_name,
-            unpacked
+            tree_name=self.tree_node.name,
+            var_name=var_name,
+            unpacked_args=unpacked,
         )])
 
     def _find_string_name(self, lazy_value):
diff --git a/jedi/inference/gradual/typing.py b/jedi/inference/gradual/typing.py
index e49a7faf..634985fd 100644
--- a/jedi/inference/gradual/typing.py
+++ b/jedi/inference/gradual/typing.py
@@ -17,7 +17,7 @@ from jedi.inference.filters import FilterWrapper
 from jedi.inference.names import NameWrapper, ValueName
 from jedi.inference.value.klass import ClassMixin
 from jedi.inference.gradual.base import BaseTypingValue, \
-    BaseTypingClassWithGenerics, BaseTypingInstance
+    BaseTypingClassWithGenerics, BaseTypingInstance, ValueWrapper
 from jedi.inference.gradual.type_var import TypeVarClass
 from jedi.inference.gradual.generics import LazyGenericManager, TupleGenericManager
 
@@ -63,8 +63,8 @@ class TypingModuleName(NameWrapper):
             # have any effects there (because it's never executed).
             return
         elif name == 'TypeVar':
-            yield TypeVarClass.create_cached(
-                inference_state, self.parent_context, self.tree_name)
+            cls, = self._wrapped_name.infer()
+            yield TypeVarClass.create_cached(inference_state, cls)
         elif name == 'Any':
             yield AnyClass.create_cached(
                 inference_state, self.parent_context, self.tree_name)
@@ -79,8 +79,8 @@ class TypingModuleName(NameWrapper):
             yield NewTypeFunction.create_cached(
                 inference_state, self.parent_context, self.tree_name)
         elif name == 'cast':
-            yield CastFunction.create_cached(
-                inference_state, self.parent_context, self.tree_name)
+            cast_fn, = self._wrapped_name.infer()
+            yield CastFunction.create_cached(inference_state, cast_fn)
         elif name == 'TypedDict':
             # TODO doesn't even exist in typeshed/typing.py, yet. But will be
             # added soon.
@@ -430,7 +430,7 @@ class NewType(Value):
         return CompiledValueName(self, 'NewType')
 
 
-class CastFunction(BaseTypingValue):
+class CastFunction(ValueWrapper):
     @repack_with_argument_clinic('type, object, /')
     def py__call__(self, type_value_set, object_value_set):
         return type_value_set.execute_annotation()
diff --git a/test/test_inference/test_signature.py b/test/test_inference/test_signature.py
index a07bfa01..f6004b0c 100644
--- a/test/test_inference/test_signature.py
+++ b/test/test_inference/test_signature.py
@@ -104,6 +104,14 @@ class X:
         (partialmethod_code + 'X.d(', None),
 
         ('import contextlib\n@contextlib.contextmanager\ndef f(x): pass\nf(', 'f(x)'),
+
+        ('from typing import cast\ncast(', {
+            'cast(typ: object, val: Any) -> Any',
+            'cast(typ: str, val: Any) -> Any',
+            'cast(typ: Type[_T], val: Any) -> _T'}),
+        ('from typing import TypeVar\nTypeVar(',
+         'TypeVar(name: str, *constraints: Type[Any], bound: Union[None, Type[Any], str]=..., '
+         'covariant: bool=..., contravariant: bool=...)'),
     ]
 )
 def test_tree_signature(Script, environment, code, expected):
@@ -114,8 +122,10 @@ def test_tree_signature(Script, environment, code, expected):
     if expected is None:
         assert not Script(code).get_signatures()
     else:
-        sig, = Script(code).get_signatures()
-        assert expected == sig.to_string()
+        actual = {sig.to_string() for sig in Script(code).get_signatures()}
+        if not isinstance(expected, set):
+            expected = {expected}
+        assert expected == actual
 
 
 @pytest.mark.parametrize(