Add a Snapshot abstraction for deferring and restoring visitor context (#4353)

This commit is contained in:
Charlie Marsh 2023-05-10 12:50:47 -04:00 committed by GitHub
parent fd34797d0f
commit ea3d3a655d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 81 additions and 84 deletions

View file

@ -2,26 +2,20 @@ use ruff_text_size::TextRange;
use rustpython_parser::ast::Expr; use rustpython_parser::ast::Expr;
use ruff_python_semantic::analyze::visibility::{Visibility, VisibleScope}; use ruff_python_semantic::analyze::visibility::{Visibility, VisibleScope};
use ruff_python_semantic::node::NodeId; use ruff_python_semantic::context::Snapshot;
use ruff_python_semantic::scope::ScopeId;
use crate::checkers::ast::AnnotationContext;
use crate::docstrings::definition::Definition; use crate::docstrings::definition::Definition;
/// A snapshot of the current scope and statement, which will be restored when visiting any
/// deferred definitions.
type Context<'a> = (ScopeId, Option<NodeId>);
/// A collection of AST nodes that are deferred for later analysis. /// A collection of AST nodes that are deferred for later analysis.
/// Used to, e.g., store functions, whose bodies shouldn't be analyzed until all /// Used to, e.g., store functions, whose bodies shouldn't be analyzed until all
/// module-level definitions have been analyzed. /// module-level definitions have been analyzed.
#[derive(Default)] #[derive(Default)]
pub struct Deferred<'a> { pub struct Deferred<'a> {
pub definitions: Vec<(Definition<'a>, Visibility, Context<'a>)>, pub definitions: Vec<(Definition<'a>, Visibility, Snapshot)>,
pub string_type_definitions: Vec<(TextRange, &'a str, AnnotationContext, Context<'a>)>, pub string_type_definitions: Vec<(TextRange, &'a str, Snapshot)>,
pub type_definitions: Vec<(&'a Expr, AnnotationContext, Context<'a>)>, pub type_definitions: Vec<(&'a Expr, Snapshot)>,
pub functions: Vec<(Context<'a>, VisibleScope)>, pub functions: Vec<(Snapshot, VisibleScope)>,
pub lambdas: Vec<(&'a Expr, Context<'a>)>, pub lambdas: Vec<(&'a Expr, Snapshot)>,
pub for_loops: Vec<Context<'a>>, pub for_loops: Vec<Snapshot>,
pub assignments: Vec<Context<'a>>, pub assignments: Vec<Snapshot>,
} }

View file

@ -55,8 +55,6 @@ use crate::{autofix, docstrings, noqa, warn_user};
mod deferred; mod deferred;
type AnnotationContext = (bool, bool);
pub struct Checker<'a> { pub struct Checker<'a> {
// Settings, static metadata, etc. // Settings, static metadata, etc.
path: &'a Path, path: &'a Path,
@ -1706,9 +1704,7 @@ where
.. ..
} => { } => {
if self.settings.rules.enabled(Rule::UnusedLoopControlVariable) { if self.settings.rules.enabled(Rule::UnusedLoopControlVariable) {
self.deferred self.deferred.for_loops.push(self.ctx.snapshot());
.for_loops
.push((self.ctx.scope_id, self.ctx.stmt_id));
} }
if self if self
.settings .settings
@ -1994,11 +1990,9 @@ where
pyupgrade::rules::yield_in_for_loop(self, stmt); pyupgrade::rules::yield_in_for_loop(self, stmt);
} }
let scope = transition_scope(self.ctx.visible_scope, stmt, Documentable::Function); let scope = transition_scope(self.ctx.visible_scope, stmt, Documentable::Function);
self.deferred.definitions.push(( self.deferred
definition, .definitions
scope.visibility, .push((definition, scope.visibility, self.ctx.snapshot()));
(self.ctx.scope_id, self.ctx.stmt_id),
));
self.ctx.visible_scope = scope; self.ctx.visible_scope = scope;
// If any global bindings don't already exist in the global scope, add it. // If any global bindings don't already exist in the global scope, add it.
@ -2033,10 +2027,9 @@ where
globals, globals,
})); }));
self.deferred.functions.push(( self.deferred
(self.ctx.scope_id, self.ctx.stmt_id), .functions
self.ctx.visible_scope, .push((self.ctx.snapshot(), self.ctx.visible_scope));
));
} }
StmtKind::ClassDef { StmtKind::ClassDef {
body, body,
@ -2056,11 +2049,9 @@ where
Documentable::Class, Documentable::Class,
); );
let scope = transition_scope(self.ctx.visible_scope, stmt, Documentable::Class); let scope = transition_scope(self.ctx.visible_scope, stmt, Documentable::Class);
self.deferred.definitions.push(( self.deferred
definition, .definitions
scope.visibility, .push((definition, scope.visibility, self.ctx.snapshot()));
(self.ctx.scope_id, self.ctx.stmt_id),
));
self.ctx.visible_scope = scope; self.ctx.visible_scope = scope;
// If any global bindings don't already exist in the global scope, add it. // If any global bindings don't already exist in the global scope, add it.
@ -2273,15 +2264,12 @@ where
self.deferred.string_type_definitions.push(( self.deferred.string_type_definitions.push((
expr.range(), expr.range(),
value, value,
(self.ctx.in_annotation, self.ctx.in_type_checking_block), self.ctx.snapshot(),
(self.ctx.scope_id, self.ctx.stmt_id),
)); ));
} else { } else {
self.deferred.type_definitions.push(( self.deferred
expr, .type_definitions
(self.ctx.in_annotation, self.ctx.in_type_checking_block), .push((expr, self.ctx.snapshot()));
(self.ctx.scope_id, self.ctx.stmt_id),
));
} }
return; return;
} }
@ -3526,8 +3514,7 @@ where
self.deferred.string_type_definitions.push(( self.deferred.string_type_definitions.push((
expr.range(), expr.range(),
value, value,
(self.ctx.in_annotation, self.ctx.in_type_checking_block), self.ctx.snapshot(),
(self.ctx.scope_id, self.ctx.stmt_id),
)); ));
} }
if self if self
@ -3648,9 +3635,7 @@ where
// Recurse. // Recurse.
match &expr.node { match &expr.node {
ExprKind::Lambda { .. } => { ExprKind::Lambda { .. } => {
self.deferred self.deferred.lambdas.push((expr, self.ctx.snapshot()));
.lambdas
.push((expr, (self.ctx.scope_id, self.ctx.stmt_id)));
} }
ExprKind::IfExp { test, body, orelse } => { ExprKind::IfExp { test, body, orelse } => {
visit_boolean_test!(self, test); visit_boolean_test!(self, test);
@ -4793,7 +4778,7 @@ impl<'a> Checker<'a> {
docstring, docstring,
}, },
self.ctx.visible_scope.visibility, self.ctx.visible_scope.visibility,
(self.ctx.scope_id, self.ctx.stmt_id), self.ctx.snapshot(),
)); ));
docstring.is_some() docstring.is_some()
} }
@ -4801,13 +4786,9 @@ impl<'a> Checker<'a> {
fn check_deferred_type_definitions(&mut self) { fn check_deferred_type_definitions(&mut self) {
while !self.deferred.type_definitions.is_empty() { while !self.deferred.type_definitions.is_empty() {
let type_definitions = std::mem::take(&mut self.deferred.type_definitions); let type_definitions = std::mem::take(&mut self.deferred.type_definitions);
for (expr, (in_annotation, in_type_checking_block), (scope_id, stmt_id)) in for (expr, snapshot) in type_definitions {
type_definitions self.ctx.restore(snapshot);
{
self.ctx.scope_id = scope_id;
self.ctx.stmt_id = stmt_id;
self.ctx.in_annotation = in_annotation;
self.ctx.in_type_checking_block = in_type_checking_block;
self.ctx.in_type_definition = true; self.ctx.in_type_definition = true;
self.ctx.in_deferred_type_definition = true; self.ctx.in_deferred_type_definition = true;
self.visit_expr(expr); self.visit_expr(expr);
@ -4820,11 +4801,15 @@ impl<'a> Checker<'a> {
fn check_deferred_string_type_definitions(&mut self, allocator: &'a typed_arena::Arena<Expr>) { fn check_deferred_string_type_definitions(&mut self, allocator: &'a typed_arena::Arena<Expr>) {
while !self.deferred.string_type_definitions.is_empty() { while !self.deferred.string_type_definitions.is_empty() {
let type_definitions = std::mem::take(&mut self.deferred.string_type_definitions); let type_definitions = std::mem::take(&mut self.deferred.string_type_definitions);
for (range, value, (in_annotation, in_type_checking_block), (scope_id, stmt_id)) in for (range, value, snapshot) in type_definitions {
type_definitions
{
if let Ok((expr, kind)) = parse_type_annotation(value, range, self.locator) { if let Ok((expr, kind)) = parse_type_annotation(value, range, self.locator) {
if in_annotation && self.ctx.annotations_future_enabled { let expr = allocator.alloc(expr);
self.ctx.restore(snapshot);
self.ctx.in_type_definition = true;
self.ctx.in_deferred_string_type_definition = Some(kind);
if self.ctx.in_annotation && self.ctx.annotations_future_enabled {
if self.settings.rules.enabled(Rule::QuotedAnnotation) { if self.settings.rules.enabled(Rule::QuotedAnnotation) {
pyupgrade::rules::quoted_annotation(self, value, range); pyupgrade::rules::quoted_annotation(self, value, range);
} }
@ -4834,16 +4819,8 @@ impl<'a> Checker<'a> {
flake8_pyi::rules::quoted_annotation_in_stub(self, value, range); flake8_pyi::rules::quoted_annotation_in_stub(self, value, range);
} }
} }
let expr = allocator.alloc(expr);
self.ctx.scope_id = scope_id;
self.ctx.stmt_id = stmt_id;
self.ctx.in_annotation = in_annotation;
self.ctx.in_type_checking_block = in_type_checking_block;
self.ctx.in_type_definition = true;
self.ctx.in_deferred_string_type_definition = Some(kind);
self.visit_expr(expr); self.visit_expr(expr);
self.ctx.in_deferred_string_type_definition = None; self.ctx.in_deferred_string_type_definition = None;
self.ctx.in_type_definition = false; self.ctx.in_type_definition = false;
} else { } else {
@ -4867,9 +4844,8 @@ impl<'a> Checker<'a> {
fn check_deferred_functions(&mut self) { fn check_deferred_functions(&mut self) {
while !self.deferred.functions.is_empty() { while !self.deferred.functions.is_empty() {
let deferred_functions = std::mem::take(&mut self.deferred.functions); let deferred_functions = std::mem::take(&mut self.deferred.functions);
for ((scope_id, stmt_id), visibility) in deferred_functions { for (snapshot, visibility) in deferred_functions {
self.ctx.scope_id = scope_id; self.ctx.restore(snapshot);
self.ctx.stmt_id = stmt_id;
self.ctx.visible_scope = visibility; self.ctx.visible_scope = visibility;
match &self.ctx.stmt().node { match &self.ctx.stmt().node {
@ -4883,7 +4859,7 @@ impl<'a> Checker<'a> {
} }
} }
self.deferred.assignments.push((scope_id, stmt_id)); self.deferred.assignments.push(snapshot);
} }
} }
} }
@ -4891,9 +4867,8 @@ impl<'a> Checker<'a> {
fn check_deferred_lambdas(&mut self) { fn check_deferred_lambdas(&mut self) {
while !self.deferred.lambdas.is_empty() { while !self.deferred.lambdas.is_empty() {
let lambdas = std::mem::take(&mut self.deferred.lambdas); let lambdas = std::mem::take(&mut self.deferred.lambdas);
for (expr, (scope_id, stmt_id)) in lambdas { for (expr, snapshot) in lambdas {
self.ctx.scope_id = scope_id; self.ctx.restore(snapshot);
self.ctx.stmt_id = stmt_id;
if let ExprKind::Lambda { args, body } = &expr.node { if let ExprKind::Lambda { args, body } = &expr.node {
self.visit_arguments(args); self.visit_arguments(args);
@ -4902,7 +4877,7 @@ impl<'a> Checker<'a> {
unreachable!("Expected ExprKind::Lambda"); unreachable!("Expected ExprKind::Lambda");
} }
self.deferred.assignments.push((scope_id, stmt_id)); self.deferred.assignments.push(snapshot);
} }
} }
} }
@ -4910,13 +4885,15 @@ impl<'a> Checker<'a> {
fn check_deferred_assignments(&mut self) { fn check_deferred_assignments(&mut self) {
while !self.deferred.assignments.is_empty() { while !self.deferred.assignments.is_empty() {
let assignments = std::mem::take(&mut self.deferred.assignments); let assignments = std::mem::take(&mut self.deferred.assignments);
for (scope_id, ..) in assignments { for snapshot in assignments {
self.ctx.restore(snapshot);
// pyflakes // pyflakes
if self.settings.rules.enabled(Rule::UnusedVariable) { if self.settings.rules.enabled(Rule::UnusedVariable) {
pyflakes::rules::unused_variable(self, scope_id); pyflakes::rules::unused_variable(self, self.ctx.scope_id);
} }
if self.settings.rules.enabled(Rule::UnusedAnnotation) { if self.settings.rules.enabled(Rule::UnusedAnnotation) {
pyflakes::rules::unused_annotation(self, scope_id); pyflakes::rules::unused_annotation(self, self.ctx.scope_id);
} }
if !self.is_stub { if !self.is_stub {
@ -4928,7 +4905,7 @@ impl<'a> Checker<'a> {
Rule::UnusedStaticMethodArgument, Rule::UnusedStaticMethodArgument,
Rule::UnusedLambdaArgument, Rule::UnusedLambdaArgument,
]) { ]) {
let scope = &self.ctx.scopes[scope_id]; let scope = &self.ctx.scopes[self.ctx.scope_id];
let parent = &self.ctx.scopes[scope.parent.unwrap()]; let parent = &self.ctx.scopes[scope.parent.unwrap()];
self.diagnostics self.diagnostics
.extend(flake8_unused_arguments::rules::unused_arguments( .extend(flake8_unused_arguments::rules::unused_arguments(
@ -4947,9 +4924,8 @@ impl<'a> Checker<'a> {
while !self.deferred.for_loops.is_empty() { while !self.deferred.for_loops.is_empty() {
let for_loops = std::mem::take(&mut self.deferred.for_loops); let for_loops = std::mem::take(&mut self.deferred.for_loops);
for (scope_id, stmt_id) in for_loops { for snapshot in for_loops {
self.ctx.scope_id = scope_id; self.ctx.restore(snapshot);
self.ctx.stmt_id = stmt_id;
if let StmtKind::For { target, body, .. } if let StmtKind::For { target, body, .. }
| StmtKind::AsyncFor { target, body, .. } = &self.ctx.stmt().node | StmtKind::AsyncFor { target, body, .. } = &self.ctx.stmt().node
@ -5442,9 +5418,8 @@ impl<'a> Checker<'a> {
let mut overloaded_name: Option<String> = None; let mut overloaded_name: Option<String> = None;
while !self.deferred.definitions.is_empty() { while !self.deferred.definitions.is_empty() {
let definitions = std::mem::take(&mut self.deferred.definitions); let definitions = std::mem::take(&mut self.deferred.definitions);
for (definition, visibility, (scope_id, stmt_id)) in definitions { for (definition, visibility, snapshot) in definitions {
self.ctx.scope_id = scope_id; self.ctx.restore(snapshot);
self.ctx.stmt_id = stmt_id;
// flake8-annotations // flake8-annotations
if enforce_annotations { if enforce_annotations {

View file

@ -19,6 +19,15 @@ use crate::binding::{
use crate::node::{NodeId, Nodes}; use crate::node::{NodeId, Nodes};
use crate::scope::{Scope, ScopeId, ScopeKind, Scopes}; use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
/// A snapshot of the [`Context`] at a given point in the AST traversal.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot {
scope_id: ScopeId,
stmt_id: Option<NodeId>,
in_annotation: bool,
in_type_checking_block: bool,
}
#[allow(clippy::struct_excessive_bools)] #[allow(clippy::struct_excessive_bools)]
pub struct Context<'a> { pub struct Context<'a> {
pub typing_modules: &'a [String], pub typing_modules: &'a [String],
@ -435,4 +444,22 @@ impl<'a> Context<'a> {
} }
exceptions exceptions
} }
/// Generate a [`Snapshot`] of the current context.
pub fn snapshot(&self) -> Snapshot {
Snapshot {
scope_id: self.scope_id,
stmt_id: self.stmt_id,
in_annotation: self.in_annotation,
in_type_checking_block: self.in_type_checking_block,
}
}
/// Restore the context to the given [`Snapshot`].
pub fn restore(&mut self, snapshot: Snapshot) {
self.scope_id = snapshot.scope_id;
self.stmt_id = snapshot.stmt_id;
self.in_annotation = snapshot.in_annotation;
self.in_type_checking_block = snapshot.in_type_checking_block;
}
} }

View file

@ -30,6 +30,7 @@ impl From<NodeId> for usize {
} }
} }
/// A [`Node`] represents a statement in a program, along with a pointer to its parent (if any).
#[derive(Debug)] #[derive(Debug)]
struct Node<'a> { struct Node<'a> {
/// The statement this node represents. /// The statement this node represents.