Skip to content
This repository was archived by the owner on Jul 27, 2023. It is now read-only.

Commit c2a0a29

Browse files
MichaReiserMichaReiser
committed
Add Decorator node
1 parent 6c9e93f commit c2a0a29

13 files changed

+2170
-1941
lines changed

‎ast-pyo3/src/gen/to_py_ast.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -944,6 +944,14 @@ impl<R> PyNode for ast::TypeIgnoreTypeIgnore<R> {
944944
}
945945
}
946946

947+
impl<R> PyNode for ast::Decorator<R> {
948+
#[inline]
949+
fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
950+
static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
951+
&PY_TYPE
952+
}
953+
}
954+
947955
impl ToPyAst for ast::ExprContext {
948956
#[inline]
949957
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -2605,6 +2613,22 @@ impl ToPyAst for ast::TypeIgnoreTypeIgnore<TextRange> {
26052613
}
26062614
}
26072615

2616+
impl ToPyAst for ast::Decorator<TextRange> {
2617+
#[inline]
2618+
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
2619+
let cache = Self::py_type_cache().get().unwrap();
2620+
2621+
let Self {
2622+
expression,
2623+
range: _range,
2624+
} = self;
2625+
2626+
let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
2627+
2628+
Ok(instance)
2629+
}
2630+
}
2631+
26082632
impl ToPyAst for ast::Mod<SourceRange> {
26092633
#[inline]
26102634
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -4717,6 +4741,22 @@ impl ToPyAst for ast::TypeIgnoreTypeIgnore<SourceRange> {
47174741
}
47184742
}
47194743

4744+
impl ToPyAst for ast::Decorator<SourceRange> {
4745+
#[inline]
4746+
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
4747+
let cache = Self::py_type_cache().get().unwrap();
4748+
4749+
let Self {
4750+
expression,
4751+
range: _range,
4752+
} = self;
4753+
4754+
let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
4755+
4756+
Ok(instance)
4757+
}
4758+
}
4759+
47204760
fn init_types(py: Python) -> PyResult<()> {
47214761
let ast_module = PyModule::import(py, "_ast")?;
47224762
cache_py_type::<ast::Mod>(ast_module)?;
@@ -4837,5 +4877,6 @@ fn init_types(py: Python) -> PyResult<()> {
48374877
cache_py_type::<ast::PatternMatchOr>(ast_module)?;
48384878
cache_py_type::<ast::TypeIgnore>(ast_module)?;
48394879
cache_py_type::<ast::TypeIgnoreTypeIgnore>(ast_module)?;
4880+
cache_py_type::<ast::Decorator>(ast_module)?;
48404881
Ok(())
48414882
}

‎ast-pyo3/src/gen/wrapper_located.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3993,6 +3993,39 @@ impl TypeIgnoreTypeIgnore {
39933993
}
39943994
}
39953995

3996+
#[pyclass(module="rustpython_ast.located", name="_decorator", extends=super::Ast, frozen)]
3997+
#[derive(Clone, Debug)]
3998+
pub struct Decorator(pub &'static ast::Decorator<SourceRange>);
3999+
4000+
impl From<&'static ast::Decorator<SourceRange>> for Decorator {
4001+
fn from(node: &'static ast::Decorator<SourceRange>) -> Self {
4002+
Decorator(node)
4003+
}
4004+
}
4005+
4006+
impl ToPyObject for Decorator {
4007+
fn to_object(&self, py: Python) -> PyObject {
4008+
let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());
4009+
Py::new(py, initializer).unwrap().into_py(py)
4010+
}
4011+
}
4012+
4013+
impl ToPyWrapper for ast::Decorator<SourceRange> {
4014+
#[inline]
4015+
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
4016+
Ok(Decorator(self).to_object(py))
4017+
}
4018+
}
4019+
4020+
#[pymethods]
4021+
impl Decorator {
4022+
#[getter]
4023+
#[inline]
4024+
fn get_expression(&self, py: Python) -> PyResult<PyObject> {
4025+
self.0.expression.to_py_wrapper(py)
4026+
}
4027+
}
4028+
39964029
impl ToPyWrapper for ast::ExprContext {
39974030
#[inline]
39984031
fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4409,5 +4442,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
44094442
super::init_type::<PatternMatchOr, ast::PatternMatchOr>(py, m)?;
44104443
super::init_type::<TypeIgnore, ast::TypeIgnore>(py, m)?;
44114444
super::init_type::<TypeIgnoreTypeIgnore, ast::TypeIgnoreTypeIgnore>(py, m)?;
4445+
super::init_type::<Decorator, ast::Decorator>(py, m)?;
44124446
Ok(())
44134447
}

‎ast-pyo3/src/gen/wrapper_ranged.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3993,6 +3993,39 @@ impl TypeIgnoreTypeIgnore {
39933993
}
39943994
}
39953995

3996+
#[pyclass(module="rustpython_ast.ranged", name="_decorator", extends=super::Ast, frozen)]
3997+
#[derive(Clone, Debug)]
3998+
pub struct Decorator(pub &'static ast::Decorator<TextRange>);
3999+
4000+
impl From<&'static ast::Decorator<TextRange>> for Decorator {
4001+
fn from(node: &'static ast::Decorator<TextRange>) -> Self {
4002+
Decorator(node)
4003+
}
4004+
}
4005+
4006+
impl ToPyObject for Decorator {
4007+
fn to_object(&self, py: Python) -> PyObject {
4008+
let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());
4009+
Py::new(py, initializer).unwrap().into_py(py)
4010+
}
4011+
}
4012+
4013+
impl ToPyWrapper for ast::Decorator<TextRange> {
4014+
#[inline]
4015+
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
4016+
Ok(Decorator(self).to_object(py))
4017+
}
4018+
}
4019+
4020+
#[pymethods]
4021+
impl Decorator {
4022+
#[getter]
4023+
#[inline]
4024+
fn get_expression(&self, py: Python) -> PyResult<PyObject> {
4025+
self.0.expression.to_py_wrapper(py)
4026+
}
4027+
}
4028+
39964029
pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
39974030
super::init_module(py, m)?;
39984031
super::init_type::<Mod, ast::Mod>(py, m)?;
@@ -4113,5 +4146,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
41134146
super::init_type::<PatternMatchOr, ast::PatternMatchOr>(py, m)?;
41144147
super::init_type::<TypeIgnore, ast::TypeIgnore>(py, m)?;
41154148
super::init_type::<TypeIgnoreTypeIgnore, ast::TypeIgnoreTypeIgnore>(py, m)?;
4149+
super::init_type::<Decorator, ast::Decorator>(py, m)?;
41164150
Ok(())
41174151
}

‎ast/Python.asdl

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
-- ASDL's 4 builtin types are:
2-
-- identifier, int, string, constant
2+
-- identifier, int, string, constant, decorator
33

44
module Python
55
{
@@ -9,17 +9,17 @@ module Python
99
| FunctionType(expr* argtypes, expr returns)
1010

1111
stmt = FunctionDef(identifier name, arguments args,
12-
stmt* body, expr* decorator_list, expr? returns,
12+
stmt* body, decorator* decorator_list, expr? returns,
1313
string? type_comment)
1414
| AsyncFunctionDef(identifier name, arguments args,
15-
stmt* body, expr* decorator_list, expr? returns,
15+
stmt* body, decorator* decorator_list, expr? returns,
1616
string? type_comment)
1717

1818
| ClassDef(identifier name,
1919
expr* bases,
2020
keyword* keywords,
2121
stmt* body,
22-
expr* decorator_list)
22+
decorator* decorator_list)
2323
| Return(expr? value)
2424

2525
| Delete(expr* targets)
@@ -142,4 +142,6 @@ module Python
142142
attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
143143

144144
type_ignore = TypeIgnore(int lineno, string tag)
145+
146+
decorator = (expr expression)
145147
}

‎ast/src/gen/fold.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,12 @@ pub trait Fold<U> {
507507
) -> Result<TypeIgnoreTypeIgnore<Self::TargetU>, Self::Error> {
508508
fold_type_ignore_type_ignore(self, node)
509509
}
510+
fn fold_decorator(
511+
&mut self,
512+
node: Decorator<U>,
513+
) -> Result<Decorator<Self::TargetU>, Self::Error> {
514+
fold_decorator(self, node)
515+
}
510516
}
511517
impl<T, U> Foldable<T, U> for Mod<T> {
512518
type Mapped = Mod<U>;
@@ -2791,3 +2797,22 @@ pub fn fold_type_ignore_type_ignore<U, F: Fold<U> + ?Sized>(
27912797
let range = folder.map_user_cfg(range, context)?;
27922798
Ok(TypeIgnoreTypeIgnore { lineno, tag, range })
27932799
}
2800+
impl<T, U> Foldable<T, U> for Decorator<T> {
2801+
type Mapped = Decorator<U>;
2802+
fn fold<F: Fold<T, TargetU = U> + ?Sized>(
2803+
self,
2804+
folder: &mut F,
2805+
) -> Result<Self::Mapped, F::Error> {
2806+
folder.fold_decorator(self)
2807+
}
2808+
}
2809+
pub fn fold_decorator<U, F: Fold<U> + ?Sized>(
2810+
#[allow(unused)] folder: &mut F,
2811+
node: Decorator<U>,
2812+
) -> Result<Decorator<F::TargetU>, F::Error> {
2813+
let Decorator { expression, range } = node;
2814+
let context = folder.will_map_user_cfg(&range);
2815+
let expression = Foldable::fold(expression, folder)?;
2816+
let range = folder.map_user_cfg(range, context)?;
2817+
Ok(Decorator { expression, range })
2818+
}

‎ast/src/gen/generic.rs

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub enum Ast<R = TextRange> {
2323
MatchCase(MatchCase<R>),
2424
Pattern(Pattern<R>),
2525
TypeIgnore(TypeIgnore<R>),
26+
Decorator(Decorator<R>),
2627
}
2728
impl<R> Node for Ast<R> {
2829
const NAME: &'static str = "AST";
@@ -137,6 +138,12 @@ impl<R> From<TypeIgnore<R>> for Ast<R> {
137138
}
138139
}
139140

141+
impl<R> From<Decorator<R>> for Ast<R> {
142+
fn from(node: Decorator<R>) -> Self {
143+
Ast::Decorator(node)
144+
}
145+
}
146+
140147
#[derive(Clone, Debug, PartialEq)]
141148
pub struct ModModule<R = TextRange> {
142149
pub range: OptionalRange<R>,
@@ -242,7 +249,7 @@ pub struct StmtFunctionDef<R = TextRange> {
242249
pub name: Identifier,
243250
pub args: Box<Arguments<R>>,
244251
pub body: Vec<Stmt<R>>,
245-
pub decorator_list: Vec<Expr<R>>,
252+
pub decorator_list: Vec<Decorator<R>>,
246253
pub returns: Option<Box<Expr<R>>>,
247254
pub type_comment: Option<String>,
248255
}
@@ -275,7 +282,7 @@ pub struct StmtAsyncFunctionDef<R = TextRange> {
275282
pub name: Identifier,
276283
pub args: Box<Arguments<R>>,
277284
pub body: Vec<Stmt<R>>,
278-
pub decorator_list: Vec<Expr<R>>,
285+
pub decorator_list: Vec<Decorator<R>>,
279286
pub returns: Option<Box<Expr<R>>>,
280287
pub type_comment: Option<String>,
281288
}
@@ -309,7 +316,7 @@ pub struct StmtClassDef<R = TextRange> {
309316
pub bases: Vec<Expr<R>>,
310317
pub keywords: Vec<Keyword<R>>,
311318
pub body: Vec<Stmt<R>>,
312-
pub decorator_list: Vec<Expr<R>>,
319+
pub decorator_list: Vec<Decorator<R>>,
313320
}
314321

315322
impl<R> Node for StmtClassDef<R> {
@@ -2987,3 +2994,14 @@ impl<R> Node for TypeIgnore<R> {
29872994
const NAME: &'static str = "type_ignore";
29882995
const FIELD_NAMES: &'static [&'static str] = &[];
29892996
}
2997+
2998+
#[derive(Clone, Debug, PartialEq)]
2999+
pub struct Decorator<R = TextRange> {
3000+
pub range: OptionalRange<R>,
3001+
pub expression: Expr<R>,
3002+
}
3003+
3004+
impl<R> Node for Decorator<R> {
3005+
const NAME: &'static str = "decorator";
3006+
const FIELD_NAMES: &'static [&'static str] = &["expression"];
3007+
}

‎ast/src/gen/located.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -815,3 +815,13 @@ impl Located for TypeIgnore {
815815
}
816816
}
817817
}
818+
819+
pub type Decorator = crate::generic::Decorator<SourceRange>;
820+
821+
#[cfg(feature = "all-nodes-with-ranges")]
822+
823+
impl Located for Decorator {
824+
fn range(&self) -> SourceRange {
825+
self.range
826+
}
827+
}

‎ast/src/gen/ranged.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,3 +495,10 @@ impl Ranged for crate::TypeIgnore {
495495
}
496496
}
497497
}
498+
499+
#[cfg(feature = "all-nodes-with-ranges")]
500+
impl Ranged for crate::generic::Decorator<TextRange> {
501+
fn range(&self) -> TextRange {
502+
self.range
503+
}
504+
}

‎ast/src/gen/visitor.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
4848
self.visit_stmt(value);
4949
}
5050
for value in node.decorator_list {
51-
self.visit_expr(value);
51+
self.visit_decorator(value);
5252
}
5353
if let Some(value) = node.returns {
5454
self.visit_expr(*value);
@@ -66,7 +66,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
6666
self.visit_stmt(value);
6767
}
6868
for value in node.decorator_list {
69-
self.visit_expr(value);
69+
self.visit_decorator(value);
7070
}
7171
if let Some(value) = node.returns {
7272
self.visit_expr(*value);
@@ -86,7 +86,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
8686
self.visit_stmt(value);
8787
}
8888
for value in node.decorator_list {
89-
self.visit_expr(value);
89+
self.visit_decorator(value);
9090
}
9191
}
9292
fn visit_stmt_Return(&mut self, node: StmtReturn<R>) {
@@ -810,4 +810,8 @@ pub trait Visitor<R = crate::text_size::TextRange> {
810810
self.visit_pattern(value);
811811
}
812812
}
813+
fn visit_decorator(&mut self, node: Decorator<R>) {
814+
self.generic_visit_decorator(node)
815+
}
816+
fn generic_visit_decorator(&mut self, node: Decorator<R>) {}
813817
}

‎parser/src/parser.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1035,7 +1035,7 @@ def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...
10351035
}
10361036

10371037
#[test]
1038-
#[cfg(not(feature = "all-nodes-with-ranges"))]
1038+
#[cfg(feature = "all-nodes-with-ranges")]
10391039
fn decorator_ranges() {
10401040
let parse_ast = ast::Suite::parse(
10411041
r#"

‎parser/src/python.lalrpop

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1162,9 +1162,9 @@ ClassDef: ast::Stmt = {
11621162
};
11631163

11641164
// Decorators:
1165-
Decorator: ast::Expr = {
1166-
<location:@L> "@" <p:NamedExpressionTest> "\n" => {
1167-
p
1165+
Decorator: ast::Decorator = {
1166+
<location:@L> "@" <p:NamedExpressionTest> <end_location:@R> "\n" => {
1167+
ast::Decorator { range: optional_range(location, end_location), expression: p }
11681168
},
11691169
};
11701170

0 commit comments

Comments
 (0)