Enable UTC-import for datetime-utc-alias fix

astral-sh · Jun 14, 2023 · 336396f · 336396f
1 parent 1350b8f
commit 336396f
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 15 deletions.
diff --git a/crates/ruff/src/rules/pyupgrade/rules/datetime_utc_alias.rs b/crates/ruff/src/rules/pyupgrade/rules/datetime_utc_alias.rs
@@ -2,9 +2,9 @@ use rustpython_parser::ast::{Expr, Ranged};
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::call_path::collect_call_path;
 
 use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
 use crate::registry::AsRule;
 
 #[violation]
@@ -34,17 +34,15 @@ pub(crate) fn datetime_utc_alias(checker: &mut Checker, expr: &Expr) {
     {
         let mut diagnostic = Diagnostic::new(DatetimeTimezoneUTC, expr.range());
         if checker.patch(diagnostic.kind.rule()) {
-            // If the reference was structured as, e.g., `datetime.timezone.utc`, then we can
-            // replace it with `datetime.UTC`. If `timezone` was imported via `from datetime import
-            // timezone`, then the replacement is more complicated.
-            if collect_call_path(expr).map_or(false, |call_path| {
-                matches!(call_path.as_slice(), ["datetime", "timezone", "utc"])
-            }) {
-                diagnostic.set_fix(Fix::suggested(Edit::range_replacement(
-                    "datetime.UTC".to_string(),
-                    expr.range(),
-                )));
-            }
+            diagnostic.try_set_fix(|| {
+                let (import_edit, binding) = checker.importer.get_or_import_symbol(
+                    &ImportRequest::import_from("datetime", "UTC"),
+                    expr.start(),
+                    checker.semantic(),
+                )?;
+                let reference_edit = Edit::range_replacement(binding, expr.range());
+                Ok(Fix::suggested_edits(import_edit, [reference_edit]))
+            });
         }
         checker.diagnostics.push(diagnostic);
     }

diff --git a/...rc/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__datetime_utc_alias_py311.snap b/...rc/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__datetime_utc_alias_py311.snap
@@ -1,7 +1,7 @@
 ---
 source: crates/ruff/src/rules/pyupgrade/mod.rs
 ---
-UP017.py:7:7: UP017 Use `datetime.UTC` alias
+UP017.py:7:7: UP017 [*] Use `datetime.UTC` alias
   |
 6 | print(datetime.timezone(-1))
 7 | print(timezone.utc)
@@ -10,7 +10,17 @@ UP017.py:7:7: UP017 Use `datetime.UTC` alias
   |
   = help: Convert to `datetime.UTC` alias
 
-UP017.py:8:7: UP017 Use `datetime.UTC` alias
+ℹ Suggested fix
+4 4 | from datetime import timezone as tz
+5 5 | 
+6 6 | print(datetime.timezone(-1))
+7   |-print(timezone.utc)
+  7 |+print(datetime.UTC)
+8 8 | print(tz.utc)
+9 9 | 
+10 10 | print(datetime.timezone.utc)
+
+UP017.py:8:7: UP017 [*] Use `datetime.UTC` alias
    |
  6 | print(datetime.timezone(-1))
  7 | print(timezone.utc)
@@ -21,6 +31,16 @@ UP017.py:8:7: UP017 Use `datetime.UTC` alias
    |
    = help: Convert to `datetime.UTC` alias
 
+ℹ Suggested fix
+5 5 | 
+6 6 | print(datetime.timezone(-1))
+7 7 | print(timezone.utc)
+8   |-print(tz.utc)
+  8 |+print(datetime.UTC)
+9 9 | 
+10 10 | print(datetime.timezone.utc)
+11 11 | print(dt.timezone.utc)
+
 UP017.py:10:7: UP017 [*] Use `datetime.UTC` alias
    |
  8 | print(tz.utc)
@@ -39,12 +59,19 @@ UP017.py:10:7: UP017 [*] Use `datetime.UTC` alias
    10 |+print(datetime.UTC)
 11 11 | print(dt.timezone.utc)
 
-UP017.py:11:7: UP017 Use `datetime.UTC` alias
+UP017.py:11:7: UP017 [*] Use `datetime.UTC` alias
    |
 10 | print(datetime.timezone.utc)
 11 | print(dt.timezone.utc)
    |       ^^^^^^^^^^^^^^^ UP017
    |
    = help: Convert to `datetime.UTC` alias
 
+ℹ Suggested fix
+8  8  | print(tz.utc)
+9  9  | 
+10 10 | print(datetime.timezone.utc)
+11    |-print(dt.timezone.utc)
+   11 |+print(datetime.UTC)
+