Avoid triggering DTZ001-006 when using .astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ001.py

@@ -19,3 +19,6 @@
 
 # no args unqualified
 datetime(2000, 1, 1, 0, 0, 0)
+
+# uses `astimezone` method
+datetime(2000, 1, 1, 0, 0, 0).astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ002.py

@@ -7,3 +7,6 @@
 
 # unqualified
 datetime.today()
+
+# uses `astimezone` method
+datetime.today().astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ003.py

@@ -7,3 +7,6 @@
 
 # unqualified
 datetime.utcnow()
+
+# uses `astimezone` method
+datetime.utcnow().astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ004.py

@@ -7,3 +7,6 @@
 
 # unqualified
 datetime.utcfromtimestamp(1234)
+
+# uses `astimezone` method
+datetime.utcfromtimestamp(1234).astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ005.py

@@ -16,3 +16,6 @@
 
 # no args unqualified
 datetime.now()
+
+# uses `astimezone` method
+datetime.now().astimezone()

crates/ruff/resources/test/fixtures/flake8_datetimez/DTZ006.py

@@ -16,3 +16,6 @@
 
 # no args unqualified
 datetime.fromtimestamp(1234)
+
+# uses `astimezone` method
+datetime.fromtimestamp(1234).astimezone()

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_fromtimestamp.rs

@@ -7,6 +7,8 @@ use ruff_python_ast::helpers::{has_non_none_keyword, is_const_none};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeFromtimestamp;
 
@@ -40,6 +42,10 @@ pub(crate) fn call_datetime_fromtimestamp(
         return;
     }
 
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
+    }
+
     // no args / no args unqualified
     if args.len() < 2 && keywords.is_empty() {
         checker

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_now_without_tzinfo.rs

@@ -7,6 +7,8 @@ use ruff_python_ast::helpers::{has_non_none_keyword, is_const_none};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeNowWithoutTzinfo;
 
@@ -35,6 +37,10 @@ pub(crate) fn call_datetime_now_without_tzinfo(
         return;
     }
 
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
+    }
+
     // no args / no args unqualified
     if args.is_empty() && keywords.is_empty() {
         checker

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_today.rs

@@ -6,6 +6,8 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeToday;
 
@@ -26,15 +28,21 @@ impl Violation for CallDatetimeToday {
 /// It uses the system local timezone.
 /// Use `datetime.datetime.now(tz=)` instead.
 pub(crate) fn call_datetime_today(checker: &mut Checker, func: &Expr, location: TextRange) {
-    if checker
+    if !checker
         .semantic()
         .resolve_call_path(func)
         .map_or(false, |call_path| {
             matches!(call_path.as_slice(), ["datetime", "datetime", "today"])
         })
     {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(CallDatetimeToday, location));
+        return;
+    }
+
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
     }
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(CallDatetimeToday, location));
 }

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_utcfromtimestamp.rs

@@ -6,6 +6,8 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeUtcfromtimestamp;
 
@@ -33,7 +35,7 @@ pub(crate) fn call_datetime_utcfromtimestamp(
     func: &Expr,
     location: TextRange,
 ) {
-    if checker
+    if !checker
         .semantic()
         .resolve_call_path(func)
         .map_or(false, |call_path| {
@@ -43,8 +45,14 @@ pub(crate) fn call_datetime_utcfromtimestamp(
             )
         })
     {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(CallDatetimeUtcfromtimestamp, location));
+        return;
+    }
+
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
     }
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(CallDatetimeUtcfromtimestamp, location));
 }

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_utcnow.rs

@@ -6,6 +6,8 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeUtcnow;
 
@@ -28,15 +30,21 @@ impl Violation for CallDatetimeUtcnow {
 /// UTC. As such, the recommended way to create an object representing the
 /// current time in UTC is by calling `datetime.now(timezone.utc)`.
 pub(crate) fn call_datetime_utcnow(checker: &mut Checker, func: &Expr, location: TextRange) {
-    if checker
+    if !checker
         .semantic()
         .resolve_call_path(func)
         .map_or(false, |call_path| {
             matches!(call_path.as_slice(), ["datetime", "datetime", "utcnow"])
         })
     {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(CallDatetimeUtcnow, location));
+        return;
+    }
+
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
     }
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(CallDatetimeUtcnow, location));
 }

crates/ruff/src/rules/flake8_datetimez/rules/call_datetime_without_tzinfo.rs

@@ -7,6 +7,8 @@ use ruff_python_ast::helpers::{has_non_none_keyword, is_const_none};
 
 use crate::checkers::ast::Checker;
 
+use super::helpers;
+
 #[violation]
 pub struct CallDatetimeWithoutTzinfo;
 
@@ -34,6 +36,10 @@ pub(crate) fn call_datetime_without_tzinfo(
         return;
     }
 
+    if helpers::parent_expr_is_astimezone(checker) {
+        return;
+    }
+
     // No positional arg: keyword is missing or constant None.
     if args.len() < 8 && !has_non_none_keyword(keywords, "tzinfo") {
         checker

crates/ruff/src/rules/flake8_datetimez/rules/helpers.rs

@@ -0,0 +1,11 @@
+use rustpython_parser::ast::{Expr, ExprAttribute};
+
+use crate::checkers::ast::Checker;
+
+/// Check if the parent expression is a call to `astimezone`. This assumes that
+/// the current expression is a `datetime.datetime` object.
+pub(crate) fn parent_expr_is_astimezone(checker: &Checker) -> bool {
+    checker.semantic().expr_parent().map_or(false, |parent| {
+        matches!(parent, Expr::Attribute(ExprAttribute { attr, .. }) if attr.as_str() == "astimezone")
+    })
+}

crates/ruff/src/rules/flake8_datetimez/rules/mod.rs

@@ -17,3 +17,4 @@ mod call_datetime_today;
 mod call_datetime_utcfromtimestamp;
 mod call_datetime_utcnow;
 mod call_datetime_without_tzinfo;
+mod helpers;

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ001_DTZ001.py.snap

@@ -42,6 +42,8 @@ DTZ001.py:21:1: DTZ001 The use of `datetime.datetime()` without `tzinfo` argumen
 20 | # no args unqualified
 21 | datetime(2000, 1, 1, 0, 0, 0)
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DTZ001
+22 | 
+23 | # uses `astimezone` method
    |
 
 

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ002_DTZ002.py.snap

@@ -11,10 +11,12 @@ DTZ002.py:4:1: DTZ002 The use of `datetime.datetime.today()` is not allowed, use
   |
 
 DTZ002.py:9:1: DTZ002 The use of `datetime.datetime.today()` is not allowed, use `datetime.datetime.now(tz=)` instead
-  |
-8 | # unqualified
-9 | datetime.today()
-  | ^^^^^^^^^^^^^^^^ DTZ002
-  |
+   |
+ 8 | # unqualified
+ 9 | datetime.today()
+   | ^^^^^^^^^^^^^^^^ DTZ002
+10 | 
+11 | # uses `astimezone` method
+   |
 
 

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ003_DTZ003.py.snap

@@ -11,10 +11,12 @@ DTZ003.py:4:1: DTZ003 The use of `datetime.datetime.utcnow()` is not allowed, us
   |
 
 DTZ003.py:9:1: DTZ003 The use of `datetime.datetime.utcnow()` is not allowed, use `datetime.datetime.now(tz=)` instead
-  |
-8 | # unqualified
-9 | datetime.utcnow()
-  | ^^^^^^^^^^^^^^^^^ DTZ003
-  |
+   |
+ 8 | # unqualified
+ 9 | datetime.utcnow()
+   | ^^^^^^^^^^^^^^^^^ DTZ003
+10 | 
+11 | # uses `astimezone` method
+   |
 
 

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ004_DTZ004.py.snap

@@ -11,10 +11,12 @@ DTZ004.py:4:1: DTZ004 The use of `datetime.datetime.utcfromtimestamp()` is not a
   |
 
 DTZ004.py:9:1: DTZ004 The use of `datetime.datetime.utcfromtimestamp()` is not allowed, use `datetime.datetime.fromtimestamp(ts, tz=)` instead
-  |
-8 | # unqualified
-9 | datetime.utcfromtimestamp(1234)
-  | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DTZ004
-  |
+   |
+ 8 | # unqualified
+ 9 | datetime.utcfromtimestamp(1234)
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DTZ004
+10 | 
+11 | # uses `astimezone` method
+   |
 
 

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ005_DTZ005.py.snap

@@ -42,6 +42,8 @@ DTZ005.py:18:1: DTZ005 The use of `datetime.datetime.now()` without `tz` argumen
 17 | # no args unqualified
 18 | datetime.now()
    | ^^^^^^^^^^^^^^ DTZ005
+19 | 
+20 | # uses `astimezone` method
    |
 
 

crates/ruff/src/rules/flake8_datetimez/snapshots/ruff__rules__flake8_datetimez__tests__DTZ006_DTZ006.py.snap

@@ -42,6 +42,8 @@ DTZ006.py:18:1: DTZ006 The use of `datetime.datetime.fromtimestamp()` without `t
 17 | # no args unqualified
 18 | datetime.fromtimestamp(1234)
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DTZ006
+19 | 
+20 | # uses `astimezone` method
    |