Allow breaking after = in named arguments when name = value exceeds line-width

crates/air_r_formatter/src/r/auxiliary/argument.rs

@@ -1,7 +1,16 @@
 use crate::prelude::*;
+use air_r_syntax::AnyRExpression;
 use air_r_syntax::RArgument;
 use air_r_syntax::RArgumentFields;
+use air_r_syntax::RArgumentList;
+use air_r_syntax::RCall;
+use air_r_syntax::RCallArguments;
+use biome_formatter::FormatOptions;
+use biome_formatter::VecBuffer;
+use biome_formatter::format_element::LineMode;
+use biome_formatter::format_element::tag::Tag;
 use biome_formatter::write;
+use biome_rowan::AstNode;
 
 #[derive(Debug, Clone, Default)]
 pub(crate) struct FormatRArgument;
@@ -36,7 +45,134 @@ pub(crate) fn fmt_argument_fields(node: &RArgument, f: &mut RFormatter) -> Forma
         // Named argument with a value
         // `foo(name = value)`
         (Some(name_clause), Some(value)) => {
-            write!(f, [name_clause.format(), space(), value.format()])
+            if is_in_test_call(node)
+                || matches!(
+                    value,
+                    AnyRExpression::RBracedExpressions(_)
+                        | AnyRExpression::RFunctionDefinition(_)
+                        | AnyRExpression::RIfStatement(_)
+                        | AnyRExpression::RForStatement(_)
+                        | AnyRExpression::RWhileStatement(_)
+                        | AnyRExpression::RRepeatStatement(_)
+                )
+                || is_call_with_long_callee(&value, f)
+            {
+                return write!(f, [name_clause.format(), space(), value.format()]);
+            }
+
+            let is_binary = matches!(value, AnyRExpression::RBinaryExpression(_));
+            let format_value = format_with(|f| {
+                let mut buffer = VecBuffer::new(f.state_mut());
+                write!(buffer, [value.format()])?;
+                let mut elements = buffer.into_vec();
+                flatten_argument_value_soft_lines(&mut elements, is_binary, f);
+                f.write_elements(elements)
+            });
+
+            let group_id = f.group_id("named_argument_value");
+            write!(
+                f,
+                [
+                    name_clause.format(),
+                    group(&indent(&soft_line_break_or_space())).with_group_id(Some(group_id)),
+                    indent_if_group_breaks(&format_value, group_id)
+                ]
+            )
+        }
+    }
+}
+
+/// Flattens soft line breaks in a named argument's value (or its `<left> <op>` prefix
+/// for a multi-line binary expression) when that span has no forced breaks and fits on a
+/// single indented line on its own, so `FitsMeasurer` measures the full single-line
+/// expression (or `<left> <op>` prefix) rather than stopping at an earlier soft line break.
+fn flatten_argument_value_soft_lines(
+    elements: &mut Vec<FormatElement>,
+    is_binary: bool,
+    f: &RFormatter,
+) {
+    let max_flat_width = usize::from(f.options().line_width().value())
+        .saturating_sub(usize::from(f.options().indent_width().value()) * 4);
+
+    let prefix_len = if !elements.will_break() && flat_text_width(elements) <= max_flat_width {
+        elements.len()
+    } else if is_binary
+        && let Some(idx) = first_binary_op_line_idx(elements)
+        && flat_text_width(&elements[..idx]) <= max_flat_width
+    {
+        idx
+    } else {
+        return;
+    };
+
+    let mut remaining = prefix_len;
+    elements.retain_mut(|item| {
+        if remaining == 0 {
+            return true;
+        }
+        remaining -= 1;
+        if matches!(item, FormatElement::Line(LineMode::SoftOrSpace)) {
+            *item = FormatElement::Space;
+        }
+        !matches!(item, FormatElement::Line(LineMode::Soft))
+    });
+}
+
+fn first_binary_op_line_idx(slice: &[FormatElement]) -> Option<usize> {
+    let mut group_depth = 0usize;
+    let mut in_expanded_group = false;
+    for (idx, item) in slice.iter().enumerate() {
+        match item {
+            FormatElement::Tag(Tag::StartGroup(group)) => {
+                group_depth += 1;
+                if group_depth > 1 && !group.mode().is_flat() {
+                    in_expanded_group = true;
+                }
+            }
+            FormatElement::Tag(Tag::EndGroup) => {
+                group_depth = group_depth.saturating_sub(1);
+            }
+            FormatElement::Line(_) if group_depth == 1 || in_expanded_group => {
+                return Some(idx);
+            }
+            other if other.will_break() => return None,
+            _ => {}
         }
     }
+    None
+}
+
+fn flat_text_width(slice: &[FormatElement]) -> usize {
+    slice
+        .iter()
+        .map(|item| match item {
+            FormatElement::StaticText { text } => text.len(),
+            FormatElement::DynamicText { text, .. } => text.len(),
+            FormatElement::LocatedTokenText { slice, .. } => usize::from(slice.len()),
+            FormatElement::Space
+            | FormatElement::HardSpace
+            | FormatElement::Line(LineMode::SoftOrSpace) => 1,
+            _ => 0,
+        })
+        .sum()
+}
+
+fn is_in_test_call(node: &RArgument) -> bool {
+    node.parent::<RArgumentList>()
+        .and_then(|list| list.parent::<RCallArguments>())
+        .and_then(|args| args.parent::<RCall>())
+        .is_some_and(|call| call.is_test_call().unwrap_or(false))
+}
+
+fn is_call_with_long_callee(value: &AnyRExpression, f: &RFormatter) -> bool {
+    let callee = match value {
+        AnyRExpression::RCall(call) => call.function(),
+        AnyRExpression::RSubset(subset) => subset.function(),
+        AnyRExpression::RSubset2(subset2) => subset2.function(),
+        _ => return false,
+    };
+    callee.is_ok_and(|callee| {
+        usize::from(callee.syntax().text_trimmed().len()) + 1
+            > usize::from(f.options().line_width().value())
+    })
 }

crates/air_r_formatter/tests/specs/r/call.R

@@ -865,3 +865,117 @@ foo( bar[
   2
 ]
 )
+
+# ------------------------------------------------------------------------
+# Break after `=` - https://github.com/posit-dev/air/issues/527
+
+foo(
+  x,
+  long_argument_name = "a_very_long_string_value_that_exceeds_eighty_columns_when_on_one_line"
+)
+
+# Already broken after `=` remains idempotent
+foo(
+  x,
+  long_argument_name =
+    "a_very_long_string_value_that_exceeds_eighty_columns_when_on_one_line"
+)
+
+# Fits on one line once enclosing call expands (does not break after `=`)
+foo(medium_name = "a_string_that_fits_on_indented_line_once_foo_call_expands_ok")
+
+# Collapses newline after `=` when `name = value` fits on one line
+foo(
+  x =
+    "short_value"
+)
+
+# Short name with a long call expands call arguments rather than breaking at `=`
+foo(
+  x = another_function(
+    long_argument_one,
+    long_argument_two,
+    long_argument_three
+  )
+)
+
+# Pipe chain where `name = <lhs> |>` fits on the line stays on the `=` line
+process_data(
+  data = starwars |>
+    filter(species == "Human") |>
+    select(name, height, mass)
+)
+
+# Pipe chain where `name = <lhs> |>` exceeds 80 columns breaks after `=`
+process_data(
+  filtered_human_character_subset_from_starwars_dataset = starwars_character_data |>
+    filter(species == "Human") |>
+    select(name, height, mass)
+)
+
+# Formulas (`~`):
+# 1. Short name + multi-line formula keeps `lhs ~` on the `=` line
+# 2. Long name + single-line formula breaks after `=` and keeps formula flat
+# 3. Long name + multi-line formula where `name = <lhs> ~` > 80 breaks after `=`
+# 4. Long name + formula with RHS chain where `name = <lhs> ~ <rhs1> +` > 80
+fit_models(
+  short_formula = Surv(dm_py, diabetes) ~
+    AGE +
+    SEX,
+  proportional_hazards_regression_model_formula = Surv(dm_py, diabetes) ~ AGE + SEX,
+  proportional_hazards_regression_model_formula_specification = Surv(dm_py, diabetes) ~
+    AGE +
+    SEX,
+  proportional_hazards_regression_rhs_chain_formula = Surv(dm_py, diabetes) ~ AGE +
+    SEX
+)
+
+# Braced expressions, functions, and control flow never break after `=`
+foo(
+  long_argument_name_that_never_breaks_before_braces = {
+    x + 1
+  },
+  long_argument_name_that_never_breaks_before_anon_fn = function(a, b) a + b,
+  long_argument_name_that_never_breaks_before_lambda = \(a, b) a + b,
+  long_argument_name_that_never_breaks_before_if_stmt = if (cond) a else b,
+  long_argument_name_that_never_breaks_before_for_stmt = for (i in xs) print(i)
+)
+
+# Comments around `=` and argument value
+foo(
+  long_argument_name = # trailing comment on `=`
+    "a_very_long_string_value_that_exceeds_eighty_columns_when_on_one_line",
+  another_long_name =
+    # leading comment before value
+    "a_very_long_string_value_that_exceeds_eighty_columns_when_on_one_line"
+)
+
+check_inputs <- function(x, y, n, level, min.level, max.level, file_path, lookup_table, allow_null = FALSE) {
+  stopifnot(
+    is.character(x),
+    `Provide equal-length x and y vectors, or a single y value` = lengths_match(x, y) || length(y) == 1L,
+    "wait_done should be a logical scalar" =
+      IsScalarLogical(wait_done) && !is.na(wait_done),
+    "min.level must be an integer in [0, 30]." = length(min.level) == 1 && .IsValidLevel(min.level),
+    `n must be a non-missing numeric scalar (length 1) or Inf` = is.numeric(n) &&
+      length(n) == 1L &&
+      !is.na(n),
+    `Expected a custom_error condition with a matching diagnostic message` = inherits(x, "custom_error") &&
+      grepl("expected diagnostic message", x$message),
+    `file_path should be a single file name or an open connection` = is_character_or_connection(file_path),
+    "Level must be an integer between 0 and 30 inclusive." = .IsValidLevel(level),
+    "Level must not be an invalid integer outside 0 and 30." = !.IsValidLevel(level),
+    "Level bounds must be ordered and within valid range" = (min.level <= max.level && max.level <= 30L),
+    "Configuration object for file_path must have is_valid flag" = get_configuration(file_path)$is_valid,
+    `Configuration options must be valid for both x and y inputs` = is_valid_configuration(
+      x,
+      y,
+      allow_null = allow_null,
+      strict_mode = TRUE
+    ),
+    "file_path entry in lookup_table must be marked valid" = lookup_table[file_path, "valid", drop = TRUE],
+    `file_path must exist as an entry in the provided lookup_table` = lookup_table[[file_path]],
+    `allow_null is not supported; provide a non-null input value` = !allow_null,
+    `allow_null is not supported; please provide a non-null input value` = !allow_null
+  )
+}