PY-55044 TypedDict **kwargs typing

GitOrigin-RevId: 5dff382b321ac166ed4b8cf90e431ea719995608

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -112,6 +112,8 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
   public static final String REQUIRED_EXT = "typing_extensions.Required";
   public static final String NOT_REQUIRED = "typing.NotRequired";
   public static final String NOT_REQUIRED_EXT = "typing_extensions.NotRequired";
+  private static final String UNPACK = "typing.Unpack";
+  private static final String UNPACK_EXT = "typing_extensions.Unpack";
 
   public static final String SELF = "typing.Self";
   public static final String SELF_EXT = "typing_extensions.Self";
@@ -257,6 +259,12 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
     if (typeHint == null) {
       return null;
     }
+    if (param.isKeywordContainer()) {
+      Ref<PyType> type = getTypeFromUnpackOperator(typeHint, context.myContext);
+      if (type != null) {
+        return type.get() instanceof PyTypedDictType ? type : null;
+      }
+    }
     if (typeHint instanceof PyReferenceExpression ref && ref.isQualified() && (
       param.isPositionalContainer() && "args".equals(ref.getReferencedName()) ||
       param.isKeywordContainer() && "kwargs".equals(ref.getReferencedName())
@@ -1626,6 +1634,17 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
     return null;
   }
 
+  @Nullable
+  private static Ref<@Nullable PyType> getTypeFromUnpackOperator(@NotNull PsiElement element, @NotNull TypeEvalContext context) {
+    if (!(element instanceof PySubscriptionExpression subscriptionExpr)) return null;
+    if (!resolvesToQualifiedNames(subscriptionExpr.getOperand(), context, UNPACK, UNPACK_EXT)) {
+      return null;
+    }
+    PyExpression indexExpression = subscriptionExpr.getIndexExpression();
+    if (indexExpression == null) return null;
+    return Ref.create(Ref.deref(getType(indexExpression, context)));
+  }
+
   @Nullable
   private static PyParamSpecType getParamSpecType(@NotNull PsiElement element, @NotNull Context context) {
     if (!(element instanceof PyCallExpression assignedCall)) return null;

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java

@@ -363,7 +363,7 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
         return type;
       }
     }
-    return new PyFunctionTypeImpl(this);
+    return PyFunctionTypeImpl.create(this, context);
   }
 
   @Override

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java

@@ -36,7 +36,7 @@ public class PyLambdaExpressionImpl extends PyElementImpl implements PyLambdaExp
         return type;
       }
     }
-    return new PyFunctionTypeImpl(this);
+    return PyFunctionTypeImpl.create(this, context);
   }
 
 

python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCallableParameterImpl.java

@@ -172,8 +172,7 @@ public final class PyCallableParameterImpl implements PyCallableParameter {
 
   @Override
   public boolean isSelf() {
-    final PyParameter parameter = PyUtil.as(myElement, PyParameter.class);
-    return parameter != null && parameter.isSelf();
+    return myElement != null && myElement.isSelf();
   }
 
   @NotNull

python/python-psi-impl/src/com/jetbrains/python/psi/types/PyFunctionTypeImpl.java

@@ -12,6 +12,7 @@ import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -20,12 +21,31 @@ import static com.jetbrains.python.psi.PyUtil.as;
 
 /**
  * Type of a particular function that is represented as a {@link PyCallable} in the PSI tree.
- *
  */
 public class PyFunctionTypeImpl implements PyFunctionType {
+  public static PyFunctionType create(@NotNull PyCallable callable, @NotNull TypeEvalContext context) {
+    List<PyCallableParameter> parameters = new ArrayList<>();
+    for (PyParameter parameter : callable.getParameterList().getParameters()) {
+      if (parameter instanceof PyNamedParameter namedParameter &&
+          namedParameter.isKeywordContainer() &&
+          context.getType(namedParameter) instanceof PyTypedDictType typedDictType) {
+        List<PyCallableParameter> typedDictParameters = typedDictType.toClass().getParameters(context);
+        parameters.addAll(ContainerUtil.notNullize(typedDictParameters));
+      }
+      else {
+        parameters.add(PyCallableParameterImpl.psi(parameter));
+      }
+    }
+    return new PyFunctionTypeImpl(callable, parameters);
+  }
+
   @NotNull private final PyCallable myCallable;
   @NotNull private final List<PyCallableParameter> myParameters;
 
+  /**
+   * @deprecated Use {@link PyFunctionTypeImpl#create(PyCallable, TypeEvalContext)}
+   */
+  @Deprecated
   public PyFunctionTypeImpl(@NotNull PyCallable callable) {
     this(callable, ContainerUtil.map(callable.getParameterList().getParameters(), PyCallableParameterImpl::psi));
   }