Add Decorator node (#7)

MichaReiser · charliermarsh · commit 6f65c5cba7f0 · 2023-06-19T12:53:05.000-04:00
diff --git a/ast-pyo3/src/gen/to_py_ast.rs b/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(())
 }
diff --git a/ast-pyo3/src/gen/wrapper_located.rs b/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(())
 }
diff --git a/ast-pyo3/src/gen/wrapper_ranged.rs b/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(())
 }
diff --git a/ast/Python.asdl b/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)
 }
diff --git a/ast/src/gen/fold.rs b/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)
+    }
     fn fold_arg_with_default(
         &mut self,
         node: ArgWithDefault<U>,
@@ -2791,6 +2797,25 @@ 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 })
+}
 impl<T, U> Foldable<T, U> for ArgWithDefault<T> {
     type Mapped = ArgWithDefault<U>;
     fn fold<F: Fold<T, TargetU = U> + ?Sized>(
diff --git a/ast/src/gen/generic.rs b/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)
+    }
+}
+
 /// 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> {
@@ -307,7 +314,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>,
 }
@@ -341,7 +348,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>,
 }
@@ -376,7 +383,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> {
@@ -3074,6 +3081,18 @@ impl<R> Node for TypeIgnore<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.
diff --git a/ast/src/gen/located.rs b/ast/src/gen/located.rs
@@ -1367,6 +1367,21 @@ 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 Arguments = crate::generic::Arguments<SourceRange>;
 
 #[cfg(feature = "all-nodes-with-ranges")]
diff --git a/ast/src/gen/ranged.rs b/ast/src/gen/ranged.rs
@@ -496,6 +496,12 @@ impl Ranged for crate::TypeIgnore {
     }
 }
 
+#[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 {
diff --git a/ast/src/gen/visitor.rs b/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>) {}
 }
diff --git a/parser/src/parser.rs b/parser/src/parser.rs
@@ -1102,7 +1102,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#"
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -1151,9 +1151,9 @@ ClassDef: ast::Stmt = {
 };
 
 // Decorators:
-Decorator: ast::Expr = {
-    <location:@L> "@" <p:NamedExpressionTest> "\n" => {
-        p
+Decorator: ast::Decorator = {
+    <location:@L> "@" <p:NamedExpressionTest> <end_location:@R> "\n" => {
+        ast::Decorator { range: optional_range(location, end_location), expression: p }
     },
 };
 
diff --git a/parser/src/python.rs b/parser/src/python.rs

Original file line number	Diff line number	Diff line change
`@@ -496,6 +496,12 @@ impl Ranged for crate::TypeIgnore {`
`496`	`496`	`}`
`497`	`497`	`}`
`498`	`498`
	`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`	`+}`
`499`	`505`	`#[cfg(feature = "all-nodes-with-ranges")]`
`500`	`506`	`impl Ranged for crate::generic::Arguments<TextRange> {`
`501`	`507`	`fn range(&self) -> TextRange {`
Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {`
`48`	`48`	`self.visit_stmt(value);`
`49`	`49`	`}`
`50`	`50`	`for value in node.decorator_list {`
`51`		`- self.visit_expr(value);`
	`51`	`+ self.visit_decorator(value);`
`52`	`52`	`}`
`53`	`53`	`if let Some(value) = node.returns {`
`54`	`54`	`self.visit_expr(*value);`
`@@ -66,7 +66,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {`
`66`	`66`	`self.visit_stmt(value);`
`67`	`67`	`}`
`68`	`68`	`for value in node.decorator_list {`
`69`		`- self.visit_expr(value);`
	`69`	`+ self.visit_decorator(value);`
`70`	`70`	`}`
`71`	`71`	`if let Some(value) = node.returns {`
`72`	`72`	`self.visit_expr(*value);`
`@@ -86,7 +86,7 @@ pub trait Visitor<R = crate::text_size::TextRange> {`
`86`	`86`	`self.visit_stmt(value);`
`87`	`87`	`}`
`88`	`88`	`for value in node.decorator_list {`
`89`		`- self.visit_expr(value);`
	`89`	`+ self.visit_decorator(value);`
`90`	`90`	`}`
`91`	`91`	`}`
`92`	`92`	`fn visit_stmt_return(&mut self, node: StmtReturn<R>) {`
`@@ -810,4 +810,8 @@ pub trait Visitor<R = crate::text_size::TextRange> {`
`810`	`810`	`self.visit_pattern(value);`
`811`	`811`	`}`
`812`	`812`	`}`
	`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>) {}`
`813`	`817`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1102,7 +1102,7 @@ def args_to_tuple(args: Ts) -> Tuple[*Ts]: ...`
`1102`	`1102`	`}`
`1103`	`1103`
`1104`	`1104`	`#[test]`
`1105`		`- #[cfg(not(feature = "all-nodes-with-ranges"))]`
	`1105`	`+ #[cfg(feature = "all-nodes-with-ranges")]`
`1106`	`1106`	`fn decorator_ranges() {`
`1107`	`1107`	`let parse_ast = parse_program(`
`1108`	`1108`	`r#"`