From 7dade46f5b45bf2f4bfd824e9246dd511eeb42fb Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 14 Jun 2023 15:38:42 -0400 Subject: [PATCH] Enable UTC-import for datetime-utc-alias fix --- .../pyupgrade/rules/datetime_utc_alias.rs | 22 ++++++------- ...rade__tests__datetime_utc_alias_py311.snap | 33 +++++++++++++++++-- 2 files changed, 40 insertions(+), 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 index 9840c8086d630..5083269d1dc21 100644 --- 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/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__datetime_utc_alias_py311.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__datetime_utc_alias_py311.snap index f6a18934b07f9..7c1bbecb220bb 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__datetime_utc_alias_py311.snap +++ b/crates/ruff/src/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,7 +59,7 @@ 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) @@ -47,4 +67,11 @@ UP017.py:11:7: UP017 Use `datetime.UTC` alias | = 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) +