Refactor, fix build, and mark uninit twice

Author: fangyi-zhou

pyrefly/lib/binding/stmt.rs

@@ -35,6 +35,7 @@ use crate::binding::binding::KeyExpect;
 use crate::binding::binding::LinkedKey;
 use crate::binding::binding::RaisedException;
 use crate::binding::bindings::BindingsBuilder;
+use crate::binding::bindings::LookupKind;
 use crate::binding::bindings::MutableCaptureLookupKind;
 use crate::binding::expr::Usage;
 use crate::binding::narrow::NarrowOps;
@@ -247,6 +248,25 @@ impl<'a> BindingsBuilder<'a> {
         self.scopes.mark_flow_termination();
     }
 
+    fn mark_exception_name_uninitialized(&mut self, name: &Identifier) {
+        // https://docs.python.org/3/reference/compound_stmts.html#except-clause
+        // Mark the handler name as uninitialized, so that it can't be used.
+        let idx = self.lookup_name(
+            Hashed::new(&name.id),
+            LookupKind::Regular,
+            &mut Usage::MutableLookup,
+        );
+        if let Ok(idx) = idx {
+            self.scopes.upsert_flow_info(
+                Hashed::new(&name.id),
+                idx,
+                Some(FlowStyle::Uninitialized),
+            );
+        } else {
+            panic!("Should have found the exception name `{name}` in the current scope");
+        }
+    }
+
     /// Evaluate the statements and update the bindings.
     /// Every statement should end up in the bindings, perhaps with a location that is never used.
     pub fn stmt(&mut self, x: Stmt) {
@@ -731,8 +751,12 @@ impl<'a> BindingsBuilder<'a> {
                 self.scopes.swap_current_flow_with(&mut base);
                 branches.push(base);
 
-                // Store the last name of the exception, so that we can mark it uninitialized in the next scope.
-                let mut last_name: Option<Identifier> = None;
+                // Store the last caught exception name, so that we can mark any
+                // caught exception to be uninitialized at the end of an
+                // `except` block, and at the beginning of the next block. This
+                // allows us to ensure that the caught exception name goes out
+                // of scope after the corresponding `except` block.
+                let mut last_caught_exception_name: Option<Identifier> = None;
 
                 for h in x.handlers {
                     base = self.scopes.clone_current_flow();
@@ -758,49 +782,38 @@ impl<'a> BindingsBuilder<'a> {
                             Binding::ExceptionHandler(type_, x.is_star),
                         );
                     }
-                    if let Some(ref name) = last_name {
+
+                    if let Some(ref name) = last_caught_exception_name {
                         // https://docs.python.org/3/reference/compound_stmts.html#except-clause
-                        // Mark the previous handler name as uninitialized in
-                        // the current scope, so that it can't be used.
-                        let idx = self.lookup_name(Hashed::new(&name.id), LookupKind::Regular);
-                        if let Ok(idx) = idx {
-                            self.scopes.upsert_flow_info(
-                                Hashed::new(&name.id),
-                                idx,
-                                Some(FlowStyle::Uninitialized),
-                            );
-                        } else {
-                            panic!(
-                                "Should have found the handler name {name} in the current scope"
-                            );
-                        }
+                        // Mark the previous caught exception name as
+                        // uninitialized in the current scope, so that it can't
+                        // be used.
+                        self.mark_exception_name_uninitialized(name);
                     }
                     self.stmts(h.body);
+
+                    if let Some(ref name) = h.name {
+                        // Mark the current caught exception name as
+                        // uninitialized in the current scope, so that the flow
+                        // merges correctly.
+                        self.mark_exception_name_uninitialized(name);
+                    }
+
                     self.scopes.swap_current_flow_with(&mut base);
 
                     if let Some(name) = h.name {
-                        last_name = Some(name.clone());
+                        last_caught_exception_name = Some(name.clone());
                     } else {
-                        last_name = None;
+                        last_caught_exception_name = None;
                     }
                     branches.push(base);
                 }
 
                 self.set_current_flow_to_merged_branches(branches, range);
-                if let Some(ref name) = last_name {
-                    // https://docs.python.org/3/reference/compound_stmts.html#except-clause
-                    // Mark the previous handler name as uninitialized in
+                if let Some(ref name) = last_caught_exception_name {
+                    // Mark the last caught exception name as uninitialized in
                     // the current scope, so that it can't be used.
-                    let idx = self.lookup_name(Hashed::new(&name.id), LookupKind::Regular);
-                    if let Ok(idx) = idx {
-                        self.scopes.upsert_flow_info(
-                            Hashed::new(&name.id),
-                            idx,
-                            Some(FlowStyle::Uninitialized),
-                        );
-                    } else {
-                        panic!("Should have found the handler name {name} in the current scope");
-                    }
+                    self.mark_exception_name_uninitialized(name);
                 }
                 self.stmts(x.finalbody);
             }

pyrefly/lib/test/scope.rs

@@ -468,14 +468,32 @@ def g():
     finally:
         e2 # E: `e2` is uninitialized
 
+    e2 # E: `e2` is uninitialized
+
 def h():
     try:
         1 / 0
     except OSError as e3:
         pass
     except Exception:
         e3 # E: `e3` is uninitialized
+
+    e3 # E: `e3` is uninitialized
+
+def i():
+    try:
+        1 / 0
+    except OSError as e4:
+        pass
+    except Exception:
+        pass
+    except OSError as e5:
+        pass
     finally:
-        e3 # E: `e3` may be uninitialized
+        e4 # E: `e4` is uninitialized
+        e5 # E: `e5` is uninitialized
+
+    e4 # E: `e4` is uninitialized
+    e5 # E: `e5` is uninitialized
 "#,
 );