[ruff] Implement post-init-default (RUF033) (#13192)

Co-authored-by: Alex Waygood <[email protected]>

crates/ruff_linter/resources/test/fixtures/ruff/RUF033.py

@@ -0,0 +1,67 @@
+from dataclasses import InitVar, dataclass
+
+
+# OK
+@dataclass
+class Foo:
+    bar: InitVar[int] = 0
+    baz: InitVar[int] = 1
+
+    def __post_init__(self, bar, baz) -> None: ...
+
+
+# RUF033
+@dataclass
+class Foo:
+    bar: InitVar[int] = 0
+    baz: InitVar[int] = 1
+
+    def __post_init__(self, bar = 11, baz = 11) -> None: ...
+
+
+# RUF033
+@dataclass
+class Foo:
+    def __post_init__(self, bar = 11, baz = 11) -> None: ...
+
+
+# OK
+@dataclass
+class Foo:
+    def __something_else__(self, bar = 11, baz = 11) -> None: ...
+
+
+# OK
+def __post_init__(foo: bool = True) -> None: ...
+
+
+# OK
+class Foo:
+    def __post_init__(self, x="hmm") -> None: ...
+
+
+# RUF033
+@dataclass
+class Foo:
+    def __post_init__(self, bar: int = 11, baz: Something[Whatever | None] = 11) -> None: ...
+
+
+# RUF033
+@dataclass
+class Foo:
+    """A very helpful docstring.
+
+    Docstrings are very important and totally not a waste of time.
+    """
+
+    ping = "pong"
+
+    def __post_init__(self, bar: int = 11, baz: int = 12) -> None: ...
+
+
+# RUF033
+@dataclass
+class Foo:
+    bar = "should've used attrs"
+
+    def __post_init__(self, bar: str = "ahhh", baz: str = "hmm") -> None: ...

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -380,6 +380,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::WhitespaceAfterDecorator) {
                 pycodestyle::rules::whitespace_after_decorator(checker, decorator_list);
             }
+            if checker.enabled(Rule::PostInitDefault) {
+                ruff::rules::post_init_default(checker, function_def);
+            }
         }
         Stmt::Return(_) => {
             if checker.enabled(Rule::ReturnOutsideFunction) {

crates/ruff_linter/src/codes.rs

@@ -960,6 +960,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Ruff, "030") => (RuleGroup::Preview, rules::ruff::rules::AssertWithPrintMessage),
         (Ruff, "031") => (RuleGroup::Preview, rules::ruff::rules::IncorrectlyParenthesizedTupleInSubscript),
         (Ruff, "032") => (RuleGroup::Preview, rules::ruff::rules::DecimalFromFloatLiteral),
+        (Ruff, "033") => (RuleGroup::Preview, rules::ruff::rules::PostInitDefault),
         (Ruff, "100") => (RuleGroup::Stable, rules::ruff::rules::UnusedNOQA),
         (Ruff, "101") => (RuleGroup::Stable, rules::ruff::rules::RedirectedNOQA),
 

crates/ruff_linter/src/rules/ruff/mod.rs

@@ -59,6 +59,7 @@ mod tests {
     #[test_case(Rule::IncorrectlyParenthesizedTupleInSubscript, Path::new("RUF031.py"))]
     #[test_case(Rule::DecimalFromFloatLiteral, Path::new("RUF032.py"))]
     #[test_case(Rule::RedirectedNOQA, Path::new("RUF101.py"))]
+    #[test_case(Rule::PostInitDefault, Path::new("RUF033.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

crates/ruff_linter/src/rules/ruff/rules/mod.rs

@@ -74,3 +74,6 @@ pub(crate) enum Context {
     Docstring,
     Comment,
 }
+pub(crate) use post_init_default::*;
+
+mod post_init_default;

crates/ruff_linter/src/rules/ruff/rules/post_init_default.rs

@@ -0,0 +1,187 @@
+use anyhow::Context;
+
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast as ast;
+use ruff_python_semantic::{Scope, ScopeKind};
+use ruff_python_trivia::{indentation_at_offset, textwrap};
+use ruff_text_size::Ranged;
+
+use crate::{checkers::ast::Checker, importer::ImportRequest};
+
+use super::helpers::is_dataclass;
+
+/// ## What it does
+/// Checks for `__post_init__` dataclass methods with parameter defaults.
+///
+/// ## Why is this bad?
+/// Adding a default value to a parameter in a `__post_init__` method has no
+/// impact on whether the parameter will have a default value in the dataclass's
+/// generated `__init__` method. To create an init-only dataclass parameter with
+/// a default value, you should use an `InitVar` field in the dataclass's class
+/// body and give that `InitVar` field a default value.
+///
+/// As the [documentation] states:
+///
+/// > Init-only fields are added as parameters to the generated `__init__()`
+/// > method, and are passed to the optional `__post_init__()` method. They are
+/// > not otherwise used by dataclasses.
+///
+/// ## Example
+/// ```python
+/// from dataclasses import InitVar, dataclass
+///
+///
+/// @dataclass
+/// class Foo:
+///     bar: InitVar[int] = 0
+///
+///     def __post_init__(self, bar: int = 1, baz: int = 2) -> None:
+///         print(bar, baz)
+///
+///
+/// foo = Foo()  # Prints '0 2'.
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from dataclasses import InitVar, dataclass
+///
+///
+/// @dataclass
+/// class Foo:
+///     bar: InitVar[int] = 1
+///     baz: InitVar[int] = 2
+///
+///     def __post_init__(self, bar: int, baz: int) -> None:
+///         print(bar, baz)
+///
+///
+/// foo = Foo()  # Prints '1 2'.
+/// ```
+///
+/// ## References
+/// - [Python documentation: Post-init processing](https://docs.python.org/3/library/dataclasses.html#post-init-processing)
+/// - [Python documentation: Init-only variables](https://docs.python.org/3/library/dataclasses.html#init-only-variables)
+///
+/// [documentation]: https://docs.python.org/3/library/dataclasses.html#init-only-variables
+#[violation]
+pub struct PostInitDefault;
+
+impl Violation for PostInitDefault {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("`__post_init__` method with argument defaults")
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        Some(format!("Use `dataclasses.InitVar` instead"))
+    }
+}
+
+/// RUF033
+pub(crate) fn post_init_default(checker: &mut Checker, function_def: &ast::StmtFunctionDef) {
+    if &function_def.name != "__post_init__" {
+        return;
+    }
+
+    let current_scope = checker.semantic().current_scope();
+    match current_scope.kind {
+        ScopeKind::Class(class_def) => {
+            if !is_dataclass(class_def, checker.semantic()) {
+                return;
+            }
+        }
+        _ => return,
+    }
+
+    let mut stopped_fixes = false;
+    let mut diagnostics = vec![];
+
+    for ast::ParameterWithDefault {
+        parameter,
+        default,
+        range: _,
+    } in function_def.parameters.iter_non_variadic_params()
+    {
+        let Some(default) = default else {
+            continue;
+        };
+        let mut diagnostic = Diagnostic::new(PostInitDefault, default.range());
+
+        if !stopped_fixes {
+            diagnostic.try_set_fix(|| {
+                use_initvar(current_scope, function_def, parameter, default, checker)
+            });
+            // Need to stop fixes as soon as there is a parameter we cannot fix.
+            // Otherwise, we risk a syntax error (a parameter without a default
+            // following parameter with a default).
+            stopped_fixes |= diagnostic.fix.is_none();
+        }
+
+        diagnostics.push(diagnostic);
+    }
+
+    checker.diagnostics.extend(diagnostics);
+}
+
+/// Generate a [`Fix`] to transform a `__post_init__` default argument into a
+/// `dataclasses.InitVar` pseudo-field.
+fn use_initvar(
+    current_scope: &Scope,
+    post_init_def: &ast::StmtFunctionDef,
+    parameter: &ast::Parameter,
+    default: &ast::Expr,
+    checker: &Checker,
+) -> anyhow::Result<Fix> {
+    if current_scope.has(&parameter.name) {
+        return Err(anyhow::anyhow!(
+            "Cannot add a `{}: InitVar` field to the class body, as a field by that name already exists",
+            parameter.name
+        ));
+    }
+
+    // Ensure that `dataclasses.InitVar` is accessible. For example,
+    // + `from dataclasses import InitVar`
+    let (import_edit, initvar_binding) = checker.importer().get_or_import_symbol(
+        &ImportRequest::import("dataclasses", "InitVar"),
+        default.start(),
+        checker.semantic(),
+    )?;
+
+    // Delete the default value. For example,
+    // - def __post_init__(self, foo: int = 0) -> None: ...
+    // + def __post_init__(self, foo: int) -> None: ...
+    let default_edit = Edit::deletion(parameter.end(), default.end());
+
+    // Add `dataclasses.InitVar` field to class body.
+    let locator = checker.locator();
+
+    let content = {
+        let parameter_name = locator.slice(&parameter.name);
+        let default = locator.slice(default);
+        let line_ending = checker.stylist().line_ending().as_str();
+
+        if let Some(annotation) = &parameter
+            .annotation
+            .as_deref()
+            .map(|annotation| locator.slice(annotation))
+        {
+            format!("{parameter_name}: {initvar_binding}[{annotation}] = {default}{line_ending}")
+        } else {
+            format!("{parameter_name}: {initvar_binding} = {default}{line_ending}")
+        }
+    };
+
+    let indentation = indentation_at_offset(post_init_def.start(), checker.locator())
+        .context("Failed to calculate leading indentation of `__post_init__` method")?;
+    let content = textwrap::indent(&content, indentation);
+
+    let initvar_edit = Edit::insertion(
+        content.into_owned(),
+        locator.line_start(post_init_def.start()),
+    );
+    Ok(Fix::unsafe_edits(import_edit, [default_edit, initvar_edit]))
+}