[red-knot] Add definitions and limited type inference for exception handlers (#13267)

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -880,6 +880,30 @@ where
 
         self.current_match_case.as_mut().unwrap().index += 1;
     }
+
+    fn visit_except_handler(&mut self, except_handler: &'ast ast::ExceptHandler) {
+        let ast::ExceptHandler::ExceptHandler(except_handler) = except_handler;
+        let ast::ExceptHandlerExceptHandler {
+            name: symbol_name,
+            type_: handled_exceptions,
+            body,
+            range: _,
+        } = except_handler;
+
+        if let Some(handled_exceptions) = handled_exceptions {
+            self.visit_expr(handled_exceptions);
+        }
+
+        // If `handled_exceptions` above was `None`, it's something like `except as e:`,
+        // which is invalid syntax. However, it's still pretty obvious here that the user
+        // *wanted* `e` to be bound, so we should still create a definition here nonetheless.
+        if let Some(symbol_name) = symbol_name {
+            let symbol = self.add_or_update_symbol(symbol_name.id.clone(), SymbolFlags::IS_DEFINED);
+            self.add_definition(symbol, except_handler);
+        }
+
+        self.visit_body(body);
+    }
 }
 
 #[derive(Copy, Clone, Debug)]

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -50,6 +50,7 @@ pub(crate) enum DefinitionNodeRef<'a> {
     Parameter(ast::AnyParameterRef<'a>),
     WithItem(WithItemDefinitionNodeRef<'a>),
     MatchPattern(MatchPatternDefinitionNodeRef<'a>),
+    ExceptHandler(&'a ast::ExceptHandlerExceptHandler),
 }
 
 impl<'a> From<&'a ast::StmtFunctionDef> for DefinitionNodeRef<'a> {
@@ -130,6 +131,12 @@ impl<'a> From<MatchPatternDefinitionNodeRef<'a>> for DefinitionNodeRef<'a> {
     }
 }
 
+impl<'a> From<&'a ast::ExceptHandlerExceptHandler> for DefinitionNodeRef<'a> {
+    fn from(node: &'a ast::ExceptHandlerExceptHandler) -> Self {
+        Self::ExceptHandler(node)
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct ImportFromDefinitionNodeRef<'a> {
     pub(crate) node: &'a ast::StmtImportFrom,
@@ -248,6 +255,9 @@ impl DefinitionNodeRef<'_> {
                 identifier: AstNodeRef::new(parsed, identifier),
                 index,
             }),
+            DefinitionNodeRef::ExceptHandler(handler) => {
+                DefinitionKind::ExceptHandler(AstNodeRef::new(parsed, handler))
+            }
         }
     }
 
@@ -280,6 +290,7 @@ impl DefinitionNodeRef<'_> {
             Self::MatchPattern(MatchPatternDefinitionNodeRef { identifier, .. }) => {
                 identifier.into()
             }
+            Self::ExceptHandler(handler) => handler.into(),
         }
     }
 }
@@ -300,6 +311,7 @@ pub enum DefinitionKind {
     ParameterWithDefault(AstNodeRef<ast::ParameterWithDefault>),
     WithItem(WithItemDefinitionKind),
     MatchPattern(MatchPatternDefinitionKind),
+    ExceptHandler(AstNodeRef<ast::ExceptHandlerExceptHandler>),
 }
 
 #[derive(Clone, Debug)]
@@ -478,3 +490,9 @@ impl From<&ast::Identifier> for DefinitionNodeKey {
         Self(NodeKey::from_node(identifier))
     }
 }
+
+impl From<&ast::ExceptHandlerExceptHandler> for DefinitionNodeKey {
+    fn from(handler: &ast::ExceptHandlerExceptHandler) -> Self {
+        Self(NodeKey::from_node(handler))
+    }
+}

crates/red_knot_python_semantic/src/types/infer.rs

@@ -425,6 +425,9 @@ impl<'db> TypeInferenceBuilder<'db> {
                     definition,
                 );
             }
+            DefinitionKind::ExceptHandler(handler) => {
+                self.infer_except_handler_definition(handler, definition);
+            }
         }
     }
 
@@ -743,11 +746,29 @@ impl<'db> TypeInferenceBuilder<'db> {
         } = try_statement;
 
         self.infer_body(body);
+
         for handler in handlers {
             let ast::ExceptHandler::ExceptHandler(handler) = handler;
-            self.infer_optional_expression(handler.type_.as_deref());
+            let ast::ExceptHandlerExceptHandler {
-            self.infer_body(&handler.body);
+                type_: handled_exceptions,
+                name: symbol_name,
+                body,
+                range: _,
+            } = handler;
+
+            // If `symbol_name` is `Some()` and `handled_exceptions` is `None`,
+            // it's invalid syntax (something like `except as e:`).
+            // However, it's obvious that the user *wanted* `e` to be bound here,
+            // so we'll have created a definition in the semantic-index stage anyway.
+            if symbol_name.is_some() {
+                self.infer_definition(handler);
+            } else {
+                self.infer_optional_expression(handled_exceptions.as_deref());
+            }
+
+            self.infer_body(body);
         }
+
         self.infer_body(orelse);
         self.infer_body(finalbody);
     }
@@ -797,6 +818,29 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.types.definitions.insert(definition, context_expr_ty);
     }
 
+    fn infer_except_handler_definition(
+        &mut self,
+        handler: &'db ast::ExceptHandlerExceptHandler,
+        definition: Definition<'db>,
+    ) {
+        let node_ty = handler
+            .type_
+            .as_deref()
+            .map(|ty| self.infer_expression(ty))
+            .unwrap_or(Type::Unknown);
+
+        // TODO: anything that's a consistent subtype of
+        // `type[BaseException] | tuple[type[BaseException], ...]` should be valid;
+        // anything else should be invalid --Alex
+        let symbol_ty = match node_ty {
+            Type::Any | Type::Unknown => node_ty,
+            Type::Class(class_ty) => Type::Instance(class_ty),
+            _ => Type::Unknown,
+        };
+
+        self.types.definitions.insert(definition, symbol_ty);
+    }
+
     fn infer_match_statement(&mut self, match_statement: &ast::StmtMatch) {
         let ast::StmtMatch {
             range: _,
@@ -4180,6 +4224,108 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn except_handler_single_exception() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            import re
+
+            try:
+                x
+            except NameError as e:
+                pass
+            except re.error as f:
+                pass
+            ",
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "e", "NameError");
+        assert_public_ty(&db, "src/a.py", "f", "error");
+        assert_file_diagnostics(&db, "src/a.py", &[]);
+
+        Ok(())
+    }
+
+    #[test]
+    fn unknown_type_in_except_handler_does_not_cause_spurious_diagnostic() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            from nonexistent_module import foo
+
+            try:
+                x
+            except foo as e:
+                pass
+            ",
+        )?;
+
+        assert_file_diagnostics(
+            &db,
+            "src/a.py",
+            &["Cannot resolve import 'nonexistent_module'."],
+        );
+        assert_public_ty(&db, "src/a.py", "foo", "Unknown");
+        assert_public_ty(&db, "src/a.py", "e", "Unknown");
+
+        Ok(())
+    }
+
+    #[test]
+    fn except_handler_multiple_exceptions() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            EXCEPTIONS = (AttributeError, TypeError)
+
+            try:
+                x
+            except (RuntimeError, OSError) as e:
+                pass
+            except EXCEPTIONS as f:
+                pass
+            ",
+        )?;
+
+        assert_file_diagnostics(&db, "src/a.py", &[]);
+
+        // For these TODOs we need support for `tuple` types:
+
+        // TODO: Should be `RuntimeError | OSError` --Alex
+        assert_public_ty(&db, "src/a.py", "e", "Unknown");
+        // TODO: Should be `AttributeError | TypeError` --Alex
+        assert_public_ty(&db, "src/a.py", "e", "Unknown");
+
+        Ok(())
+    }
+
+    #[test]
+    fn exception_handler_with_invalid_syntax() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            try:
+                x
+            except as e:
+                pass
+            ",
+        )?;
+
+        assert_file_diagnostics(&db, "src/a.py", &[]);
+        assert_public_ty(&db, "src/a.py", "e", "Unknown");
+
+        Ok(())
+    }
+
     #[test]
     fn invalid_iterable() {
         let mut db = setup_db();

crates/ruff_benchmark/benches/red_knot.rs

@@ -33,7 +33,6 @@ static EXPECTED_DIAGNOSTICS: &[&str] = &[
     "Use double quotes for strings",
     "Use double quotes for strings",
     "Use double quotes for strings",
-    "/src/tomllib/_parser.py:628:75: Name 'e' used when not defined.",
 ];
 
 fn get_test_file(name: &str) -> TestFile {