fix Narrow the type of a when branching on a in b #2608 (#2627)

Summary:
Fixes #2608

Implemented narrowing for a in mapping by intersecting the LHS with the mapping key type. This fixes not in branches by negation.

Pull Request resolved: #2627

Test Plan: add test

Reviewed By: rchen152

Differential Revision: D95416263

Pulled By: yangdanny97

fbshipit-source-id: d016249a6eb9075ed5593d78e9f63607e1f8aba3

Author: asukaminato0721

pyrefly/lib/alt/narrow.rs

@@ -901,7 +901,15 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                     return self.intersect(ty, &self.unions(key_types));
                 }
 
-                ty.clone()
+                // Check if the right operand is a mapping (e.g. dict[str, int]).
+                // If so, we can narrow the left operand to the mapping's key type.
+                if !self.behaves_like_any(&right_ty)
+                    && let Some((key_ty, _)) = self.unwrap_mapping(&right_ty)
+                {
+                    self.intersect(ty, &key_ty)
+                } else {
+                    ty.clone()
+                }
             }
             AtomicNarrowOp::NotIn(v) => {
                 // First, check for List, Tuple, and Set literal expressions (syntactic check,

pyrefly/lib/test/narrow.rs

@@ -1662,6 +1662,22 @@ class MyTest(TestCase):
 "#,
 );
 
+testcase!(
+    test_in_mapping_narrows_key,
+    r#"
+from typing import assert_type
+
+def parse_resource_id() -> tuple[str, str] | tuple[None, None]: ...
+
+def lookup_resource(registry: dict[str, str]) -> str | None:
+    kind, _obj_id = parse_resource_id()
+    if kind not in registry:
+        return None
+    assert_type(kind, str)
+    return registry[kind]
+"#,
+);
+
 // Make sure we catch illegal arguments to isinstance and issubclass even when we aren't narrowing.
 testcase!(
     test_validate_class_object_no_narrow,