Skip to content

Commit

Permalink
Introduce ExpressionRef (#6637)
Browse files Browse the repository at this point in the history
## Summary

This PR revives the `ExpressionRef` concept introduced in
#5644, motivated by the change we
want to make in #6575 to narrow
the type of the expression that can be passed to `parenthesized_range`.

## Test Plan

`cargo test`
  • Loading branch information
charliermarsh authored Aug 17, 2023
1 parent fa7442d commit 1334232
Show file tree
Hide file tree
Showing 6 changed files with 317 additions and 29 deletions.
291 changes: 291 additions & 0 deletions crates/ruff_python_ast/src/expression.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,291 @@
use ruff_text_size::TextRange;

use crate::node::AnyNodeRef;
use crate::{self as ast, Expr, Ranged};

/// Unowned pendant to [`ast::Expr`] that stores a reference instead of a owned value.
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum ExpressionRef<'a> {
BoolOp(&'a ast::ExprBoolOp),
NamedExpr(&'a ast::ExprNamedExpr),
BinOp(&'a ast::ExprBinOp),
UnaryOp(&'a ast::ExprUnaryOp),
Lambda(&'a ast::ExprLambda),
IfExp(&'a ast::ExprIfExp),
Dict(&'a ast::ExprDict),
Set(&'a ast::ExprSet),
ListComp(&'a ast::ExprListComp),
SetComp(&'a ast::ExprSetComp),
DictComp(&'a ast::ExprDictComp),
GeneratorExp(&'a ast::ExprGeneratorExp),
Await(&'a ast::ExprAwait),
Yield(&'a ast::ExprYield),
YieldFrom(&'a ast::ExprYieldFrom),
Compare(&'a ast::ExprCompare),
Call(&'a ast::ExprCall),
FormattedValue(&'a ast::ExprFormattedValue),
FString(&'a ast::ExprFString),
Constant(&'a ast::ExprConstant),
Attribute(&'a ast::ExprAttribute),
Subscript(&'a ast::ExprSubscript),
Starred(&'a ast::ExprStarred),
Name(&'a ast::ExprName),
List(&'a ast::ExprList),
Tuple(&'a ast::ExprTuple),
Slice(&'a ast::ExprSlice),
IpyEscapeCommand(&'a ast::ExprIpyEscapeCommand),
}

impl<'a> From<&'a Box<Expr>> for ExpressionRef<'a> {
fn from(value: &'a Box<Expr>) -> Self {
ExpressionRef::from(value.as_ref())
}
}

impl<'a> From<&'a Expr> for ExpressionRef<'a> {
fn from(value: &'a Expr) -> Self {
match value {
Expr::BoolOp(value) => ExpressionRef::BoolOp(value),
Expr::NamedExpr(value) => ExpressionRef::NamedExpr(value),
Expr::BinOp(value) => ExpressionRef::BinOp(value),
Expr::UnaryOp(value) => ExpressionRef::UnaryOp(value),
Expr::Lambda(value) => ExpressionRef::Lambda(value),
Expr::IfExp(value) => ExpressionRef::IfExp(value),
Expr::Dict(value) => ExpressionRef::Dict(value),
Expr::Set(value) => ExpressionRef::Set(value),
Expr::ListComp(value) => ExpressionRef::ListComp(value),
Expr::SetComp(value) => ExpressionRef::SetComp(value),
Expr::DictComp(value) => ExpressionRef::DictComp(value),
Expr::GeneratorExp(value) => ExpressionRef::GeneratorExp(value),
Expr::Await(value) => ExpressionRef::Await(value),
Expr::Yield(value) => ExpressionRef::Yield(value),
Expr::YieldFrom(value) => ExpressionRef::YieldFrom(value),
Expr::Compare(value) => ExpressionRef::Compare(value),
Expr::Call(value) => ExpressionRef::Call(value),
Expr::FormattedValue(value) => ExpressionRef::FormattedValue(value),
Expr::FString(value) => ExpressionRef::FString(value),
Expr::Constant(value) => ExpressionRef::Constant(value),
Expr::Attribute(value) => ExpressionRef::Attribute(value),
Expr::Subscript(value) => ExpressionRef::Subscript(value),
Expr::Starred(value) => ExpressionRef::Starred(value),
Expr::Name(value) => ExpressionRef::Name(value),
Expr::List(value) => ExpressionRef::List(value),
Expr::Tuple(value) => ExpressionRef::Tuple(value),
Expr::Slice(value) => ExpressionRef::Slice(value),
Expr::IpyEscapeCommand(value) => ExpressionRef::IpyEscapeCommand(value),
}
}
}

impl<'a> From<&'a ast::ExprBoolOp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprBoolOp) -> Self {
Self::BoolOp(value)
}
}
impl<'a> From<&'a ast::ExprNamedExpr> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprNamedExpr) -> Self {
Self::NamedExpr(value)
}
}
impl<'a> From<&'a ast::ExprBinOp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprBinOp) -> Self {
Self::BinOp(value)
}
}
impl<'a> From<&'a ast::ExprUnaryOp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprUnaryOp) -> Self {
Self::UnaryOp(value)
}
}
impl<'a> From<&'a ast::ExprLambda> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprLambda) -> Self {
Self::Lambda(value)
}
}
impl<'a> From<&'a ast::ExprIfExp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprIfExp) -> Self {
Self::IfExp(value)
}
}
impl<'a> From<&'a ast::ExprDict> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprDict) -> Self {
Self::Dict(value)
}
}
impl<'a> From<&'a ast::ExprSet> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprSet) -> Self {
Self::Set(value)
}
}
impl<'a> From<&'a ast::ExprListComp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprListComp) -> Self {
Self::ListComp(value)
}
}
impl<'a> From<&'a ast::ExprSetComp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprSetComp) -> Self {
Self::SetComp(value)
}
}
impl<'a> From<&'a ast::ExprDictComp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprDictComp) -> Self {
Self::DictComp(value)
}
}
impl<'a> From<&'a ast::ExprGeneratorExp> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprGeneratorExp) -> Self {
Self::GeneratorExp(value)
}
}
impl<'a> From<&'a ast::ExprAwait> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprAwait) -> Self {
Self::Await(value)
}
}
impl<'a> From<&'a ast::ExprYield> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprYield) -> Self {
Self::Yield(value)
}
}
impl<'a> From<&'a ast::ExprYieldFrom> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprYieldFrom) -> Self {
Self::YieldFrom(value)
}
}
impl<'a> From<&'a ast::ExprCompare> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprCompare) -> Self {
Self::Compare(value)
}
}
impl<'a> From<&'a ast::ExprCall> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprCall) -> Self {
Self::Call(value)
}
}
impl<'a> From<&'a ast::ExprFormattedValue> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprFormattedValue) -> Self {
Self::FormattedValue(value)
}
}
impl<'a> From<&'a ast::ExprFString> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprFString) -> Self {
Self::FString(value)
}
}
impl<'a> From<&'a ast::ExprConstant> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprConstant) -> Self {
Self::Constant(value)
}
}
impl<'a> From<&'a ast::ExprAttribute> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprAttribute) -> Self {
Self::Attribute(value)
}
}
impl<'a> From<&'a ast::ExprSubscript> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprSubscript) -> Self {
Self::Subscript(value)
}
}
impl<'a> From<&'a ast::ExprStarred> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprStarred) -> Self {
Self::Starred(value)
}
}
impl<'a> From<&'a ast::ExprName> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprName) -> Self {
Self::Name(value)
}
}
impl<'a> From<&'a ast::ExprList> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprList) -> Self {
Self::List(value)
}
}
impl<'a> From<&'a ast::ExprTuple> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprTuple) -> Self {
Self::Tuple(value)
}
}
impl<'a> From<&'a ast::ExprSlice> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprSlice) -> Self {
Self::Slice(value)
}
}
impl<'a> From<&'a ast::ExprIpyEscapeCommand> for ExpressionRef<'a> {
fn from(value: &'a ast::ExprIpyEscapeCommand) -> Self {
Self::IpyEscapeCommand(value)
}
}

impl<'a> From<ExpressionRef<'a>> for AnyNodeRef<'a> {
fn from(value: ExpressionRef<'a>) -> Self {
match value {
ExpressionRef::BoolOp(expression) => AnyNodeRef::ExprBoolOp(expression),
ExpressionRef::NamedExpr(expression) => AnyNodeRef::ExprNamedExpr(expression),
ExpressionRef::BinOp(expression) => AnyNodeRef::ExprBinOp(expression),
ExpressionRef::UnaryOp(expression) => AnyNodeRef::ExprUnaryOp(expression),
ExpressionRef::Lambda(expression) => AnyNodeRef::ExprLambda(expression),
ExpressionRef::IfExp(expression) => AnyNodeRef::ExprIfExp(expression),
ExpressionRef::Dict(expression) => AnyNodeRef::ExprDict(expression),
ExpressionRef::Set(expression) => AnyNodeRef::ExprSet(expression),
ExpressionRef::ListComp(expression) => AnyNodeRef::ExprListComp(expression),
ExpressionRef::SetComp(expression) => AnyNodeRef::ExprSetComp(expression),
ExpressionRef::DictComp(expression) => AnyNodeRef::ExprDictComp(expression),
ExpressionRef::GeneratorExp(expression) => AnyNodeRef::ExprGeneratorExp(expression),
ExpressionRef::Await(expression) => AnyNodeRef::ExprAwait(expression),
ExpressionRef::Yield(expression) => AnyNodeRef::ExprYield(expression),
ExpressionRef::YieldFrom(expression) => AnyNodeRef::ExprYieldFrom(expression),
ExpressionRef::Compare(expression) => AnyNodeRef::ExprCompare(expression),
ExpressionRef::Call(expression) => AnyNodeRef::ExprCall(expression),
ExpressionRef::FormattedValue(expression) => AnyNodeRef::ExprFormattedValue(expression),
ExpressionRef::FString(expression) => AnyNodeRef::ExprFString(expression),
ExpressionRef::Constant(expression) => AnyNodeRef::ExprConstant(expression),
ExpressionRef::Attribute(expression) => AnyNodeRef::ExprAttribute(expression),
ExpressionRef::Subscript(expression) => AnyNodeRef::ExprSubscript(expression),
ExpressionRef::Starred(expression) => AnyNodeRef::ExprStarred(expression),
ExpressionRef::Name(expression) => AnyNodeRef::ExprName(expression),
ExpressionRef::List(expression) => AnyNodeRef::ExprList(expression),
ExpressionRef::Tuple(expression) => AnyNodeRef::ExprTuple(expression),
ExpressionRef::Slice(expression) => AnyNodeRef::ExprSlice(expression),
ExpressionRef::IpyEscapeCommand(expression) => {
AnyNodeRef::ExprIpyEscapeCommand(expression)
}
}
}
}

impl Ranged for ExpressionRef<'_> {
fn range(&self) -> TextRange {
match self {
ExpressionRef::BoolOp(expression) => expression.range(),
ExpressionRef::NamedExpr(expression) => expression.range(),
ExpressionRef::BinOp(expression) => expression.range(),
ExpressionRef::UnaryOp(expression) => expression.range(),
ExpressionRef::Lambda(expression) => expression.range(),
ExpressionRef::IfExp(expression) => expression.range(),
ExpressionRef::Dict(expression) => expression.range(),
ExpressionRef::Set(expression) => expression.range(),
ExpressionRef::ListComp(expression) => expression.range(),
ExpressionRef::SetComp(expression) => expression.range(),
ExpressionRef::DictComp(expression) => expression.range(),
ExpressionRef::GeneratorExp(expression) => expression.range(),
ExpressionRef::Await(expression) => expression.range(),
ExpressionRef::Yield(expression) => expression.range(),
ExpressionRef::YieldFrom(expression) => expression.range(),
ExpressionRef::Compare(expression) => expression.range(),
ExpressionRef::Call(expression) => expression.range(),
ExpressionRef::FormattedValue(expression) => expression.range(),
ExpressionRef::FString(expression) => expression.range(),
ExpressionRef::Constant(expression) => expression.range(),
ExpressionRef::Attribute(expression) => expression.range(),
ExpressionRef::Subscript(expression) => expression.range(),
ExpressionRef::Starred(expression) => expression.range(),
ExpressionRef::Name(expression) => expression.range(),
ExpressionRef::List(expression) => expression.range(),
ExpressionRef::Tuple(expression) => expression.range(),
ExpressionRef::Slice(expression) => expression.range(),
ExpressionRef::IpyEscapeCommand(expression) => expression.range(),
}
}
}
2 changes: 2 additions & 0 deletions crates/ruff_python_ast/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ pub mod all;
pub mod call_path;
pub mod comparable;
pub mod docstrings;
mod expression;
pub mod hashable;
pub mod helpers;
pub mod identifier;
Expand All @@ -20,6 +21,7 @@ pub mod types;
pub mod visitor;
pub mod whitespace;

pub use expression::*;
pub use nodes::*;

pub trait Ranged {
Expand Down
13 changes: 5 additions & 8 deletions crates/ruff_python_formatter/src/expression/expr_bin_op.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
use std::iter;

use smallvec::SmallVec;

use ruff_formatter::{format_args, write, FormatOwnedWithRule, FormatRefWithRule};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{
Constant, Expr, ExprAttribute, ExprBinOp, ExprConstant, ExprUnaryOp, Operator, StringConstant,
UnaryOp,
};
use smallvec::SmallVec;

use ruff_formatter::{format_args, write, FormatOwnedWithRule, FormatRefWithRule};
use ruff_python_ast::node::{AnyNodeRef, AstNode};

use crate::comments::{trailing_comments, trailing_node_comments, SourceComment};
use crate::expression::expr_constant::ExprConstantLayout;
Expand Down Expand Up @@ -73,10 +73,7 @@ impl FormatNodeRule<ExprBinOp> for FormatExprBinOp {
let binary_chain: SmallVec<[&ExprBinOp; 4]> =
iter::successors(Some(item), |parent| {
parent.left.as_bin_op_expr().and_then(|bin_expression| {
if is_expression_parenthesized(
bin_expression.as_any_node_ref(),
source,
) {
if is_expression_parenthesized(bin_expression.into(), source) {
None
} else {
Some(bin_expression)
Expand Down
12 changes: 5 additions & 7 deletions crates/ruff_python_formatter/src/expression/expr_bool_op.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{BoolOp, Expr, ExprBoolOp};

use crate::comments::leading_comments;
use crate::expression::parentheses::{
in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
OptionalParentheses,
};
use crate::prelude::*;
use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
use ruff_python_ast::node::{AnyNodeRef, AstNode};
use ruff_python_ast::{BoolOp, Expr, ExprBoolOp};

use super::parentheses::is_expression_parenthesized;

Expand Down Expand Up @@ -95,10 +96,7 @@ impl Format<PyFormatContext<'_>> for FormatValue<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
match self.value {
Expr::BoolOp(bool_op)
if !is_expression_parenthesized(
bool_op.as_any_node_ref(),
f.context().source(),
) =>
if !is_expression_parenthesized(bool_op.into(), f.context().source()) =>
{
// Mark chained boolean operations e.g. `x and y or z` and avoid creating a new group
write!(f, [bool_op.format().with_options(BoolOpLayout::Chained)])
Expand Down
Loading

0 comments on commit 1334232

Please sign in to comment.