Add Decorator node (#7)

Author: MichaReiser

ast-pyo3/src/gen/to_py_ast.rs

@@ -984,6 +984,14 @@ impl<R> PyNode for ast::TypeParamTypeVarTuple<R> {
     }
 }
 
+impl<R> PyNode for ast::Decorator<R> {
+    #[inline]
+    fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
+        static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
+        &PY_TYPE
+    }
+}
+
 impl ToPyAst for ast::ExprContext {
     #[inline]
     fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -2736,6 +2744,23 @@ impl ToPyAst for ast::TypeParamTypeVarTuple<TextRange> {
     }
 }
 
+impl ToPyAst for ast::Decorator<TextRange> {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        let cache = Self::py_type_cache().get().unwrap();
+
+        let Self {
+            expression,
+            range: _range,
+        } = self;
+
+        let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
+
+        Ok(instance)
+    }
+}
+
+
 impl ToPyAst for ast::Mod<SourceRange> {
     #[inline]
     fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
@@ -4966,6 +4991,21 @@ impl ToPyAst for ast::TypeParamTypeVarTuple<SourceRange> {
             instance.setattr(cache.end_lineno.as_ref(py), end.row.get())?;
             instance.setattr(cache.end_col_offset.as_ref(py), end.column.get())?;
         }
+        Ok(instance)
+    }
+}
+
+impl ToPyAst for ast::Decorator<SourceRange> {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        let cache = Self::py_type_cache().get().unwrap();
+
+        let Self {
+            expression,
+            range: _range,
+        } = self;
+
+        let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((expression.to_py_ast(py)?,))?;
 
         Ok(instance)
     }
@@ -5096,5 +5136,6 @@ fn init_types(py: Python) -> PyResult<()> {
     cache_py_type::<ast::TypeParamTypeVar>(ast_module)?;
     cache_py_type::<ast::TypeParamParamSpec>(ast_module)?;
     cache_py_type::<ast::TypeParamTypeVarTuple>(ast_module)?;
+    cache_py_type::<ast::Decorator>(ast_module)?;
     Ok(())
 }

ast-pyo3/src/gen/wrapper_located.rs

@@ -4117,6 +4117,27 @@ impl ToPyWrapper for ast::TypeParamTypeVar<SourceRange> {
     #[inline]
     fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
         Ok(TypeParamTypeVar(self).to_object(py))
+#[pyclass(module="rustpython_ast.located", name="_decorator", extends=super::Ast, frozen)]
+#[derive(Clone, Debug)]
+pub struct Decorator(pub &'static ast::Decorator<SourceRange>);
+
+impl From<&'static ast::Decorator<SourceRange>> for Decorator {
+    fn from(node: &'static ast::Decorator<SourceRange>) -> Self {
+        Decorator(node)
+    }
+}
+
+impl ToPyObject for Decorator {
+    fn to_object(&self, py: Python) -> PyObject {
+        let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());
+        Py::new(py, initializer).unwrap().into_py(py)
+    }
+}
+
+impl ToPyWrapper for ast::Decorator<SourceRange> {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(Decorator(self).to_object(py))
     }
 }
 
@@ -4205,6 +4226,14 @@ impl TypeParamTypeVarTuple {
     }
 }
 
+impl Decorator {
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
+}
+
 impl ToPyWrapper for ast::ExprContext {
     #[inline]
     fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4626,5 +4655,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
     super::init_type::<TypeParamTypeVar, ast::TypeParamTypeVar>(py, m)?;
     super::init_type::<TypeParamParamSpec, ast::TypeParamParamSpec>(py, m)?;
     super::init_type::<TypeParamTypeVarTuple, ast::TypeParamTypeVarTuple>(py, m)?;
+    super::init_type::<Decorator, ast::Decorator>(py, m)?;
     Ok(())
 }

ast-pyo3/src/gen/wrapper_ranged.rs

@@ -4059,6 +4059,22 @@ impl TypeIgnoreTypeIgnore {
     }
 }
 
+#[pyclass(module="rustpython_ast.ranged", name="_decorator", extends=super::Ast, frozen)]
+#[derive(Clone, Debug)]
+pub struct Decorator(pub &'static ast::Decorator<TextRange>);
+
+impl From<&'static ast::Decorator<TextRange>> for Decorator {
+    fn from(node: &'static ast::Decorator<TextRange>) -> Self {
+        Decorator(node)
+    }
+}
+
+impl ToPyObject for Decorator {
+    fn to_object(&self, py: Python) -> PyObject {
+        let initializer = PyClassInitializer::from(Ast).add_subclass(self.clone());        
+        Py::new(py, initializer).unwrap().into_py(py)
+    }
+}
 #[pyclass(module="rustpython_ast.ranged", name="_type_param", extends=super::Ast, frozen, subclass)]
 #[derive(Clone, Debug)]
 pub struct TypeParam;
@@ -4083,6 +4099,12 @@ impl ToPyObject for TypeParam {
     }
 }
 
+impl ToPyWrapper for ast::Decorator<TextRange> {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(Decorator(self).to_object(py))
+    }
+}
 impl ToPyWrapper for ast::TypeParam<TextRange> {
     #[inline]
     fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4121,6 +4143,13 @@ impl ToPyWrapper for ast::TypeParamTypeVar<TextRange> {
 }
 
 #[pymethods]
+impl Decorator {
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
+}
 impl TypeParamTypeVar {
     #[getter]
     #[inline]
@@ -4330,5 +4359,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
     super::init_type::<TypeParamTypeVar, ast::TypeParamTypeVar>(py, m)?;
     super::init_type::<TypeParamParamSpec, ast::TypeParamParamSpec>(py, m)?;
     super::init_type::<TypeParamTypeVarTuple, ast::TypeParamTypeVarTuple>(py, m)?;
+    super::init_type::<Decorator, ast::Decorator>(py, m)?;
     Ok(())
 }

ast/Python.asdl

@@ -1,5 +1,5 @@
 -- ASDL's 4 builtin types are:
--- identifier, int, string, constant
+-- identifier, int, string, constant, decorator
 
 module Python
 {
@@ -9,17 +9,17 @@ module Python
         | FunctionType(expr* argtypes, expr returns)
 
     stmt = FunctionDef(identifier name, arguments args,
-                       stmt* body, expr* decorator_list, expr? returns,
+                       stmt* body, decorator* decorator_list, expr? returns,
                        string? type_comment, type_param* type_params)
           | AsyncFunctionDef(identifier name, arguments args,
-                             stmt* body, expr* decorator_list, expr? returns,
+                             stmt* body, decorator* decorator_list, expr? returns,
                              string? type_comment, type_param* type_params)
 
           | ClassDef(identifier name,
              expr* bases,
              keyword* keywords,
              stmt* body,
-             expr* decorator_list,
+             decorator* decorator_list,
              type_param* type_params)
           | Return(expr? value)
 
@@ -149,4 +149,6 @@ module Python
                | ParamSpec(identifier name)
                | TypeVarTuple(identifier name)
                attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
+               
+    decorator = (expr expression)
 }

ast/src/gen/fold.rs

@@ -537,6 +537,12 @@ pub trait Fold<U> {
     ) -> Result<TypeParamTypeVarTuple<Self::TargetU>, Self::Error> {
         fold_type_param_type_var_tuple(self, node)
     }
+    fn fold_decorator(
+        &mut self,
+        node: Decorator<U>,
+    ) -> Result<Decorator<Self::TargetU>, Self::Error> {
+        fold_decorator(self, node)
+    }
     fn fold_arg_with_default(
         &mut self,
         node: ArgWithDefault<U>,
@@ -2944,6 +2950,25 @@ pub fn fold_type_param_type_var_tuple<U, F: Fold<U> + ?Sized>(
     let range = folder.map_user(range, context)?;
     Ok(TypeParamTypeVarTuple { name, range })
 }
+impl<T, U> Foldable<T, U> for Decorator<T> {
+    type Mapped = Decorator<U>;
+    fn fold<F: Fold<T, TargetU = U> + ?Sized>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self::Mapped, F::Error> {
+        folder.fold_decorator(self)
+    }
+}
+pub fn fold_decorator<U, F: Fold<U> + ?Sized>(
+    #[allow(unused)] folder: &mut F,
+    node: Decorator<U>,
+) -> Result<Decorator<F::TargetU>, F::Error> {
+    let Decorator { expression, range } = node;
+    let context = folder.will_map_user_cfg(&range);
+    let expression = Foldable::fold(expression, folder)?;
+    let range = folder.map_user_cfg(range, context)?;
+    Ok(Decorator { expression, range })
+}
 impl<T, U> Foldable<T, U> for ArgWithDefault<T> {
     type Mapped = ArgWithDefault<U>;
     fn fold<F: Fold<T, TargetU = U> + ?Sized>(

ast/src/gen/generic.rs

@@ -23,7 +23,9 @@ pub enum Ast<R = TextRange> {
     Pattern(Pattern<R>),
     TypeIgnore(TypeIgnore<R>),
     TypeParam(TypeParam<R>),
+    Decorator(Decorator<R>),
 }
+
 impl<R> Node for Ast<R> {
     const NAME: &'static str = "AST";
     const FIELD_NAMES: &'static [&'static str] = &[];
@@ -143,6 +145,12 @@ impl<R> From<TypeParam<R>> for Ast<R> {
     }
 }
 
+impl<R> From<Decorator<R>> for Ast<R> {
+    fn from(node: Decorator<R>) -> Self {
+        Ast::Decorator(node)
+    }
+}
+
 /// See also [mod](https://docs.python.org/3/library/ast.html#ast.mod)
 #[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Mod<R = TextRange> {
@@ -315,7 +323,7 @@ pub struct StmtFunctionDef<R = TextRange> {
     pub name: Identifier,
     pub args: Box<Arguments<R>>,
     pub body: Vec<Stmt<R>>,
-    pub decorator_list: Vec<Expr<R>>,
+    pub decorator_list: Vec<Decorator<R>>,
     pub returns: Option<Box<Expr<R>>>,
     pub type_comment: Option<String>,
     pub type_params: Vec<TypeParam<R>>,
@@ -351,7 +359,7 @@ pub struct StmtAsyncFunctionDef<R = TextRange> {
     pub name: Identifier,
     pub args: Box<Arguments<R>>,
     pub body: Vec<Stmt<R>>,
-    pub decorator_list: Vec<Expr<R>>,
+    pub decorator_list: Vec<Decorator<R>>,
     pub returns: Option<Box<Expr<R>>>,
     pub type_comment: Option<String>,
     pub type_params: Vec<TypeParam<R>>,
@@ -388,7 +396,7 @@ pub struct StmtClassDef<R = TextRange> {
     pub bases: Vec<Expr<R>>,
     pub keywords: Vec<Keyword<R>>,
     pub body: Vec<Stmt<R>>,
-    pub decorator_list: Vec<Expr<R>>,
+    pub decorator_list: Vec<Decorator<R>>,
     pub type_params: Vec<TypeParam<R>>,
 }
 
@@ -3197,6 +3205,18 @@ impl<R> Node for TypeParam<R> {
     const FIELD_NAMES: &'static [&'static str] = &[];
 }
 
+/// See also [decorator](https://docs.python.org/3/library/ast.html#ast.decorator)
+#[derive(Clone, Debug, PartialEq)]
+pub struct Decorator<R = TextRange> {
+    pub range: OptionalRange<R>,
+    pub expression: Expr<R>,
+}
+
+impl<R> Node for Decorator<R> {
+    const NAME: &'static str = "decorator";
+    const FIELD_NAMES: &'static [&'static str] = &["expression"];
+}
+
 /// An alternative type of AST `arguments`. This is parser-friendly and human-friendly definition of function arguments.
 /// This form also has advantage to implement pre-order traverse.
 /// `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

@@ -1383,6 +1383,23 @@ impl LocatedMut for TypeIgnore {
     }
 }
 
+
+pub type Decorator = crate::generic::Decorator<SourceRange>;
+
+#[cfg(feature = "all-nodes-with-ranges")]
+impl Located for Decorator {
+    fn range(&self) -> SourceRange {
+        self.range
+    }
+}
+
+#[cfg(feature = "all-nodes-with-ranges")]
+impl LocatedMut for Decorator {
+    fn range_mut(&mut self) -> &mut SourceRange {
+        &mut self.range
+    }
+}
+
 pub type TypeParam = crate::generic::TypeParam<SourceRange>;
 
 pub type TypeParamTypeVar = crate::generic::TypeParamTypeVar<SourceRange>;

ast/src/gen/ranged.rs

@@ -527,6 +527,12 @@ impl Ranged for crate::TypeParam {
     }
 }
 
+#[cfg(feature = "all-nodes-with-ranges")]
+impl Ranged for crate::generic::Decorator<TextRange> {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}
 #[cfg(feature = "all-nodes-with-ranges")]
 impl Ranged for crate::generic::Arguments<TextRange> {
     fn range(&self) -> TextRange {

ast/src/gen/visitor.rs

@@ -49,7 +49,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
             self.visit_stmt(value);
         }
         for value in node.decorator_list {
-            self.visit_expr(value);
+            self.visit_decorator(value);
         }
         if let Some(value) = node.returns {
             self.visit_expr(*value);
@@ -70,7 +70,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
             self.visit_stmt(value);
         }
         for value in node.decorator_list {
-            self.visit_expr(value);
+            self.visit_decorator(value);
         }
         if let Some(value) = node.returns {
             self.visit_expr(*value);
@@ -93,7 +93,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {
             self.visit_stmt(value);
         }
         for value in node.decorator_list {
-            self.visit_expr(value);
+            self.visit_decorator(value);
         }
         for value in node.type_params {
             self.visit_type_param(value);
@@ -862,4 +862,8 @@ pub trait Visitor<R = crate::text_size::TextRange> {
         self.generic_visit_type_param_type_var_tuple(node)
     }
     fn generic_visit_type_param_type_var_tuple(&mut self, node: TypeParamTypeVarTuple<R>) {}
+    fn visit_decorator(&mut self, node: Decorator<R>) {
+        self.generic_visit_decorator(node)
+    }
+    fn generic_visit_decorator(&mut self, node: Decorator<R>) {}
 }

parser/src/parser.rs

@@ -1239,7 +1239,7 @@ def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...
     }
 
     #[test]
-    #[cfg(not(feature = "all-nodes-with-ranges"))]
+    #[cfg(feature = "all-nodes-with-ranges")]
     fn decorator_ranges() {
         let parse_ast = parse_program(
             r#"