PYI011: allow math constants in defaults

astral-sh · Mar 13, 2023 · e32835b · e32835b
1 parent 3a5fbd6
commit e32835b
Show file tree

Hide file tree

Showing 3 changed files with 170 additions and 2 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi
@@ -61,3 +61,43 @@ def f22(
     x: complex = -42.5j  # Error PYI011 Only simple default values allowed for typed arguments
     + 4.3j,
 ) -> None: ...
+def f23(
+    x: bool = True,  # OK
+) -> None: ...
+def f24(
+    x: float = 3.14,  # OK
+) -> None: ...
+def f25(
+    x: float = -3.14,  # OK
+) -> None: ...
+def f26(
+    x: complex = -3.14j,  # OK
+) -> None: ...
+def f27(
+    x: complex = -3 - 3.14j,  # OK
+) -> None: ...
+def f28(
+    x: float = math.tau,  # OK
+) -> None: ...
+def f29(
+    x: float = math.inf,  # OK
+) -> None: ...
+def f30(
+    x: float = -math.inf,  # OK
+) -> None: ...
+def f31(
+    x: float = inf,  # Error PYI011 Only simple default values allowed for typed arguments
+) -> None: ...
+def f32(
+    x: float = np.inf,  # Error PYI011 Only simple default values allowed for typed arguments
+) -> None: ...
+def f33(
+    x: float = math.nan,  # OK
+) -> None: ...
+def f34(
+    x: float = -math.nan,  # Error PYI011 Only simple default values allowed for typed arguments
+) -> None: ...
+def f35(
+    x: complex = math.inf  # Error PYI011 Only simple default values allowed for typed arguments
+    + 1j,
+) -> None: ...
diff --git a/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs b/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs
@@ -28,6 +28,14 @@ impl Violation for ArgumentSimpleDefaults {
     }
 }
 
+const ALLOWED_MATH_ATTRIBUTES_IN_DEFAULTS: &[&[&str]] = &[
+    &["math", "inf"],
+    &["math", "nan"],
+    &["math", "e"],
+    &["math", "pi"],
+    &["math", "tau"],
+];
+
 const ALLOWED_ATTRIBUTES_IN_DEFAULTS: &[&[&str]] = &[
     &["sys", "stdin"],
     &["sys", "stdout"],
@@ -99,6 +107,21 @@ fn is_valid_default_value_with_annotation(default: &Expr, checker: &Checker) ->
                     return checker.locator.slice(operand).len() <= 10;
                 }
             }
+            // Ex) `-math.inf`, `-math.pi`, etc.
+            if let ExprKind::Attribute { .. } = &operand.node {
+                if checker
+                    .ctx
+                    .resolve_call_path(default)
+                    .map_or(false, |call_path| {
+                        ALLOWED_MATH_ATTRIBUTES_IN_DEFAULTS.iter().any(|target| {
+                            // reject `-math.nan`
+                            call_path.as_slice() == *target && *target != ["math", "nan"]
+                        })
+                    })
+                {
+                    return true;
+                }
+            }
         }
         ExprKind::BinOp {
             left,
@@ -134,14 +157,15 @@ fn is_valid_default_value_with_annotation(default: &Expr, checker: &Checker) ->
                 }
             }
         }
-        // Ex) `sys.stdin`, etc.
+        // Ex) `math.inf`, `sys.stdin`, etc.
         ExprKind::Attribute { .. } => {
             if checker
                 .ctx
                 .resolve_call_path(default)
                 .map_or(false, |call_path| {
-                    ALLOWED_ATTRIBUTES_IN_DEFAULTS
+                    ALLOWED_MATH_ATTRIBUTES_IN_DEFAULTS
                         .iter()
+                        .chain(ALLOWED_ATTRIBUTES_IN_DEFAULTS.iter())
                         .any(|target| call_path.as_slice() == *target)
                 })
             {

diff --git a/...uff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI011_PYI011.pyi.snap b/...uff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI011_PYI011.pyi.snap
@@ -145,4 +145,108 @@ expression: diagnostics
     column: 10
   fix: ~
   parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 80
+    column: 15
+  end_location:
+    row: 80
+    column: 23
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 83
+    column: 15
+  end_location:
+    row: 83
+    column: 23
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 86
+    column: 15
+  end_location:
+    row: 86
+    column: 24
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 89
+    column: 15
+  end_location:
+    row: 89
+    column: 18
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 92
+    column: 15
+  end_location:
+    row: 92
+    column: 21
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 95
+    column: 15
+  end_location:
+    row: 95
+    column: 23
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 98
+    column: 15
+  end_location:
+    row: 98
+    column: 24
+  fix: ~
+  parent: ~
+- kind:
+    name: TypedArgumentSimpleDefaults
+    body: Only simple default values allowed for typed arguments
+    suggestion: ~
+    fixable: false
+  location:
+    row: 101
+    column: 17
+  end_location:
+    row: 102
+    column: 8
+  fix: ~
+  parent: ~