[pyrefly] Preserve literal keys in dict unions

pyrefly/lib/alt/operators.rs

@@ -19,6 +19,7 @@ use pyrefly_types::simplify::intersect;
 use pyrefly_types::type_var::Restriction;
 use pyrefly_util::prelude::VecExt;
 use ruff_python_ast::CmpOp;
+use ruff_python_ast::Expr;
 use ruff_python_ast::ExprBinOp;
 use ruff_python_ast::ExprCompare;
 use ruff_python_ast::ExprUnaryOp;
@@ -404,7 +405,14 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             rhs = self.expr_infer_with_hint(&x.right, hint, errors);
         } else {
             lhs = self.expr_infer(&x.left, errors);
-            rhs = self.expr_infer(&x.right, errors);
+            rhs = if x.op == Operator::BitOr
+                && matches!(&*x.right, Expr::Dict(_))
+                && matches!(&lhs, Type::ClassType(cls) if cls.class_object().is_builtin("dict"))
+            {
+                self.expr_infer_with_hint(&x.right, Some(HintRef::soft(&lhs)), errors)
+            } else {
+                self.expr_infer(&x.right, errors)
+            };
         }
 
         // Optimisation: If we have `Union[a, b] | Union[c, d]`, instead of unioning
@@ -679,7 +687,14 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             self.try_binop_calls(&calls_to_try, range, errors, &context)
         };
         let base = self.expr_infer(&x.target, errors);
-        let rhs = self.expr_infer(&x.value, errors);
+        let rhs = if x.op == Operator::BitOr
+            && matches!(&*x.value, Expr::Dict(_))
+            && matches!(&base, Type::ClassType(cls) if cls.class_object().is_builtin("dict"))
+        {
+            self.expr_infer_with_hint(&x.value, Some(HintRef::soft(&base)), errors)
+        } else {
+            self.expr_infer(&x.value, errors)
+        };
         if matches!(x.op, Operator::Div | Operator::FloorDiv | Operator::Mod)
             && Self::is_literal_zero(&rhs)
         {

pyrefly/lib/test/operators.rs

@@ -20,6 +20,20 @@ def f(a: int, b: int) -> None:
 "#,
 );
 
+testcase!(
+    test_dict_union_literal_keys,
+    r#"
+from typing import Literal
+
+Allowed = Literal["a", "b", "c"]
+d: dict[Allowed, int] = {"a": 0, "b": 0}
+e: dict[Allowed, int] = d | {"c": 0}
+d |= {"c": 0}
+bad_key: dict[Allowed, int] = d | {"not-allowed": 0}  # E: `dict[Literal['a', 'b', 'c'] | str, int]` is not assignable to `dict[Allowed, int]`
+bad_value: dict[Allowed, int] = d | {"a": "not-an-int"}  # E: `dict[Literal['a', 'b', 'c'] | str, int | str]` is not assignable to `dict[Allowed, int]`
+    "#,
+);
+
 testcase!(
     test_bounded_type_var_comparison,
     r#"