Enable UTC-import for `datetime-utc-alias` fix (#5100)

## Summary

Small update to leverage `get_or_import_symbol` to fix `UP017` in more
cases (e.g., when we need to import `UTC`, or access it from an alias or
something).

## Test Plan

Check out the updated snapshot.
This commit is contained in:
Charlie Marsh 2023-06-14 17:13:36 -04:00 committed by GitHub
parent 56476dfd61
commit 848f184b8c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 40 additions and 15 deletions

View File

@ -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);
}

View File

@ -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)