Allow f-strings with %z for DTZ007

astral-sh · Mar 29, 2024 · d637aa5 · d637aa5
1 parent a0263ab
commit d637aa5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_datetimez/DTZ007.py b/crates/ruff_linter/resources/test/fixtures/flake8_datetimez/DTZ007.py
@@ -33,3 +33,7 @@
 
 # no replace orastimezone unqualified
 datetime.strptime("something", "something")
+
+# F-strings
+datetime.strptime("something", f"%Y-%m-%dT%H:%M:%S{('.%f' if millis else '')}%z")
+datetime.strptime("something", f"%Y-%m-%d %H:%M:%S%z")
diff --git a/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_strptime_without_zone.rs b/crates/ruff_linter/src/rules/flake8_datetimez/rules/call_datetime_strptime_without_zone.rs
@@ -102,11 +102,23 @@ pub(crate) fn call_datetime_strptime_without_zone(checker: &mut Checker, call: &
     }
 
     // Does the `strptime` call contain a format string with a timezone specifier?
-    if let Some(Expr::StringLiteral(ast::ExprStringLiteral { value: format, .. })) =
-        call.arguments.args.get(1).as_ref()
-    {
-        if format.to_str().contains("%z") {
-            return;
+    if let Some(expr) = call.arguments.args.get(1) {
+        match expr {
+            Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
+                if value.to_str().contains("%z") {
+                    return;
+                }
+            }
+            Expr::FString(ast::ExprFString { value, .. }) => {
+                // TODO(dhruvmanila): This doesn't consider f-strings that are implicitly concatenated
+                // to strings. For example, `f"%Y-%m-%dT%H:%M:%S{('.%f' if millis else '')}" "%z"`
+                for f_string in value.f_strings() {
+                    if f_string.literals().any(|literal| literal.contains("%z")) {
+                        return;
+                    }
+                }
+            }
+            _ => {}
         }
     };
 

diff --git a/...e8_datetimez/snapshots/ruff_linter__rules__flake8_datetimez__tests__DTZ007_DTZ007.py.snap b/...e8_datetimez/snapshots/ruff_linter__rules__flake8_datetimez__tests__DTZ007_DTZ007.py.snap
@@ -46,5 +46,7 @@ DTZ007.py:35:1: DTZ007 Naive datetime constructed using `datetime.datetime.strpt
 34 | # no replace orastimezone unqualified
 35 | datetime.strptime("something", "something")
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DTZ007
+36 | 
+37 | # F-strings
    |
    = help: Call `.replace(tzinfo=<timezone>)` or `.astimezone()` to convert to an aware datetime