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

Commit 7ea94c1

Browse files
MichaReiserzanieb
MichaReiser
authored and
zanieb
committed
Add Decorator node (#7)
1 parent 403906c commit 7ea94c1

File tree

12 files changed

+2308
-2085
lines changed

12 files changed

+2308
-2085
lines changed

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -984,6 +984,14 @@ impl<R> PyNode for ast::TypeParamTypeVarTuple<R> {
984984
}
985985
}
986986

987+
impl<R> PyNode for ast::Decorator<R> {
988+
#[inline]
989+
fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
990+
static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
991+
&PY_TYPE
992+
}
993+
}
994+
987995
impl ToPyAst for ast::ExprContext {
988996
#[inline]
989997
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -2736,6 +2744,23 @@ impl ToPyAst for ast::TypeParamTypeVarTuple<TextRange> {
27362744
}
27372745
}
27382746

2747+
impl ToPyAst for ast::Decorator<TextRange> {
2748+
#[inline]
2749+
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
2750+
let cache = Self::py_type_cache().get().unwrap();
2751+
2752+
let Self {
2753+
expression,
2754+
range: _range,
2755+
} = self;
2756+
2757+
let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
2758+
2759+
Ok(instance)
2760+
}
2761+
}
2762+
2763+
27392764
impl ToPyAst for ast::Mod<SourceRange> {
27402765
#[inline]
27412766
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -4966,6 +4991,21 @@ impl ToPyAst for ast::TypeParamTypeVarTuple<SourceRange> {
49664991
instance.setattr(cache.end_lineno.as_ref(py), end.row.get())?;
49674992
instance.setattr(cache.end_col_offset.as_ref(py), end.column.get())?;
49684993
}
4994+
Ok(instance)
4995+
}
4996+
}
4997+
4998+
impl ToPyAst for ast::Decorator<SourceRange> {
4999+
#[inline]
5000+
fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
5001+
let cache = Self::py_type_cache().get().unwrap();
5002+
5003+
let Self {
5004+
expression,
5005+
range: _range,
5006+
} = self;
5007+
5008+
let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
49695009

49705010
Ok(instance)
49715011
}
@@ -5096,5 +5136,6 @@ fn init_types(py: Python) -> PyResult<()> {
50965136
cache_py_type::<ast::TypeParamTypeVar>(ast_module)?;
50975137
cache_py_type::<ast::TypeParamParamSpec>(ast_module)?;
50985138
cache_py_type::<ast::TypeParamTypeVarTuple>(ast_module)?;
5139+
cache_py_type::<ast::Decorator>(ast_module)?;
50995140
Ok(())
51005141
}

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4117,6 +4117,27 @@ impl ToPyWrapper for ast::TypeParamTypeVar<SourceRange> {
41174117
#[inline]
41184118
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
41194119
Ok(TypeParamTypeVar(self).to_object(py))
4120+
#[pyclass(module="rustpython_ast.located", name="_decorator", extends=super::Ast, frozen)]
4121+
#[derive(Clone, Debug)]
4122+
pub struct Decorator(pub &'static ast::Decorator<SourceRange>);
4123+
4124+
impl From<&'static ast::Decorator<SourceRange>> for Decorator {
4125+
fn from(node: &'static ast::Decorator<SourceRange>) -> Self {
4126+
Decorator(node)
4127+
}
4128+
}
4129+
4130+
impl ToPyObject for Decorator {
4131+
fn to_object(&self, py: Python) -> PyObject {
4132+
let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());
4133+
Py::new(py, initializer).unwrap().into_py(py)
4134+
}
4135+
}
4136+
4137+
impl ToPyWrapper for ast::Decorator<SourceRange> {
4138+
#[inline]
4139+
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
4140+
Ok(Decorator(self).to_object(py))
41204141
}
41214142
}
41224143

@@ -4205,6 +4226,14 @@ impl TypeParamTypeVarTuple {
42054226
}
42064227
}
42074228

4229+
impl Decorator {
4230+
#[getter]
4231+
#[inline]
4232+
fn get_expression(&self, py: Python) -> PyResult<PyObject> {
4233+
self.0.expression.to_py_wrapper(py)
4234+
}
4235+
}
4236+
42084237
impl ToPyWrapper for ast::ExprContext {
42094238
#[inline]
42104239
fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4626,5 +4655,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
46264655
super::init_type::<TypeParamTypeVar, ast::TypeParamTypeVar>(py, m)?;
46274656
super::init_type::<TypeParamParamSpec, ast::TypeParamParamSpec>(py, m)?;
46284657
super::init_type::<TypeParamTypeVarTuple, ast::TypeParamTypeVarTuple>(py, m)?;
4658+
super::init_type::<Decorator, ast::Decorator>(py, m)?;
46294659
Ok(())
46304660
}

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4059,6 +4059,22 @@ impl TypeIgnoreTypeIgnore {
40594059
}
40604060
}
40614061

4062+
#[pyclass(module="rustpython_ast.ranged", name="_decorator", extends=super::Ast, frozen)]
4063+
#[derive(Clone, Debug)]
4064+
pub struct Decorator(pub &'static ast::Decorator<TextRange>);
4065+
4066+
impl From<&'static ast::Decorator<TextRange>> for Decorator {
4067+
fn from(node: &'static ast::Decorator<TextRange>) -> Self {
4068+
Decorator(node)
4069+
}
4070+
}
4071+
4072+
impl ToPyObject for Decorator {
4073+
fn to_object(&self, py: Python) -> PyObject {
4074+
let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());
4075+
Py::new(py, initializer).unwrap().into_py(py)
4076+
}
4077+
}
40624078
#[pyclass(module="rustpython_ast.ranged", name="_type_param", extends=super::Ast, frozen, subclass)]
40634079
#[derive(Clone, Debug)]
40644080
pub struct TypeParam;
@@ -4083,6 +4099,12 @@ impl ToPyObject for TypeParam {
40834099
}
40844100
}
40854101

4102+
impl ToPyWrapper for ast::Decorator<TextRange> {
4103+
#[inline]
4104+
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
4105+
Ok(Decorator(self).to_object(py))
4106+
}
4107+
}
40864108
impl ToPyWrapper for ast::TypeParam<TextRange> {
40874109
#[inline]
40884110
fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4121,6 +4143,13 @@ impl ToPyWrapper for ast::TypeParamTypeVar<TextRange> {
41214143
}
41224144

41234145
#[pymethods]
4146+
impl Decorator {
4147+
#[getter]
4148+
#[inline]
4149+
fn get_expression(&self, py: Python) -> PyResult<PyObject> {
4150+
self.0.expression.to_py_wrapper(py)
4151+
}
4152+
}
41244153
impl TypeParamTypeVar {
41254154
#[getter]
41264155
#[inline]
@@ -4330,5 +4359,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
43304359
super::init_type::<TypeParamTypeVar, ast::TypeParamTypeVar>(py, m)?;
43314360
super::init_type::<TypeParamParamSpec, ast::TypeParamParamSpec>(py, m)?;
43324361
super::init_type::<TypeParamTypeVarTuple, ast::TypeParamTypeVarTuple>(py, m)?;
4362+
super::init_type::<Decorator, ast::Decorator>(py, m)?;
43334363
Ok(())
43344364
}

‎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, type_param* type_params)
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, type_param* type_params)
1717

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

@@ -149,4 +149,6 @@ module Python
149149
| ParamSpec(identifier name)
150150
| TypeVarTuple(identifier name)
151151
attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
152+
153+
decorator = (expr expression)
152154
}

‎ast/src/gen/fold.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,12 @@ pub trait Fold<U> {
537537
) -> Result<TypeParamTypeVarTuple<Self::TargetU>, Self::Error> {
538538
fold_type_param_type_var_tuple(self, node)
539539
}
540+
fn fold_decorator(
541+
&mut self,
542+
node: Decorator<U>,
543+
) -> Result<Decorator<Self::TargetU>, Self::Error> {
544+
fold_decorator(self, node)
545+
}
540546
fn fold_arg_with_default(
541547
&mut self,
542548
node: ArgWithDefault<U>,
@@ -2944,6 +2950,25 @@ pub fn fold_type_param_type_var_tuple<U, F: Fold<U> + ?Sized>(
29442950
let range = folder.map_user(range, context)?;
29452951
Ok(TypeParamTypeVarTuple { name, range })
29462952
}
2953+
impl<T, U> Foldable<T, U> for Decorator<T> {
2954+
type Mapped = Decorator<U>;
2955+
fn fold<F: Fold<T, TargetU = U> + ?Sized>(
2956+
self,
2957+
folder: &mut F,
2958+
) -> Result<Self::Mapped, F::Error> {
2959+
folder.fold_decorator(self)
2960+
}
2961+
}
2962+
pub fn fold_decorator<U, F: Fold<U> + ?Sized>(
2963+
#[allow(unused)] folder: &mut F,
2964+
node: Decorator<U>,
2965+
) -> Result<Decorator<F::TargetU>, F::Error> {
2966+
let Decorator { expression, range } = node;
2967+
let context = folder.will_map_user_cfg(&range);
2968+
let expression = Foldable::fold(expression, folder)?;
2969+
let range = folder.map_user_cfg(range, context)?;
2970+
Ok(Decorator { expression, range })
2971+
}
29472972
impl<T, U> Foldable<T, U> for ArgWithDefault<T> {
29482973
type Mapped = ArgWithDefault<U>;
29492974
fn fold<F: Fold<T, TargetU = U> + ?Sized>(

‎ast/src/gen/generic.rs

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ pub enum Ast<R = TextRange> {
2323
Pattern(Pattern<R>),
2424
TypeIgnore(TypeIgnore<R>),
2525
TypeParam(TypeParam<R>),
26+
Decorator(Decorator<R>),
2627
}
28+
2729
impl<R> Node for Ast<R> {
2830
const NAME: &'static str = "AST";
2931
const FIELD_NAMES: &'static [&'static str] = &[];
@@ -143,6 +145,12 @@ impl<R> From<TypeParam<R>> for Ast<R> {
143145
}
144146
}
145147

148+
impl<R> From<Decorator<R>> for Ast<R> {
149+
fn from(node: Decorator<R>) -> Self {
150+
Ast::Decorator(node)
151+
}
152+
}
153+
146154
/// See also [mod](https://docs.python.org/3/library/ast.html#ast.mod)
147155
#[derive(Clone, Debug, PartialEq, is_macro::Is)]
148156
pub enum Mod<R = TextRange> {
@@ -315,7 +323,7 @@ pub struct StmtFunctionDef<R = TextRange> {
315323
pub name: Identifier,
316324
pub args: Box<Arguments<R>>,
317325
pub body: Vec<Stmt<R>>,
318-
pub decorator_list: Vec<Expr<R>>,
326+
pub decorator_list: Vec<Decorator<R>>,
319327
pub returns: Option<Box<Expr<R>>>,
320328
pub type_comment: Option<String>,
321329
pub type_params: Vec<TypeParam<R>>,
@@ -351,7 +359,7 @@ pub struct StmtAsyncFunctionDef<R = TextRange> {
351359
pub name: Identifier,
352360
pub args: Box<Arguments<R>>,
353361
pub body: Vec<Stmt<R>>,
354-
pub decorator_list: Vec<Expr<R>>,
362+
pub decorator_list: Vec<Decorator<R>>,
355363
pub returns: Option<Box<Expr<R>>>,
356364
pub type_comment: Option<String>,
357365
pub type_params: Vec<TypeParam<R>>,
@@ -388,7 +396,7 @@ pub struct StmtClassDef<R = TextRange> {
388396
pub bases: Vec<Expr<R>>,
389397
pub keywords: Vec<Keyword<R>>,
390398
pub body: Vec<Stmt<R>>,
391-
pub decorator_list: Vec<Expr<R>>,
399+
pub decorator_list: Vec<Decorator<R>>,
392400
pub type_params: Vec<TypeParam<R>>,
393401
}
394402

@@ -3197,6 +3205,18 @@ impl<R> Node for TypeParam<R> {
31973205
const FIELD_NAMES: &'static [&'static str] = &[];
31983206
}
31993207

3208+
/// See also [decorator](https://docs.python.org/3/library/ast.html#ast.decorator)
3209+
#[derive(Clone, Debug, PartialEq)]
3210+
pub struct Decorator<R = TextRange> {
3211+
pub range: OptionalRange<R>,
3212+
pub expression: Expr<R>,
3213+
}
3214+
3215+
impl<R> Node for Decorator<R> {
3216+
const NAME: &'static str = "decorator";
3217+
const FIELD_NAMES: &'static [&'static str] = &["expression"];
3218+
}
3219+
32003220
/// An alternative type of AST `arguments`. This is parser-friendly and human-friendly definition of function arguments.
32013221
/// This form also has advantage to implement pre-order traverse.
32023222
/// `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: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1383,6 +1383,23 @@ impl LocatedMut for TypeIgnore {
13831383
}
13841384
}
13851385

1386+
1387+
pub type Decorator = crate::generic::Decorator<SourceRange>;
1388+
1389+
#[cfg(feature = "all-nodes-with-ranges")]
1390+
impl Located for Decorator {
1391+
fn range(&self) -> SourceRange {
1392+
self.range
1393+
}
1394+
}
1395+
1396+
#[cfg(feature = "all-nodes-with-ranges")]
1397+
impl LocatedMut for Decorator {
1398+
fn range_mut(&mut self) -> &mut SourceRange {
1399+
&mut self.range
1400+
}
1401+
}
1402+
13861403
pub type TypeParam = crate::generic::TypeParam<SourceRange>;
13871404

13881405
pub type TypeParamTypeVar = crate::generic::TypeParamTypeVar<SourceRange>;

‎ast/src/gen/ranged.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,12 @@ impl Ranged for crate::TypeParam {
527527
}
528528
}
529529

530+
#[cfg(feature = "all-nodes-with-ranges")]
531+
impl Ranged for crate::generic::Decorator<TextRange> {
532+
fn range(&self) -> TextRange {
533+
self.range
534+
}
535+
}
530536
#[cfg(feature = "all-nodes-with-ranges")]
531537
impl Ranged for crate::generic::Arguments<TextRange> {
532538
fn range(&self) -> TextRange {

‎ast/src/gen/visitor.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
4949
self.visit_stmt(value);
5050
}
5151
for value in node.decorator_list {
52-
self.visit_expr(value);
52+
self.visit_decorator(value);
5353
}
5454
if let Some(value) = node.returns {
5555
self.visit_expr(*value);
@@ -70,7 +70,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
7070
self.visit_stmt(value);
7171
}
7272
for value in node.decorator_list {
73-
self.visit_expr(value);
73+
self.visit_decorator(value);
7474
}
7575
if let Some(value) = node.returns {
7676
self.visit_expr(*value);
@@ -93,7 +93,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
9393
self.visit_stmt(value);
9494
}
9595
for value in node.decorator_list {
96-
self.visit_expr(value);
96+
self.visit_decorator(value);
9797
}
9898
for value in node.type_params {
9999
self.visit_type_param(value);
@@ -862,4 +862,8 @@ pub trait Visitor<R = crate::text_size::TextRange> {
862862
self.generic_visit_type_param_type_var_tuple(node)
863863
}
864864
fn generic_visit_type_param_type_var_tuple(&mut self, node: TypeParamTypeVarTuple<R>) {}
865+
fn visit_decorator(&mut self, node: Decorator<R>) {
866+
self.generic_visit_decorator(node)
867+
}
868+
fn generic_visit_decorator(&mut self, node: Decorator<R>) {}
865869
}

‎parser/src/parser.rs

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

12411241
#[test]
1242-
#[cfg(not(feature = "all-nodes-with-ranges"))]
1242+
#[cfg(feature = "all-nodes-with-ranges")]
12431243
fn decorator_ranges() {
12441244
let parse_ast = parse_program(
12451245
r#"

0 commit comments

Comments
 (0)