[flake8-datetime] Ignore .replace() calls while looking for .astimezone (#16050)

Co-authored-by: Alex Waygood <alex.waygood@gmail.com>

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

@@ -19,3 +19,21 @@ datetime.now()
 
 # uses `astimezone` method
 datetime.now().astimezone()
+datetime.now().astimezone
+
+
+# https://github.com/astral-sh/ruff/issues/15998
+
+## Errors
+datetime.now().replace.astimezone()
+datetime.now().replace[0].astimezone()
+datetime.now()().astimezone()
+datetime.now().replace(datetime.now()).astimezone()
+
+foo.replace(datetime.now().replace).astimezone()
+
+## No errors
+datetime.now().replace(microsecond=0).astimezone()
+datetime.now().replace(0).astimezone()
+datetime.now().replace(0).astimezone
+datetime.now().replace(0).replace(1).astimezone

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

@@ -87,7 +87,7 @@ pub(crate) fn call_datetime_fromtimestamp(checker: &Checker, call: &ast::ExprCal
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -82,7 +82,7 @@ pub(crate) fn call_datetime_now_without_tzinfo(checker: &Checker, call: &ast::Ex
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -71,7 +71,7 @@ pub(crate) fn call_datetime_today(checker: &Checker, func: &Expr, location: Text
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -78,7 +78,7 @@ pub(crate) fn call_datetime_utcfromtimestamp(checker: &Checker, func: &Expr, loc
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -78,7 +78,7 @@ pub(crate) fn call_datetime_utcnow(checker: &Checker, func: &Expr, location: Tex
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -76,7 +76,7 @@ pub(crate) fn call_datetime_without_tzinfo(checker: &Checker, call: &ast::ExprCa
         return;
     }
 
-    if helpers::parent_expr_is_astimezone(checker) {
+    if helpers::followed_by_astimezone(checker) {
         return;
     }
 

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

@@ -1,4 +1,4 @@
-use ruff_python_ast::{Expr, ExprAttribute};
+use ruff_python_ast::{AnyNodeRef, Expr, ExprAttribute, ExprCall};
 
 use crate::checkers::ast::Checker;
 
@@ -8,10 +8,101 @@ pub(super) enum DatetimeModuleAntipattern {
     NonePassedToTzArgument,
 }
 
-/// Check if the parent expression is a call to `astimezone`.
-/// This assumes that the current expression is a `datetime.datetime` object.
-pub(super) fn parent_expr_is_astimezone(checker: &Checker) -> bool {
-    checker.semantic().current_expression_parent().is_some_and(|parent| {
-        matches!(parent, Expr::Attribute(ExprAttribute { attr, .. }) if attr.as_str() == "astimezone")
-    })
+/// Check if the "current expression" being visited is followed
+/// in the source code by a chain of `.replace()` calls followed by `.astimezone`.
+/// The function operates on the assumption that the current expression
+/// is a [`datetime.datetime`][datetime] object.
+///
+/// For example, given the following Python source code:
+///
+/// ```py
+/// import datetime
+///
+/// datetime.now().replace(hours=4).replace(minutes=46).astimezone()
+/// ```
+///
+/// The last line will produce an AST looking something like this
+/// (this is pseudocode approximating our AST):
+///
+/// ```rs
+/// Call {
+///     func: Attribute {
+///         value: Call {
+///             func: Attribute {
+///                 value: Call {
+///                    func: Attribute {
+///                        value: Call {                    // We are visiting this
+///                            func: Attribute {            // expression node here
+///                                value: Call {            //
+///                                    func: Name {         //
+///                                        id: "datetime",  //
+///                                    },                   //
+///                                },                       //
+///                                attr: "now"              //
+///                            },                           //
+///                        },                               //
+///                        attr: "replace"
+///                    },
+///                 },
+///                 attr: "replace"
+///             },
+///         },
+///         attr: "astimezone"
+///     },
+/// }
+/// ```
+///
+/// The node we are visiting as the "current expression" is deeply
+/// nested inside many other expressions. As such, in order to check
+/// whether the `datetime.now()` call is followed by 0-or-more `.replace()`
+/// calls and then an `.astimezone()` call, we must iterate up through the
+/// "parent expressions" in the semantic model, checking if they match this
+/// AST pattern.
+///
+/// [datetime]: https://docs.python.org/3/library/datetime.html#datetime-objects
+pub(super) fn followed_by_astimezone(checker: &Checker) -> bool {
+    let semantic = checker.semantic();
+    let mut last = None;
+
+    for (index, expr) in semantic.current_expressions().enumerate() {
+        if index == 0 {
+            // datetime.now(...).replace(...).astimezone
+            // ^^^^^^^^^^^^^^^^^
+            continue;
+        }
+
+        if index % 2 == 1 {
+            // datetime.now(...).replace(...).astimezone
+            //                   ^^^^^^^      ^^^^^^^^^^
+            let Expr::Attribute(ExprAttribute { attr, .. }) = expr else {
+                return false;
+            };
+
+            match attr.as_str() {
+                "replace" => last = Some(AnyNodeRef::from(expr)),
+                "astimezone" => return true,
+                _ => return false,
+            }
+        } else {
+            // datetime.now(...).replace(...).astimezone
+            //                          ^^^^^
+            let Expr::Call(ExprCall { func, .. }) = expr else {
+                return false;
+            };
+
+            // Without this branch, we would fail to emit a diagnostic on code like this:
+            //
+            // ```py
+            // foo.replace(datetime.now().replace).astimezone()
+            //           # ^^^^^^^^^^^^^^  Diagnostic should be emitted here
+            //           #                 since the `datetime.now()` call is not followed
+            //           #                 by `.astimezone()`
+            // ```
+            if !last.is_some_and(|it| it.ptr_eq(AnyNodeRef::from(&**func))) {
+                return false;
+            }
+        }
+    }
+
+    false
 }

crates/ruff_linter/src/rules/flake8_datetimez/snapshots/ruff_linter__rules__flake8_datetimez__tests__DTZ005_DTZ005.py.snap

@@ -50,3 +50,56 @@ DTZ005.py:18:1: DTZ005 `datetime.datetime.now()` called without a `tz` argument
 20 | # uses `astimezone` method
    |
    = help: Pass a `datetime.timezone` object to the `tz` parameter
+
+DTZ005.py:28:1: DTZ005 `datetime.datetime.now()` called without a `tz` argument
+   |
+27 | ## Errors
+28 | datetime.now().replace.astimezone()
+   | ^^^^^^^^^^^^^^ DTZ005
+29 | datetime.now().replace[0].astimezone()
+30 | datetime.now()().astimezone()
+   |
+   = help: Pass a `datetime.timezone` object to the `tz` parameter
+
+DTZ005.py:29:1: DTZ005 `datetime.datetime.now()` called without a `tz` argument
+   |
+27 | ## Errors
+28 | datetime.now().replace.astimezone()
+29 | datetime.now().replace[0].astimezone()
+   | ^^^^^^^^^^^^^^ DTZ005
+30 | datetime.now()().astimezone()
+31 | datetime.now().replace(datetime.now()).astimezone()
+   |
+   = help: Pass a `datetime.timezone` object to the `tz` parameter
+
+DTZ005.py:30:1: DTZ005 `datetime.datetime.now()` called without a `tz` argument
+   |
+28 | datetime.now().replace.astimezone()
+29 | datetime.now().replace[0].astimezone()
+30 | datetime.now()().astimezone()
+   | ^^^^^^^^^^^^^^ DTZ005
+31 | datetime.now().replace(datetime.now()).astimezone()
+   |
+   = help: Pass a `datetime.timezone` object to the `tz` parameter
+
+DTZ005.py:31:24: DTZ005 `datetime.datetime.now()` called without a `tz` argument
+   |
+29 | datetime.now().replace[0].astimezone()
+30 | datetime.now()().astimezone()
+31 | datetime.now().replace(datetime.now()).astimezone()
+   |                        ^^^^^^^^^^^^^^ DTZ005
+32 |
+33 | foo.replace(datetime.now().replace).astimezone()
+   |
+   = help: Pass a `datetime.timezone` object to the `tz` parameter
+
+DTZ005.py:33:13: DTZ005 `datetime.datetime.now()` called without a `tz` argument
+   |
+31 | datetime.now().replace(datetime.now()).astimezone()
+32 |
+33 | foo.replace(datetime.now().replace).astimezone()
+   |             ^^^^^^^^^^^^^^ DTZ005
+34 |
+35 | ## No errors
+   |
+   = help: Pass a `datetime.timezone` object to the `tz` parameter