Add Decorator node

Author: MichaReiser

ast-pyo3/src/gen/to_py_ast.rs

@@ -944,6 +944,14 @@ impl<R> PyNode for ast::TypeIgnoreTypeIgnore<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> {
@@ -2605,6 +2613,22 @@ impl ToPyAst for ast::TypeIgnoreTypeIgnore<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> {
@@ -4717,6 +4741,22 @@ impl ToPyAst for ast::TypeIgnoreTypeIgnore<SourceRange> {
     }
 }
 
+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)
+    }
+}
+
 fn init_types(py: Python) -> PyResult<()> {
     let ast_module = PyModule::import(py, "_ast")?;
     cache_py_type::<ast::Mod>(ast_module)?;
@@ -4837,5 +4877,6 @@ fn init_types(py: Python) -> PyResult<()> {
     cache_py_type::<ast::PatternMatchOr>(ast_module)?;
     cache_py_type::<ast::TypeIgnore>(ast_module)?;
     cache_py_type::<ast::TypeIgnoreTypeIgnore>(ast_module)?;
+    cache_py_type::<ast::Decorator>(ast_module)?;
     Ok(())
 }

ast-pyo3/src/gen/wrapper_located.rs

@@ -3993,6 +3993,39 @@ impl TypeIgnoreTypeIgnore {
     }
 }
 
+#[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))
+    }
+}
+
+#[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 {
     #[inline]
     fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {
@@ -4409,5 +4442,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
     super::init_type::<PatternMatchOr, ast::PatternMatchOr>(py, m)?;
     super::init_type::<TypeIgnore, ast::TypeIgnore>(py, m)?;
     super::init_type::<TypeIgnoreTypeIgnore, ast::TypeIgnoreTypeIgnore>(py, m)?;
+    super::init_type::<Decorator, ast::Decorator>(py, m)?;
     Ok(())
 }

ast-pyo3/src/gen/wrapper_ranged.rs

@@ -3993,6 +3993,39 @@ 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)
+    }
+}
+
+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 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<()> {
     super::init_module(py, m)?;
     super::init_type::<Mod, ast::Mod>(py, m)?;
@@ -4113,5 +4146,6 @@ pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
     super::init_type::<PatternMatchOr, ast::PatternMatchOr>(py, m)?;
     super::init_type::<TypeIgnore, ast::TypeIgnore>(py, m)?;
     super::init_type::<TypeIgnoreTypeIgnore, ast::TypeIgnoreTypeIgnore>(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)
           | AsyncFunctionDef(identifier name, arguments args,
-                             stmt* body, expr* decorator_list, expr? returns,
+                             stmt* body, decorator* decorator_list, expr? returns,
                              string? type_comment)
 
           | ClassDef(identifier name,
              expr* bases,
              keyword* keywords,
              stmt* body,
-             expr* decorator_list)
+             decorator* decorator_list)
           | Return(expr? value)
 
           | Delete(expr* targets)
@@ -142,4 +142,6 @@ module Python
              attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
 
     type_ignore = TypeIgnore(int lineno, string tag)
+
+    decorator = (expr expression)
 }

ast/src/gen/fold.rs

@@ -507,6 +507,12 @@ pub trait Fold<U> {
     ) -> Result<TypeIgnoreTypeIgnore<Self::TargetU>, Self::Error> {
         fold_type_ignore_type_ignore(self, node)
     }
+    fn fold_decorator(
+        &mut self,
+        node: Decorator<U>,
+    ) -> Result<Decorator<Self::TargetU>, Self::Error> {
+        fold_decorator(self, node)
+    }
 }
 impl<T, U> Foldable<T, U> for Mod<T> {
     type Mapped = Mod<U>;
@@ -2791,3 +2797,22 @@ pub fn fold_type_ignore_type_ignore<U, F: Fold<U> + ?Sized>(
     let range = folder.map_user_cfg(range, context)?;
     Ok(TypeIgnoreTypeIgnore { lineno, tag, 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 })
+}

ast/src/gen/generic.rs

@@ -23,6 +23,7 @@ pub enum Ast<R = TextRange> {
     MatchCase(MatchCase<R>),
     Pattern(Pattern<R>),
     TypeIgnore(TypeIgnore<R>),
+    Decorator(Decorator<R>),
 }
 impl<R> Node for Ast<R> {
     const NAME: &'static str = "AST";
@@ -137,6 +138,12 @@ impl<R> From<TypeIgnore<R>> for Ast<R> {
     }
 }
 
+impl<R> From<Decorator<R>> for Ast<R> {
+    fn from(node: Decorator<R>) -> Self {
+        Ast::Decorator(node)
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct ModModule<R = TextRange> {
     pub range: OptionalRange<R>,
@@ -242,7 +249,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>,
 }
@@ -275,7 +282,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>,
 }
@@ -309,7 +316,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>>,
 }
 
 impl<R> Node for StmtClassDef<R> {
@@ -2987,3 +2994,14 @@ impl<R> Node for TypeIgnore<R> {
     const NAME: &'static str = "type_ignore";
     const FIELD_NAMES: &'static [&'static str] = &[];
 }
+
+#[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"];
+}

ast/src/gen/located.rs

@@ -815,3 +815,13 @@ impl Located 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
+    }
+}

ast/src/gen/ranged.rs

@@ -495,3 +495,10 @@ impl Ranged for crate::TypeIgnore {
         }
     }
 }
+
+#[cfg(feature = "all-nodes-with-ranges")]
+impl Ranged for crate::generic::Decorator<TextRange> {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}

ast/src/gen/visitor.rs

@@ -48,7 +48,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);
@@ -66,7 +66,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);
@@ -86,7 +86,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);
         }
     }
     fn visit_stmt_Return(&mut self, node: StmtReturn<R>) {
@@ -810,4 +810,8 @@ pub trait Visitor<R = crate::text_size::TextRange> {
             self.visit_pattern(value);
         }
     }
+    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

@@ -1035,7 +1035,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#"