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

Commit 6f65c5c

Browse files
MichaReisercharliermarsh
MichaReiser
authored and
charliermarsh
committed
Add Decorator node (#7)
1 parent 8a415fa commit 6f65c5c

File tree

12 files changed

+2154
-1927
lines changed

12 files changed

+2154
-1927
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
fn fold_arg_with_default(
511517
&mut self,
512518
node: ArgWithDefault<U>,
@@ -2791,6 +2797,25 @@ 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+
}
27942819
impl<T, U> Foldable<T, U> for ArgWithDefault<T> {
27952820
type Mapped = ArgWithDefault<U>;
27962821
fn fold<F: Fold<T, TargetU = U> + ?Sized>(

‎ast/src/gen/generic.rs

Lines changed: 22 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
/// See also [mod](https://docs.python.org/3/library/ast.html#ast.mod)
141148
#[derive(Clone, Debug, PartialEq, is_macro::Is)]
142149
pub enum Mod<R = TextRange> {
@@ -307,7 +314,7 @@ pub struct StmtFunctionDef<R = TextRange> {
307314
pub name: Identifier,
308315
pub args: Box<Arguments<R>>,
309316
pub body: Vec<Stmt<R>>,
310-
pub decorator_list: Vec<Expr<R>>,
317+
pub decorator_list: Vec<Decorator<R>>,
311318
pub returns: Option<Box<Expr<R>>>,
312319
pub type_comment: Option<String>,
313320
}
@@ -341,7 +348,7 @@ pub struct StmtAsyncFunctionDef<R = TextRange> {
341348
pub name: Identifier,
342349
pub args: Box<Arguments<R>>,
343350
pub body: Vec<Stmt<R>>,
344-
pub decorator_list: Vec<Expr<R>>,
351+
pub decorator_list: Vec<Decorator<R>>,
345352
pub returns: Option<Box<Expr<R>>>,
346353
pub type_comment: Option<String>,
347354
}
@@ -376,7 +383,7 @@ pub struct StmtClassDef<R = TextRange> {
376383
pub bases: Vec<Expr<R>>,
377384
pub keywords: Vec<Keyword<R>>,
378385
pub body: Vec<Stmt<R>>,
379-
pub decorator_list: Vec<Expr<R>>,
386+
pub decorator_list: Vec<Decorator<R>>,
380387
}
381388

382389
impl<R> Node for StmtClassDef<R> {
@@ -3074,6 +3081,18 @@ impl<R> Node for TypeIgnore<R> {
30743081
const FIELD_NAMES: &'static [&'static str] = &[];
30753082
}
30763083

3084+
/// See also [decorator](https://docs.python.org/3/library/ast.html#ast.decorator)
3085+
#[derive(Clone, Debug, PartialEq)]
3086+
pub struct Decorator<R = TextRange> {
3087+
pub range: OptionalRange<R>,
3088+
pub expression: Expr<R>,
3089+
}
3090+
3091+
impl<R> Node for Decorator<R> {
3092+
const NAME: &'static str = "decorator";
3093+
const FIELD_NAMES: &'static [&'static str] = &["expression"];
3094+
}
3095+
30773096
/// An alternative type of AST `arguments`. This is parser-friendly and human-friendly definition of function arguments.
30783097
/// This form also has advantage to implement pre-order traverse.
30793098
/// `defaults` and `kw_defaults` fields are removed and the default values are placed under each `arg_with_default` typed argument.

‎ast/src/gen/located.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1367,6 +1367,21 @@ impl LocatedMut for TypeIgnore {
13671367
}
13681368
}
13691369

1370+
pub type Decorator = crate::generic::Decorator<SourceRange>;
1371+
1372+
#[cfg(feature = "all-nodes-with-ranges")]
1373+
impl Located for Decorator {
1374+
fn range(&self) -> SourceRange {
1375+
self.range
1376+
}
1377+
}
1378+
#[cfg(feature = "all-nodes-with-ranges")]
1379+
impl LocatedMut for Decorator {
1380+
fn range_mut(&mut self) -> &mut SourceRange {
1381+
&mut self.range
1382+
}
1383+
}
1384+
13701385
pub type Arguments = crate::generic::Arguments<SourceRange>;
13711386

13721387
#[cfg(feature = "all-nodes-with-ranges")]

‎ast/src/gen/ranged.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,12 @@ impl Ranged for crate::TypeIgnore {
496496
}
497497
}
498498

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+
}
499505
#[cfg(feature = "all-nodes-with-ranges")]
500506
impl Ranged for crate::generic::Arguments<TextRange> {
501507
fn range(&self) -> TextRange {

‎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
@@ -1102,7 +1102,7 @@ def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...
11021102
}
11031103

11041104
#[test]
1105-
#[cfg(not(feature = "all-nodes-with-ranges"))]
1105+
#[cfg(feature = "all-nodes-with-ranges")]
11061106
fn decorator_ranges() {
11071107
let parse_ast = parse_program(
11081108
r#"

‎parser/src/python.lalrpop

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

11531153
// Decorators:
1154-
Decorator: ast::Expr = {
1155-
<location:@L> "@" <p:NamedExpressionTest> "\n" => {
1156-
p
1154+
Decorator: ast::Decorator = {
1155+
<location:@L> "@" <p:NamedExpressionTest> <end_location:@R> "\n" => {
1156+
ast::Decorator { range: optional_range(location, end_location), expression: p }
11571157
},
11581158
};
11591159

0 commit comments

Comments
 (0)