impl code action for narrowing the value with an is not None #4169

pyrefly/lib/alt/answers_solver.rs

@@ -3251,24 +3251,27 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 .with_quick_fix(ErrorQuickFix::ReplaceWithEnumMember { replacement });
         }
         if Self::type_contains_none(got) && !Self::type_contains_none(want) {
-            let hint = match tcc().kind {
+            let (hint, offer_narrowing_fix) = match tcc().kind {
                 TypeCheckKind::ExplicitFunctionReturn
                 | TypeCheckKind::AnnAssign
                 | TypeCheckKind::AnnotatedName(_)
                     if !got.is_none() =>
                 {
-                    Some(format!(
-                        "Consider narrowing the value with an `is not None` check or changing the declared type to `{} | None`",
-                        self.for_display(want.clone()),
-                    ))
+                    (
+                        Some(format!(
+                            "Consider narrowing the value with an `is not None` check or changing the declared type to `{} | None`",
+                            self.for_display(want.clone()),
+                        )),
+                        true,
+                    )
                 }
                 TypeCheckKind::ImplicitFunctionReturn(_) => {
                     // For implicit returns (missing return statement), the error message
                     // "Function declared to return X, but one or more paths are missing
                     // an explicit return" is already clear. Adding a None hint here is
                     // confusing because the user didn't explicitly return None.
                     // Skip the hint for implicit returns.
-                    None
+                    (None, false)
                 }
                 TypeCheckKind::Attribute(_)
                 | TypeCheckKind::CallArgument(..)
@@ -3279,22 +3282,34 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         // Skip the hint. Narrowing the value doesn't make sense if there's no
                         // non-None part to narrow to, and changing the attribute or parameter type
                         // is often unactionable, since the definition may be in third-party code.
-                        None
+                        (None, false)
                     } else {
                         // We only suggest narrowing. Changing the attribute or parameter type is
                         // often unactionable, since the definition may be in third-party code.
-                        Some("Consider narrowing the value with an `is not None` check".to_owned())
+                        (
+                            Some(
+                                "Consider narrowing the value with an `is not None` check"
+                                    .to_owned(),
+                            ),
+                            true,
+                        )
                     }
                 }
-                _ => Some(format!(
-                    "Consider changing the declared type to `{} | None`",
-                    self.for_display(want.clone())
-                )),
+                _ => (
+                    Some(format!(
+                        "Consider changing the declared type to `{} | None`",
+                        self.for_display(want.clone())
+                    )),
+                    false,
+                ),
             };
             if let Some(hint) = hint {
                 builder = builder
                     .with_detail(format!("The declared type does not allow `None`. {hint}."));
             }
+            if offer_narrowing_fix {
+                builder = builder.with_quick_fix(ErrorQuickFix::AssertNotNone);
+            }
         }
         builder.emit();
     }

pyrefly/lib/alt/attr.rs

@@ -41,6 +41,7 @@ use crate::error::collector::ErrorCollector;
 use crate::error::context::ErrorContext;
 use crate::error::context::TypeCheckContext;
 use crate::error::context::TypeCheckKind;
+use crate::error::error::ErrorQuickFix;
 use crate::error::style::ErrorStyle;
 use crate::solver::solver::SubsetError;
 use crate::state::loader::FindingOrError;
@@ -605,10 +606,20 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         // Check if we have a partial union failure (attribute exists on some union members
         // but not others) before consuming the vectors. This helps us decide whether to suggest.
         let is_partial_union_failure = !found.is_empty() && !not_found.is_empty();
+        let mut can_narrow_none = is_partial_union_failure
+            && error.is_empty()
+            && not_found.iter().all(|missing| {
+                matches!(
+                    missing,
+                    NotFoundOn::ClassInstance(class, _)
+                        if class == self.stdlib.none_type().class_object()
+                )
+            });
         for (attr, _) in found {
             match self.resolve_get_access(attr_name, attr, range, errors, context) {
                 Ok(ty) => types.push(ty),
                 Err(err) => {
+                    can_narrow_none = false;
                     error_messages.push(err.to_error_msg(attr_name));
                     success = false;
                 }
@@ -644,11 +655,14 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 msg.push(format!("Did you mean `{suggestion}`?"));
             }
             let (header, details) = msg.split_off_first();
-            errors
+            let mut builder = errors
                 .error_builder(range, ErrorKind::MissingAttribute, header)
                 .with_details(details)
-                .with_context(context)
-                .emit();
+                .with_context(context);
+            if can_narrow_none {
+                builder = builder.with_quick_fix(ErrorQuickFix::AssertNotNone);
+            }
+            builder.emit();
             self.heap.mk_any_error()
         } else {
             self.heap.mk_any_error() // we've encountered internal errors (already logged above)

pyrefly/lib/error/collector.rs

@@ -67,7 +67,10 @@ impl ModuleErrors {
                 previous_range = x.range();
                 previous_start = self.items.len();
                 self.items.push(x);
-            } else if !self.items[previous_start..].contains(&x) {
+            } else if !self.items[previous_start..]
+                .iter_mut()
+                .any(|existing| existing.merge_if_same_diagnostic(&x))
+            {
                 self.items.push(x);
             }
         }

pyrefly/lib/error/error.rs

@@ -47,6 +47,7 @@ pub struct SecondaryAnnotation {
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ErrorQuickFix {
     ReplaceWithEnumMember { replacement: String },
+    AssertNotNone,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -445,6 +446,27 @@ impl Error {
         self
     }
 
+    /// Merge editor fixes when both values describe the same user-facing diagnostic.
+    pub(crate) fn merge_if_same_diagnostic(&mut self, other: &Self) -> bool {
+        if self.module != other.module
+            || self.range != other.range
+            || self.display_range != other.display_range
+            || self.error_kind != other.error_kind
+            || self.severity != other.severity
+            || self.msg_header != other.msg_header
+            || self.msg_details != other.msg_details
+            || self.secondary_annotations != other.secondary_annotations
+        {
+            return false;
+        }
+        for fix in &other.quick_fixes {
+            if !self.quick_fixes.contains(fix) {
+                self.quick_fixes.push(fix.clone());
+            }
+        }
+        true
+    }
+
     pub fn display_range(&self) -> &DisplayRange {
         &self.display_range
     }

pyrefly/lib/state/lsp.rs

@@ -2972,6 +2972,18 @@ impl<'a> Transaction<'a> {
                     other_actions.push(action);
                 }
             }
+            if error_range.contains_range(range)
+                && let Some(action) = quick_fixes::assert_not_none::assert_not_none_code_action(
+                    &module_info,
+                    &ast,
+                    &error,
+                )
+            {
+                let key = (action.0.clone(), action.2, action.3.clone());
+                if other_action_keys.insert(key) {
+                    other_actions.push(action);
+                }
+            }
             if error_range.contains_range(range)
                 && let Some(action) = quick_fixes::pyrefly_ignore::add_pyrefly_ignore_code_action(
                     &module_info,

pyrefly/lib/state/lsp/quick_fixes/assert_not_none.rs

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use dupe::Dupe;
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module::Module;
+use ruff_python_ast::AnyNodeRef;
+use ruff_python_ast::Expr;
+use ruff_python_ast::ModModule;
+use ruff_text_size::Ranged;
+use ruff_text_size::TextRange;
+
+use crate::ModuleInfo;
+use crate::error::error::Error;
+use crate::error::error::ErrorQuickFix;
+use crate::state::lsp::quick_fixes::extract_shared::find_enclosing_statement_range;
+use crate::state::lsp::quick_fixes::extract_shared::line_indent_and_start;
+
+/// Insert an assertion before a standalone statement that uses an optional name.
+pub(crate) fn assert_not_none_code_action(
+    module_info: &ModuleInfo,
+    ast: &ModModule,
+    error: &Error,
+) -> Option<(String, Module, TextRange, String)> {
+    if !error
+        .quick_fixes()
+        .iter()
+        .any(|fix| matches!(fix, ErrorQuickFix::AssertNotNone))
+    {
+        return None;
+    }
+    let name = Ast::locate_node(ast, error.range().start())
+        .into_iter()
+        .find_map(|node| match node {
+            AnyNodeRef::ExprName(name) if name.range().contains_range(error.range()) => Some(name),
+            AnyNodeRef::ExprAttribute(attribute)
+                if attribute.range().contains_range(error.range())
+                    && let Expr::Name(name) = attribute.value.as_ref() =>
+            {
+                Some(name)
+            }
+            _ => None,
+        })?;
+    let statement = find_enclosing_statement_range(ast, error.range())?;
+    let source = module_info.contents().as_str();
+    let (indent, line_start) = line_indent_and_start(source, statement.start())?;
+
+    // Inserting before a statement is only valid when it starts on its own line.
+    if source[line_start.to_usize()..statement.start().to_usize()] != indent {
+        return None;
+    }
+
+    let assertion = format!("assert {} is not None", name.id);
+    Some((
+        format!("Add `{assertion}`"),
+        module_info.dupe(),
+        TextRange::new(line_start, line_start),
+        format!("{indent}{assertion}\n"),
+    ))
+}

pyrefly/lib/state/lsp/quick_fixes/enum_member.rs

@@ -33,8 +33,10 @@ pub(crate) fn replace_with_enum_member_code_action(
 }
 
 fn enum_member_replacement(error: &Error) -> Option<&str> {
-    let ErrorQuickFix::ReplaceWithEnumMember { replacement } = error.quick_fixes().first()?;
-    Some(replacement.as_str())
+    error.quick_fixes().iter().find_map(|fix| match fix {
+        ErrorQuickFix::ReplaceWithEnumMember { replacement } => Some(replacement.as_str()),
+        ErrorQuickFix::AssertNotNone => None,
+    })
 }
 
 fn enclosing_string_literal_range(ast: &ModModule, error_range: TextRange) -> Option<TextRange> {

pyrefly/lib/state/lsp/quick_fixes/mod.rs

@@ -6,6 +6,7 @@
  */
 
 pub(crate) mod add_override;
+pub(crate) mod assert_not_none;
 pub(crate) mod convert_dict;
 pub(crate) mod convert_star_import;
 pub(crate) mod enum_member;

pyrefly/lib/test/lsp/code_actions.rs

@@ -63,6 +63,19 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String
     report
 }
 
+fn get_test_report_with_error_count(state: &State, handle: &Handle, position: TextSize) -> String {
+    let error_count = state
+        .transaction()
+        .get_errors(vec![handle])
+        .collect_errors()
+        .ordinary
+        .len();
+    format!(
+        "Error count: {error_count}\n{}",
+        get_test_report(state, handle, position)
+    )
+}
+
 fn apply_refactor_edits_for_module(
     module: &ModuleInfo,
     edits: &[(Module, TextRange, String)],
@@ -1031,6 +1044,80 @@ takes_status("active")
     );
 }
 
+#[test]
+fn quickfix_narrow_optional_argument() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[(
+            "main",
+            r#"def g(y: int) -> None: ...
+def f(x: int | None) -> None:
+    g(x)
+#     ^
+"#,
+        )],
+        get_test_report,
+    );
+    assert!(
+        report.contains("# Title: Add `assert x is not None`"),
+        "{report}"
+    );
+    assert!(
+        report.contains("    assert x is not None\n    g(x)"),
+        "{report}"
+    );
+}
+
+#[test]
+fn quickfix_narrow_optional_attribute_receiver() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[(
+            "main",
+            r#"class T:
+    y: int
+
+def f(x: T | None) -> None:
+    x.y
+#   ^
+"#,
+        )],
+        get_test_report,
+    );
+    assert!(
+        report.contains("# Title: Add `assert x is not None`"),
+        "{report}"
+    );
+    assert!(
+        report.contains("    assert x is not None\n    x.y"),
+        "{report}"
+    );
+}
+
+#[test]
+fn quickfix_narrow_optional_attribute_augmented_assignment() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[(
+            "main",
+            r#"class T:
+    y: int
+
+def f(x: T | None) -> None:
+    x.y += 1
+#   ^
+"#,
+        )],
+        get_test_report_with_error_count,
+    );
+    assert!(report.contains("Error count: 1"), "{report}");
+    assert!(
+        report.contains("# Title: Add `assert x is not None`"),
+        "{report}"
+    );
+    assert!(
+        report.contains("    assert x is not None\n    x.y += 1"),
+        "{report}"
+    );
+}
+
 #[test]
 fn quickfix_add_pyrefly_ignore_code_with_existing_comment() {
     let report = get_batched_lsp_operations_report_allow_error(