[ty] Infer dict(**TypedDict) in TypedDict context

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/generics/pep695/aliases.md

@@ -520,6 +520,18 @@ def _():
     # dict() call with keyword args should be inferred as Dog
     animal: MaybeDog = dict(name="Buddy", breed="Labrador")
     reveal_type(animal)  # revealed: Dog
+
+def _(dog: Dog):
+    # dict() call with unpacked TypedDict kwargs should also be inferred as Dog
+    animal: MaybeDog = dict(**dog)
+    reveal_type(animal)  # revealed: Dog
+
+def returns_dog(dog: Dog) -> Dog:
+    return dict(**dog)
+
+def takes_dog(dog: Dog) -> None: ...
+def _(dog: Dog):
+    takes_dog(dict(**dog))
 ```
 
 And with set literal inference:

crates/ty_python_semantic/src/types/infer/builder/dict.rs

@@ -3,7 +3,7 @@ use ruff_python_ast::{self as ast, HasNodeIndex};
 use rustc_hash::FxHashMap;
 
 use super::{ArgExpr, TypeInferenceBuilder};
-use crate::types::typed_dict::validate_typed_dict_constructor;
+use crate::types::typed_dict::{supports_typed_dict_unpack, validate_typed_dict_constructor};
 use crate::types::{KnownClass, Type, TypeContext};
 
 impl<'db> TypeInferenceBuilder<'db, '_> {
@@ -13,44 +13,62 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         arguments: &ast::Arguments,
         call_expression_tcx: TypeContext<'db>,
     ) -> Option<Type<'db>> {
-        if !arguments.args.is_empty()
-            || arguments
-                .keywords
-                .iter()
-                .any(|keyword| keyword.arg.is_none())
-        {
+        if !arguments.args.is_empty() {
             return None;
         }
 
         // Fast-path dict(...) in TypedDict context: infer keyword values against fields,
-        // then validate and return the TypedDict type.
+        // then validate and return the TypedDict type. This also covers `dict(**src)` when `src`
+        // is `TypedDict`-shaped.
         if let Some(tcx) = call_expression_tcx.annotation
             && let Some(typed_dict) = tcx
                 .filter_union(self.db(), Type::is_typed_dict)
                 .as_typed_dict()
         {
-            let items = typed_dict.items(self.db());
+            let mut speculative_builder = self.speculate();
+            let items = typed_dict.items(speculative_builder.db());
+            let mut unpacked_keyword_types = Vec::with_capacity(arguments.keywords.len());
             for keyword in &arguments.keywords {
-                if let Some(arg_name) = &keyword.arg {
-                    let value_tcx = items
-                        .get(arg_name.id.as_str())
-                        .map(|field| TypeContext::new(Some(field.declared_ty)))
-                        .unwrap_or_default();
-                    self.infer_expression(&keyword.value, value_tcx);
-                }
+                let value_tcx = keyword
+                    .arg
+                    .as_ref()
+                    .and_then(|arg_name| items.get(arg_name.id.as_str()))
+                    .map(|field| TypeContext::new(Some(field.declared_ty)))
+                    .unwrap_or_default();
+                let value_ty = speculative_builder.infer_expression(&keyword.value, value_tcx);
+                unpacked_keyword_types.push(keyword.arg.is_none().then_some(value_ty));
+            }
+
+            let supports_typed_dict_context = unpacked_keyword_types
+                .iter()
+                .copied()
+                .flatten()
+                .all(|keyword_ty| supports_typed_dict_unpack(speculative_builder.db(), keyword_ty));
+
+            if !supports_typed_dict_context {
+                return None;
             }
 
             validate_typed_dict_constructor(
-                &self.context,
+                &speculative_builder.context,
                 typed_dict,
                 arguments,
                 func.into(),
-                |expr, _| self.expression_type(expr),
+                |expr, _| speculative_builder.expression_type(expr),
             );
+            self.extend(speculative_builder);
 
             return Some(Type::TypedDict(typed_dict));
         }
 
+        if arguments
+            .keywords
+            .iter()
+            .any(|keyword| keyword.arg.is_none())
+        {
+            return None;
+        }
+
         // Lower `dict(a=..., b=...)` to synthetic `(Literal["a"], value)` pairs so we can
         // reuse dict-literal inference. We key the synthetic name off the value node because
         // `infer_collection_literal` operates on expressions rather than keywords.

crates/ty_python_semantic/src/types/typed_dict.rs

@@ -958,6 +958,10 @@ fn extract_unpacked_typed_dict_keys<'db>(
     }
 }
 
+pub(super) fn supports_typed_dict_unpack<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {
+    ty.is_dynamic() || ty.is_never() || extract_unpacked_typed_dict_keys(db, ty).is_some()
+}
+
 /// Infers each unpacked `**kwargs` constructor argument exactly once.
 ///
 /// Mixed positional-and-keyword `TypedDict` construction needs to inspect unpacked keyword types