Add Decorator node (#7)

Author: zanieb

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> {
@@ -4954,7 +4979,7 @@ impl ToPyAst for ast::TypeParamTypeVarTuple<SourceRange> {
 
         let Self {
             name,
-            range: _range,
+            range: _range
         } = self;
 
         let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((name.to_py_ast(py)?,))?;
@@ -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

@@ -4082,6 +4082,22 @@ impl ToPyObject for TypeParam {
         Py::new(py, initializer).unwrap().into_py(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::TypeParam<SourceRange> {
     #[inline]
@@ -4119,6 +4135,12 @@ impl ToPyWrapper for ast::TypeParamTypeVar<SourceRange> {
         Ok(TypeParamTypeVar(self).to_object(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))
+    }
+}
 
 #[pymethods]
 impl TypeParamTypeVar {
@@ -4203,6 +4225,21 @@ impl TypeParamTypeVarTuple {
     fn get_name(&self, py: Python) -> PyResult<PyObject> {
         self.0.name.to_py_wrapper(py)
     }
+
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
+}
+
+#[pymethods]
+impl Decorator {
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
 }
 
 impl ToPyWrapper for ast::ExprContext {
@@ -4626,5 +4663,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

@@ -4078,7 +4078,22 @@ impl TypeParam {
 }
 impl ToPyObject for TypeParam {
     fn to_object(&self, py: Python) -> PyObject {
-        let initializer = Self::new();
+        let initializer = Self::new();Py::new(py, initializer).unwrap().into_py(py)
+    }
+}
+#[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)
     }
 }
@@ -4120,6 +4135,13 @@ impl ToPyWrapper for ast::TypeParamTypeVar<TextRange> {
     }
 }
 
+impl ToPyWrapper for ast::Decorator<TextRange> {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(Decorator(self).to_object(py))
+    }
+}
+
 #[pymethods]
 impl TypeParamTypeVar {
     #[getter]
@@ -4203,6 +4225,21 @@ impl TypeParamTypeVarTuple {
     fn get_name(&self, py: Python) -> PyResult<PyObject> {
         self.0.name.to_py_wrapper(py)
     }
+
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
+}
+
+#[pymethods]
+impl Decorator {
+    #[getter]
+    #[inline]
+    fn get_expression(&self, py: Python) -> PyResult<PyObject> {
+        self.0.expression.to_py_wrapper(py)
+    }
 }
 
 pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
@@ -4330,5 +4367,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,18 +9,18 @@ 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,
              type_param* type_params)
+             decorator* decorator_list)
           | Return(expr? value)
 
           | Delete(expr* targets)
@@ -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>,
@@ -2864,14 +2870,14 @@ pub fn fold_type_ignore_type_ignore<U, F: Fold<U> + ?Sized>(
     Ok(TypeIgnoreTypeIgnore { lineno, tag, range })
 }
 impl<T, U> Foldable<T, U> for TypeParam<T> {
-    type Mapped = TypeParam<U>;
-    fn fold<F: Fold<T, TargetU = U> + ?Sized>(
+    type Mapped = TypeParam<U>;fn fold<F: Fold<T, TargetU = U> + ?Sized>(
         self,
         folder: &mut F,
     ) -> Result<Self::Mapped, F::Error> {
         folder.fold_type_param(self)
     }
 }
+
 pub fn fold_type_param<U, F: Fold<U> + ?Sized>(
     #[allow(unused)] folder: &mut F,
     node: TypeParam<U>,
@@ -2944,6 +2950,27 @@ 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>(
@@ -2953,6 +2980,7 @@ impl<T, U> Foldable<T, U> for ArgWithDefault<T> {
         folder.fold_arg_with_default(self)
     }
 }
+
 pub fn fold_arg_with_default<U, F: Fold<U> + ?Sized>(
     #[allow(unused)] folder: &mut F,
     node: ArgWithDefault<U>,

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,8 +396,8 @@ 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 type_params: Vec<TypeParam<R>>,
+    pub decorator_list: Vec<Decorator<R>>,
 }
 
 impl<R> Node for StmtClassDef<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.