refactor(use-from-import): build fixed variant via AST (#3132)

astral-sh · Feb 22, 2023 · 6ced512 · 6ced512
1 parent 7d55b41
commit 6ced512
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/crates/ruff/src/rules/pylint/rules/use_from_import.rs b/crates/ruff/src/rules/pylint/rules/use_from_import.rs
@@ -1,6 +1,7 @@
 use ruff_macros::{define_violation, derive_message_formats};
-use rustpython_parser::ast::{Alias, Stmt};
+use rustpython_parser::ast::{Alias, AliasData, Located, Stmt, StmtKind};
 
+use crate::ast::helpers::{create_stmt, unparse_stmt};
 use crate::ast::types::Range;
 use crate::checkers::ast::Checker;
 use crate::fix::Fix;
@@ -58,7 +59,21 @@ pub fn use_from_import(checker: &mut Checker, stmt: &Stmt, alias: &Alias, names:
     );
     if fixable && checker.patch(diagnostic.kind.rule()) {
         diagnostic.amend(Fix::replacement(
-            format!("from {module} import {asname}"),
+            unparse_stmt(
+                &create_stmt(StmtKind::ImportFrom {
+                    module: Some(module.to_string()),
+                    names: vec![Located::new(
+                        stmt.location,
+                        stmt.end_location.unwrap(),
+                        AliasData {
+                            name: asname.into(),
+                            asname: None,
+                        },
+                    )],
+                    level: Some(0),
+                }),
+                checker.stylist,
+            ),
             stmt.location,
             stmt.end_location.unwrap(),
         ));