Skip to content

Commit

Permalink
Allow flake8-type-checking rules to automatically quote runtime-eva…
Browse files Browse the repository at this point in the history
…luated references (#6001)

## Summary

This allows us to fix usages like:

```python
from pandas import DataFrame

def baz() -> DataFrame:
    ...
```

By quoting the `DataFrame` in `-> DataFrame`. Without quotes, moving
`from pandas import DataFrame` into an `if TYPE_CHECKING:` block will
fail at runtime, since Python tries to evaluate the annotation to add it
to the function's `__annotations__`.

Unfortunately, this does require us to split our "annotation kind" flags
into three categories, rather than two:

- `typing-only`: The annotation is only evaluated at type-checking-time.
- `runtime-evaluated`: Python will evaluate the annotation at runtime
(like above) -- but we're willing to quote it.
- `runtime-required`: Python will evaluate the annotation at runtime
(like above), and some library (like Pydantic) needs it to be available
at runtime, so we _can't_ quote it.

This functionality is gated behind a setting
(`flake8-type-checking.quote-annotations`).

Closes #5559.
  • Loading branch information
charliermarsh committed Dec 13, 2023
1 parent 4d2ee5b commit 1a65e54
Show file tree
Hide file tree
Showing 18 changed files with 1,033 additions and 207 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
def f():
from pandas import DataFrame

def baz() -> DataFrame:
...


def f():
from pandas import DataFrame

def baz() -> DataFrame[int]:
...


def f():
from pandas import DataFrame

def baz() -> DataFrame["int"]:
...


def f():
import pandas as pd

def baz() -> pd.DataFrame:
...


def f():
import pandas as pd

def baz() -> pd.DataFrame.Extra:
...


def f():
import pandas as pd

def baz() -> pd.DataFrame | int:
...



def f():
from pandas import DataFrame

def baz() -> DataFrame():
...


def f():
from typing import Literal

from pandas import DataFrame

def baz() -> DataFrame[Literal["int"]]:
...


def f():
from typing import TYPE_CHECKING

if TYPE_CHECKING:
from pandas import DataFrame

def func(value: DataFrame):
...
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) {
flake8_type_checking::helpers::is_valid_runtime_import(
binding,
&checker.semantic,
&checker.settings.flake8_type_checking,
)
})
.collect()
Expand Down
66 changes: 66 additions & 0 deletions crates/ruff_linter/src/checkers/ast/annotation.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
use ruff_python_semantic::{ScopeKind, SemanticModel};

use crate::rules::flake8_type_checking;
use crate::settings::LinterSettings;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum AnnotationContext {
/// Python will evaluate the annotation at runtime, but it's not _required_ and, as such, could
/// be quoted to convert it into a typing-only annotation.
///
/// For example:
/// ```python
/// from pandas import DataFrame
///
/// def foo() -> DataFrame:
/// ...
/// ```
///
/// Above, Python will evaluate `DataFrame` at runtime in order to add it to `__annotations__`.
RuntimeEvaluated,
/// Python will evaluate the annotation at runtime, and it's required to be available at
/// runtime, as a library (like Pydantic) needs access to it.
RuntimeRequired,
/// The annotation is only evaluated at type-checking time.
TypingOnly,
}

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_def,
&settings.flake8_type_checking.runtime_required_base_classes,
&settings.flake8_type_checking.runtime_required_decorators,
semantic,
)
})
{
return Self::RuntimeRequired;
}

// If `__future__` annotations are enabled, then annotations are never evaluated
// at runtime, so we can treat them as typing-only.
if semantic.future_annotations() {
return Self::TypingOnly;
}

// Otherwise, if we're in a class or module scope, then the annotation needs to
// be available at runtime.
// See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements
if matches!(
semantic.current_scope().kind,
ScopeKind::Class(_) | ScopeKind::Module
) {
return Self::RuntimeEvaluated;
}

Self::TypingOnly
}
}
68 changes: 28 additions & 40 deletions crates/ruff_linter/src/checkers/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ use ruff_python_semantic::{
use ruff_python_stdlib::builtins::{IPYTHON_BUILTINS, MAGIC_GLOBALS, PYTHON_BUILTINS};
use ruff_source_file::Locator;

use crate::checkers::ast::annotation::AnnotationContext;
use crate::checkers::ast::deferred::Deferred;
use crate::docstrings::extraction::ExtractionTarget;
use crate::importer::Importer;
Expand All @@ -68,6 +69,7 @@ use crate::settings::{flags, LinterSettings};
use crate::{docstrings, noqa};

mod analyze;
mod annotation;
mod deferred;

pub(crate) struct Checker<'a> {
Expand Down Expand Up @@ -515,8 +517,10 @@ where
.chain(&parameters.kwonlyargs)
{
if let Some(expr) = &parameter_with_default.parameter.annotation {
if runtime_annotation || singledispatch {
self.visit_runtime_annotation(expr);
if singledispatch {
self.visit_runtime_required_annotation(expr);
} else if runtime_annotation {
self.visit_runtime_evaluated_annotation(expr);
} else {
self.visit_annotation(expr);
};
Expand All @@ -529,7 +533,7 @@ where
if let Some(arg) = &parameters.vararg {
if let Some(expr) = &arg.annotation {
if runtime_annotation {
self.visit_runtime_annotation(expr);
self.visit_runtime_evaluated_annotation(expr);
} else {
self.visit_annotation(expr);
};
Expand All @@ -538,15 +542,15 @@ where
if let Some(arg) = &parameters.kwarg {
if let Some(expr) = &arg.annotation {
if runtime_annotation {
self.visit_runtime_annotation(expr);
self.visit_runtime_evaluated_annotation(expr);
} else {
self.visit_annotation(expr);
};
}
}
for expr in returns {
if runtime_annotation {
self.visit_runtime_annotation(expr);
self.visit_runtime_evaluated_annotation(expr);
} else {
self.visit_annotation(expr);
};
Expand Down Expand Up @@ -677,40 +681,16 @@ where
value,
..
}) => {
// If we're in a class or module scope, then the annotation needs to be
// available at runtime.
// See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements
let runtime_annotation = if self.semantic.future_annotations() {
self.semantic
.current_scope()
.kind
.as_class()
.is_some_and(|class_def| {
flake8_type_checking::helpers::runtime_evaluated_class(
class_def,
&self
.settings
.flake8_type_checking
.runtime_evaluated_base_classes,
&self
.settings
.flake8_type_checking
.runtime_evaluated_decorators,
&self.semantic,
)
})
} else {
matches!(
self.semantic.current_scope().kind,
ScopeKind::Class(_) | ScopeKind::Module
)
};

if runtime_annotation {
self.visit_runtime_annotation(annotation);
} else {
self.visit_annotation(annotation);
match AnnotationContext::from_model(&self.semantic, self.settings) {
AnnotationContext::RuntimeRequired => {
self.visit_runtime_required_annotation(annotation);
}
AnnotationContext::RuntimeEvaluated => {
self.visit_runtime_evaluated_annotation(annotation);
}
AnnotationContext::TypingOnly => self.visit_annotation(annotation),
}

if let Some(expr) = value {
if self.semantic.match_typing_expr(annotation, "TypeAlias") {
self.visit_type_definition(expr);
Expand Down Expand Up @@ -1527,10 +1507,18 @@ impl<'a> Checker<'a> {
self.semantic.flags = snapshot;
}

/// Visit an [`Expr`], and treat it as a runtime-evaluated type annotation.
fn visit_runtime_evaluated_annotation(&mut self, expr: &'a Expr) {
let snapshot = self.semantic.flags;
self.semantic.flags |= SemanticModelFlags::RUNTIME_EVALUATED_ANNOTATION;
self.visit_type_definition(expr);
self.semantic.flags = snapshot;
}

/// Visit an [`Expr`], and treat it as a runtime-required type annotation.
fn visit_runtime_annotation(&mut self, expr: &'a Expr) {
fn visit_runtime_required_annotation(&mut self, expr: &'a Expr) {
let snapshot = self.semantic.flags;
self.semantic.flags |= SemanticModelFlags::RUNTIME_ANNOTATION;
self.semantic.flags |= SemanticModelFlags::RUNTIME_REQUIRED_ANNOTATION;
self.visit_type_definition(expr);
self.semantic.flags = snapshot;
}
Expand Down
Loading

0 comments on commit 1a65e54

Please sign in to comment.