Skip to content

Commit

Permalink
Respect runtime-required decorators on functions
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Dec 29, 2023
1 parent 2895e7d commit 2ddfa62
Show file tree
Hide file tree
Showing 6 changed files with 91 additions and 30 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from array import array
from dataclasses import dataclass
from uuid import UUID # TCH003
from collections.abc import Sequence
from pydantic import validate_call

import attrs
from attrs import frozen
Expand All @@ -22,3 +24,8 @@ class B:
@dataclass
class C:
x: UUID


@validate_call(config={'arbitrary_types_allowed': True})
def test(user: Sequence):
...
30 changes: 19 additions & 11 deletions crates/ruff_linter/src/checkers/ast/annotation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,21 +28,29 @@ pub(super) enum AnnotationContext {
impl AnnotationContext {
pub(super) fn from_model(semantic: &SemanticModel, settings: &LinterSettings) -> Self {
// If the annotation is in a class scope (e.g., an annotated assignment for a
// class field), and that class is marked as annotation as runtime-required.
if semantic
.current_scope()
.kind
.as_class()
.is_some_and(|class_def| {
flake8_type_checking::helpers::runtime_required_class(
// class field) or a function scope, and that class or function is marked as
// runtime-required, treat the annotation as runtime-required.
match semantic.current_scope().kind {
ScopeKind::Class(class_def)
if flake8_type_checking::helpers::runtime_required_class(
class_def,
&settings.flake8_type_checking.runtime_required_base_classes,
&settings.flake8_type_checking.runtime_required_decorators,
semantic,
)
})
{
return Self::RuntimeRequired;
) =>
{
return Self::RuntimeRequired
}
ScopeKind::Function(function_def)
if flake8_type_checking::helpers::runtime_required_function(
function_def,
&settings.flake8_type_checking.runtime_required_decorators,
semantic,
) =>
{
return Self::RuntimeRequired
}
_ => {}
}

// If `__future__` annotations are enabled, then annotations are never evaluated
Expand Down
21 changes: 17 additions & 4 deletions crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use anyhow::Result;
use ruff_diagnostics::Edit;
use ruff_python_ast::call_path::from_qualified_name;
use ruff_python_ast::helpers::map_callable;
use ruff_python_ast::{self as ast, Expr};
use ruff_python_ast::{self as ast, Decorator, Expr};
use ruff_python_codegen::{Generator, Stylist};
use ruff_python_semantic::{
analyze, Binding, BindingKind, NodeId, ResolvedReference, SemanticModel,
Expand Down Expand Up @@ -43,6 +43,19 @@ pub(crate) fn is_valid_runtime_import(
}
}

/// Returns `true` if a function's parameters should be treated as runtime-required.
pub(crate) fn runtime_required_function(
function_def: &ast::StmtFunctionDef,
decorators: &[String],
semantic: &SemanticModel,
) -> bool {
if runtime_required_decorators(&function_def.decorator_list, decorators, semantic) {
return true;
}
false
}

/// Returns `true` if a class's assignments should be treated as runtime-required.
pub(crate) fn runtime_required_class(
class_def: &ast::StmtClassDef,
base_classes: &[String],
Expand All @@ -52,7 +65,7 @@ pub(crate) fn runtime_required_class(
if runtime_required_base_class(class_def, base_classes, semantic) {
return true;
}
if runtime_required_decorators(class_def, decorators, semantic) {
if runtime_required_decorators(&class_def.decorator_list, decorators, semantic) {
return true;
}
false
Expand All @@ -72,15 +85,15 @@ fn runtime_required_base_class(
}

fn runtime_required_decorators(
class_def: &ast::StmtClassDef,
decorator_list: &[Decorator],
decorators: &[String],
semantic: &SemanticModel,
) -> bool {
if decorators.is_empty() {
return false;
}

class_def.decorator_list.iter().any(|decorator| {
decorator_list.iter().any(|decorator| {
semantic
.resolve_call_path(map_callable(&decorator.expression))
.is_some_and(|call_path| {
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,7 @@ mod tests {
runtime_required_decorators: vec![
"attrs.define".to_string(),
"attrs.frozen".to_string(),
"pydantic.validate_call".to_string(),
],
..Default::default()
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ runtime_evaluated_decorators_3.py:6:18: TCH003 [*] Move standard library import
5 | from dataclasses import dataclass
6 | from uuid import UUID # TCH003
| ^^^^ TCH003
7 |
8 | import attrs
7 | from collections.abc import Sequence
8 | from pydantic import validate_call
|
= help: Move into type-checking block

Expand All @@ -17,15 +17,44 @@ runtime_evaluated_decorators_3.py:6:18: TCH003 [*] Move standard library import
4 4 | from array import array
5 5 | from dataclasses import dataclass
6 |-from uuid import UUID # TCH003
7 6 |
8 7 | import attrs
9 8 | from attrs import frozen
9 |+from typing import TYPE_CHECKING
10 |+
11 |+if TYPE_CHECKING:
12 |+ from uuid import UUID
10 13 |
11 14 |
12 15 | @attrs.define(auto_attribs=True)
7 6 | from collections.abc import Sequence
8 7 | from pydantic import validate_call
9 8 |
10 9 | import attrs
11 10 | from attrs import frozen
11 |+from typing import TYPE_CHECKING
12 |+
13 |+if TYPE_CHECKING:
14 |+ from uuid import UUID
12 15 |
13 16 |
14 17 | @attrs.define(auto_attribs=True)

runtime_evaluated_decorators_3.py:7:29: TCH003 [*] Move standard library import `collections.abc.Sequence` into a type-checking block
|
5 | from dataclasses import dataclass
6 | from uuid import UUID # TCH003
7 | from collections.abc import Sequence
| ^^^^^^^^ TCH003
8 | from pydantic import validate_call
|
= help: Move into type-checking block

Unsafe fix
4 4 | from array import array
5 5 | from dataclasses import dataclass
6 6 | from uuid import UUID # TCH003
7 |-from collections.abc import Sequence
8 7 | from pydantic import validate_call
9 8 |
10 9 | import attrs
11 10 | from attrs import frozen
11 |+from typing import TYPE_CHECKING
12 |+
13 |+if TYPE_CHECKING:
14 |+ from collections.abc import Sequence
12 15 |
13 16 |
14 17 | @attrs.define(auto_attribs=True)


9 changes: 6 additions & 3 deletions crates/ruff_workspace/src/options.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1632,13 +1632,16 @@ pub struct Flake8TypeCheckingOptions {
)]
pub runtime_evaluated_base_classes: Option<Vec<String>>,

/// Exempt classes decorated with any of the enumerated decorators from
/// needing to be moved into type-checking blocks.
/// Exempt classes and functions decorated with any of the enumerated
/// decorators from being moved into type-checking blocks.
///
/// Common examples include Pydantic's `@pydantic.validate_call` decorator
/// (for functions) and attrs' `@attrs.define` decorator (for classes).
#[option(
default = "[]",
value_type = "list[str]",
example = r#"
runtime-evaluated-decorators = ["attrs.define", "attrs.frozen"]
runtime-evaluated-decorators = ["pydantic.validate_call", "attrs.define"]
"#
)]
pub runtime_evaluated_decorators: Option<Vec<String>>,
Expand Down

0 comments on commit 2ddfa62

Please sign in to comment.