preserve literal types for module-scoped final vars inside function scopes

Summary: fixes #3378

Reviewed By: stroxler

Differential Revision: D105257293

fbshipit-source-id: d210327d091b05efd07aa2faf61c511027063ef7

Author: yangdanny97

pyrefly/lib/binding/expr.rs

@@ -373,7 +373,8 @@ impl<'a> BindingsBuilder<'a> {
 
                 let promote = self.scopes.in_function_scope()
                     && (is_module_scope || self.scopes.is_defined_at_module_scope(&name.id))
-                    && !is_constant_name(&name.id);
+                    && !is_constant_name(&name.id)
+                    && !self.scopes.is_final_at_module_scope(&name.id);
                 if promote {
                     self.promote_ranges.insert(name.range);
                 }

pyrefly/lib/binding/scope.rs

@@ -404,8 +404,8 @@ impl Static {
     }
 
     /// Populate static definitions from a list of statements.
-    /// Returns the set of implicit captures (names read but not locally defined)
-    /// and a map of Final variable string values.
+    /// Returns the set of implicit captures (names read but not locally defined),
+    /// the set of all Final names, and a map of Final variable string values.
     fn stmts(
         &mut self,
         x: &[Stmt],
@@ -415,7 +415,7 @@ impl Static {
         sys_info: SysInfo,
         get_annotation_idx: &mut impl FnMut(ShortIdentifier) -> Idx<KeyAnnotation>,
         scopes: Option<&Scopes>,
-    ) -> (SmallSet<Name>, SmallMap<Name, String>) {
+    ) -> (SmallSet<Name>, SmallSet<Name>, SmallMap<Name, String>) {
         let mut d = Definitions::new(
             x,
             module_info.name(),
@@ -464,12 +464,13 @@ impl Static {
                 self.upsert(name.cloned(), range, StaticStyle::MergeableImport, range)
             }
         }
+        let final_names = d.final_names.keys().cloned().collect();
         let final_string_values = d
             .final_names
             .into_iter()
             .filter_map(|(name, value)| value.map(|v| (name, v)))
             .collect();
-        (implicit_captures, final_string_values)
+        (implicit_captures, final_names, final_string_values)
     }
 
     fn expr_lvalue(&mut self, x: &Expr) {
@@ -1206,6 +1207,9 @@ pub struct Scope {
     /// from enclosing scopes. Populated during `init_current_static` from the
     /// `Definitions` phase. Used to seed flow entries for captured variables.
     implicit_captures: SmallSet<Name>,
+    /// All names marked `Final` in this scope. Used to prevent literal
+    /// promotion so that `Final` variables preserve their literal types.
+    final_names: SmallSet<Name>,
     /// Names marked `Final` with string literal values, e.g. `X: Final = "x"`.
     /// Used to resolve Final variable references in synthesized class field names.
     final_string_values: SmallMap<Name, String>,
@@ -1227,6 +1231,7 @@ impl Scope {
             finally_depth: 0,
             with_depth: 0,
             implicit_captures: SmallSet::new(),
+            final_names: SmallSet::new(),
             final_string_values: SmallMap::new(),
         }
     }
@@ -1673,7 +1678,7 @@ impl Scopes {
         get_annotation_idx: &mut impl FnMut(ShortIdentifier) -> Idx<KeyAnnotation>,
     ) {
         let mut initialize = |scope: &mut Scope, myself: Option<&Self>| {
-            let (implicit_captures, final_string_values) = scope.stat.stmts(
+            let (implicit_captures, final_names, final_string_values) = scope.stat.stmts(
                 x,
                 module_info,
                 top_level,
@@ -1683,6 +1688,7 @@ impl Scopes {
                 myself,
             );
             scope.implicit_captures = implicit_captures;
+            scope.final_names = final_names;
             scope.final_string_values = final_string_values;
             // Presize the flow, as its likely to need as much space as static
             scope.flow.info.reserve(scope.stat.0.capacity());
@@ -1701,6 +1707,11 @@ impl Scopes {
         }
     }
 
+    /// Check if a name is declared as `Final` at module scope.
+    pub fn is_final_at_module_scope(&self, name: &Name) -> bool {
+        self.scopes.first().scope.final_names.contains(name)
+    }
+
     /// Look up a Final variable's string literal value in the current scope stack.
     /// Searches from the innermost scope outward, stopping at the first scope
     /// that binds the name, even if it's not Final.

pyrefly/lib/test/narrow.rs

@@ -159,6 +159,50 @@ def f2(x: E | int):
     "#,
 );
 
+testcase!(
+    test_is_not_final_enum,
+    r#"
+from enum import Enum
+from typing import Final, Iterable, assert_type
+
+class EmptyType(Enum):
+    EMPTY = 0
+
+Empty: Final = EmptyType.EMPTY
+
+class Bar:
+    baz: int | EmptyType = 4
+
+def foo(bar: Bar) -> Iterable[int]:
+    if bar.baz is not Empty:
+        assert_type(bar.baz, int)
+        return [bar.baz]
+    return []
+    "#,
+);
+
+testcase!(
+    test_is_not_final_enum_type_alias,
+    r#"
+from enum import Enum
+from typing import Final, Iterable, Literal, TypeAlias
+
+class _EmptyEnum(Enum):
+    EMPTY = 0
+
+EmptyType: TypeAlias = Literal[_EmptyEnum.EMPTY]
+Empty: Final = _EmptyEnum.EMPTY
+
+class Bar:
+    baz: int | EmptyType = 4
+
+def foo(bar: Bar) -> Iterable[int]:
+    if bar.baz is not Empty:
+        return [bar.baz]
+    return []
+    "#,
+);
+
 testcase!(
     test_ellipsis_is,
     r#"