Implement "datetime.UTC alias" check from pyupgrade (#1341)

README.md

@@ -628,6 +628,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP014 | ConvertNamedTupleFunctionalToClass | Convert `...` from `NamedTuple` functional to class syntax | 🛠 |
 | UP015 | RedundantOpenModes | Unnecessary open mode parameters | 🛠 |
 | UP016 | RemoveSixCompat | Unnecessary `six` compatibility usage | 🛠 |
+| UP017 | DatetimeTimezoneUTC | Use `datetime.UTC` alias | 🛠 |
 
 ### pep8-naming (N)
 
@@ -1224,7 +1225,7 @@ natively, including:
 - [`pep8-naming`](https://pypi.org/project/pep8-naming/)
 - [`pydocstyle`](https://pypi.org/project/pydocstyle/)
 - [`pygrep-hooks`](https://github.com/pre-commit/pygrep-hooks) (3/10)
-- [`pyupgrade`](https://pypi.org/project/pyupgrade/) (16/33)
+- [`pyupgrade`](https://pypi.org/project/pyupgrade/) (17/33)
 - [`yesqa`](https://github.com/asottile/yesqa)
 
 Note that, in some cases, Ruff uses different error code prefixes than would be found in the
@@ -1281,7 +1282,7 @@ Today, Ruff can be used to replace Flake8 when used with any of the following pl
 Ruff can also replace [`isort`](https://pypi.org/project/isort/),
 [`yesqa`](https://github.com/asottile/yesqa), [`eradicate`](https://pypi.org/project/eradicate/),
 [`pygrep-hooks`](https://github.com/pre-commit/pygrep-hooks) (3/10), and a subset of the rules
-implemented in [`pyupgrade`](https://pypi.org/project/pyupgrade/) (16/33).
+implemented in [`pyupgrade`](https://pypi.org/project/pyupgrade/) (17/33).
 
 If you're looking to use Ruff, but rely on an unsupported Flake8 plugin, free to file an Issue.
 

resources/test/fixtures/pyupgrade/UP017.py

@@ -0,0 +1,11 @@
+import datetime
+import datetime as dt
+from datetime import timezone
+from datetime import timezone as tz
+
+print(datetime.timezone(-1))
+print(timezone.utc)
+print(tz.utc)
+
+print(datetime.timezone.utc)
+print(dt.timezone.utc)

src/checkers/ast.rs

@@ -1553,6 +1553,12 @@ where
                     pyupgrade::plugins::remove_six_compat(self, expr);
                 }
 
+                if self.settings.enabled.contains(&CheckCode::UP017)
+                    && self.settings.target_version >= PythonVersion::Py311
+                {
+                    pyupgrade::plugins::datetime_utc_alias(self, expr);
+                }
+
                 if self.settings.enabled.contains(&CheckCode::YTT202) {
                     flake8_2020::plugins::name_or_attribute(self, expr);
                 }

src/checks.rs

@@ -225,6 +225,7 @@ pub enum CheckCode {
     UP014,
     UP015,
     UP016,
+    UP017,
     // pydocstyle
     D100,
     D101,
@@ -827,6 +828,7 @@ pub enum CheckKind {
     ConvertNamedTupleFunctionalToClass(String),
     RedundantOpenModes,
     RemoveSixCompat,
+    DatetimeTimezoneUTC,
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1200,6 +1202,7 @@ impl CheckCode {
             CheckCode::UP014 => CheckKind::ConvertNamedTupleFunctionalToClass("...".to_string()),
             CheckCode::UP015 => CheckKind::RedundantOpenModes,
             CheckCode::UP016 => CheckKind::RemoveSixCompat,
+            CheckCode::UP017 => CheckKind::DatetimeTimezoneUTC,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1617,6 +1620,7 @@ impl CheckCode {
             CheckCode::UP014 => CheckCategory::Pyupgrade,
             CheckCode::UP015 => CheckCategory::Pyupgrade,
             CheckCode::UP016 => CheckCategory::Pyupgrade,
+            CheckCode::UP017 => CheckCategory::Pyupgrade,
             CheckCode::W292 => CheckCategory::Pycodestyle,
             CheckCode::W605 => CheckCategory::Pycodestyle,
             CheckCode::YTT101 => CheckCategory::Flake82020,
@@ -1827,6 +1831,7 @@ impl CheckKind {
             CheckKind::ConvertNamedTupleFunctionalToClass(_) => &CheckCode::UP014,
             CheckKind::RedundantOpenModes => &CheckCode::UP015,
             CheckKind::RemoveSixCompat => &CheckCode::UP016,
+            CheckKind::DatetimeTimezoneUTC => &CheckCode::UP017,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(_) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(_) => &CheckCode::D410,
@@ -2549,6 +2554,7 @@ impl CheckKind {
             CheckKind::UnnecessaryEncodeUTF8 => "Unnecessary call to `encode` as UTF-8".to_string(),
             CheckKind::RedundantOpenModes => "Unnecessary open mode parameters".to_string(),
             CheckKind::RemoveSixCompat => "Unnecessary `six` compatibility usage".to_string(),
+            CheckKind::DatetimeTimezoneUTC => "Use `datetime.UTC` alias".to_string(),
             CheckKind::ConvertTypedDictFunctionalToClass(name) => {
                 format!("Convert `{name}` from `TypedDict` functional to class syntax")
             }
@@ -2985,6 +2991,7 @@ impl CheckKind {
                 | CheckKind::RedundantOpenModes
                 | CheckKind::RedundantTupleInExceptionHandler(..)
                 | CheckKind::RemoveSixCompat
+                | CheckKind::DatetimeTimezoneUTC
                 | CheckKind::SectionNameEndsInColon(..)
                 | CheckKind::SectionNotOverIndented(..)
                 | CheckKind::SectionUnderlineAfterName(..)
@@ -3068,6 +3075,7 @@ pub static PREFIX_REDIRECTS: Lazy<FxHashMap<&'static str, CheckCodePrefix>> = La
         ("U014", CheckCodePrefix::UP014),
         ("U015", CheckCodePrefix::UP015),
         ("U016", CheckCodePrefix::UP016),
+        ("U017", CheckCodePrefix::UP017),
         // TODO(charlie): Remove by 2023-02-01.
         ("I252", CheckCodePrefix::TID252),
         ("M001", CheckCodePrefix::RUF100),
@@ -3142,6 +3150,7 @@ pub static CODE_REDIRECTS: Lazy<FxHashMap<&'static str, CheckCode>> = Lazy::new(
         ("U014", CheckCode::UP014),
         ("U015", CheckCode::UP015),
         ("U016", CheckCode::UP016),
+        ("U017", CheckCode::UP017),
         // TODO(charlie): Remove by 2023-02-01.
         ("I252", CheckCode::TID252),
         ("M001", CheckCode::RUF100),

src/checks_gen.rs

@@ -492,6 +492,7 @@ pub enum CheckCodePrefix {
     U014,
     U015,
     U016,
+    U017,
     UP,
     UP0,
     UP00,
@@ -511,6 +512,7 @@ pub enum CheckCodePrefix {
     UP014,
     UP015,
     UP016,
+    UP017,
     W,
     W2,
     W29,
@@ -2085,6 +2087,7 @@ impl CheckCodePrefix {
                     CheckCode::UP014,
                     CheckCode::UP015,
                     CheckCode::UP016,
+                    CheckCode::UP017,
                 ]
             }
             CheckCodePrefix::U0 => {
@@ -2110,6 +2113,7 @@ impl CheckCodePrefix {
                     CheckCode::UP014,
                     CheckCode::UP015,
                     CheckCode::UP016,
+                    CheckCode::UP017,
                 ]
             }
             CheckCodePrefix::U00 => {
@@ -2217,6 +2221,7 @@ impl CheckCodePrefix {
                     CheckCode::UP014,
                     CheckCode::UP015,
                     CheckCode::UP016,
+                    CheckCode::UP017,
                 ]
             }
             CheckCodePrefix::U010 => {
@@ -2282,6 +2287,15 @@ impl CheckCodePrefix {
                 );
                 vec![CheckCode::UP016]
             }
+            CheckCodePrefix::U017 => {
+                one_time_warning!(
+                    "{}{} {}",
+                    "warning".yellow().bold(),
+                    ":".bold(),
+                    "`U017` has been remapped to `UP017`".bold()
+                );
+                vec![CheckCode::UP017]
+            }
             CheckCodePrefix::UP => vec![
                 CheckCode::UP001,
                 CheckCode::UP003,
@@ -2298,6 +2312,7 @@ impl CheckCodePrefix {
                 CheckCode::UP014,
                 CheckCode::UP015,
                 CheckCode::UP016,
+                CheckCode::UP017,
             ],
             CheckCodePrefix::UP0 => vec![
                 CheckCode::UP001,
@@ -2315,6 +2330,7 @@ impl CheckCodePrefix {
                 CheckCode::UP014,
                 CheckCode::UP015,
                 CheckCode::UP016,
+                CheckCode::UP017,
             ],
             CheckCodePrefix::UP00 => vec![
                 CheckCode::UP001,
@@ -2342,6 +2358,7 @@ impl CheckCodePrefix {
                 CheckCode::UP014,
                 CheckCode::UP015,
                 CheckCode::UP016,
+                CheckCode::UP017,
             ],
             CheckCodePrefix::UP010 => vec![CheckCode::UP010],
             CheckCodePrefix::UP011 => vec![CheckCode::UP011],
@@ -2350,6 +2367,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP014 => vec![CheckCode::UP014],
             CheckCodePrefix::UP015 => vec![CheckCode::UP015],
             CheckCodePrefix::UP016 => vec![CheckCode::UP016],
+            CheckCodePrefix::UP017 => vec![CheckCode::UP017],
             CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
             CheckCodePrefix::W2 => vec![CheckCode::W292],
             CheckCodePrefix::W29 => vec![CheckCode::W292],
@@ -2884,6 +2902,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::U014 => SuffixLength::Three,
             CheckCodePrefix::U015 => SuffixLength::Three,
             CheckCodePrefix::U016 => SuffixLength::Three,
+            CheckCodePrefix::U017 => SuffixLength::Three,
             CheckCodePrefix::UP => SuffixLength::Zero,
             CheckCodePrefix::UP0 => SuffixLength::One,
             CheckCodePrefix::UP00 => SuffixLength::Two,
@@ -2903,6 +2922,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP014 => SuffixLength::Three,
             CheckCodePrefix::UP015 => SuffixLength::Three,
             CheckCodePrefix::UP016 => SuffixLength::Three,
+            CheckCodePrefix::UP017 => SuffixLength::Three,
             CheckCodePrefix::W => SuffixLength::Zero,
             CheckCodePrefix::W2 => SuffixLength::One,
             CheckCodePrefix::W29 => SuffixLength::Two,

src/pyupgrade/mod.rs

@@ -105,4 +105,18 @@ mod tests {
         insta::assert_yaml_snapshot!(checks);
         Ok(())
     }
+
+    #[test]
+    fn datetime_utc_alias_py311() -> Result<()> {
+        let mut checks = test_path(
+            Path::new("./resources/test/fixtures/pyupgrade/UP017.py"),
+            &settings::Settings {
+                target_version: PythonVersion::Py311,
+                ..settings::Settings::for_rule(CheckCode::UP017)
+            },
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
 }

src/pyupgrade/plugins/datetime_utc_alias.rs

@@ -0,0 +1,25 @@
+use rustpython_ast::Expr;
+
+use crate::ast::helpers::{collect_call_paths, compose_call_path, dealias_call_path};
+use crate::ast::types::Range;
+use crate::autofix::Fix;
+use crate::checkers::ast::Checker;
+use crate::checks::{Check, CheckCode, CheckKind};
+
+/// UP017
+pub fn datetime_utc_alias(checker: &mut Checker, expr: &Expr) {
+    let dealiased_call_path = dealias_call_path(collect_call_paths(expr), &checker.import_aliases);
+    if dealiased_call_path == ["datetime", "timezone", "utc"] {
+        let mut check = Check::new(CheckKind::DatetimeTimezoneUTC, Range::from_located(expr));
+        if checker.patch(&CheckCode::UP017) {
+            check.amend(Fix::replacement(
+                compose_call_path(expr)
+                    .unwrap()
+                    .replace("timezone.utc", "UTC"),
+                expr.location,
+                expr.end_location.unwrap(),
+            ));
+        }
+        checker.add_check(check);
+    }
+}

src/pyupgrade/plugins/mod.rs

@@ -1,5 +1,6 @@
 pub use convert_named_tuple_functional_to_class::convert_named_tuple_functional_to_class;
 pub use convert_typed_dict_functional_to_class::convert_typed_dict_functional_to_class;
+pub use datetime_utc_alias::datetime_utc_alias;
 pub use deprecated_unittest_alias::deprecated_unittest_alias;
 pub use redundant_open_modes::redundant_open_modes;
 pub use remove_six_compat::remove_six_compat;
@@ -15,6 +16,7 @@ pub use useless_object_inheritance::useless_object_inheritance;
 
 mod convert_named_tuple_functional_to_class;
 mod convert_typed_dict_functional_to_class;
+mod datetime_utc_alias;
 mod deprecated_unittest_alias;
 mod redundant_open_modes;
 mod remove_six_compat;

src/pyupgrade/snapshots/ruff__pyupgrade__tests__datetime_utc_alias_py311.snap

@@ -0,0 +1,35 @@
+---
+source: src/pyupgrade/mod.rs
+expression: checks
+---
+- kind: DatetimeTimezoneUTC
+  location:
+    row: 10
+    column: 6
+  end_location:
+    row: 10
+    column: 27
+  fix:
+    content: datetime.UTC
+    location:
+      row: 10
+      column: 6
+    end_location:
+      row: 10
+      column: 27
+- kind: DatetimeTimezoneUTC
+  location:
+    row: 11
+    column: 6
+  end_location:
+    row: 11
+    column: 21
+  fix:
+    content: dt.UTC
+    location:
+      row: 11
+      column: 6
+    end_location:
+      row: 11
+      column: 21
+