From 947fb53d0b41fec465db3d8e725bdb2eec1299ec Mon Sep 17 00:00:00 2001
From: "Jeong, YunWon" <69878+youknowone@users.noreply.github.com>
Date: Thu, 11 May 2023 19:44:20 +0900
Subject: [PATCH] Field accessor and utilities (#20)

* Apply is-macro to Constant and ast nodes
---
 Cargo.toml             |  1 +
 ast/Cargo.toml         |  1 +
 ast/asdl_rs.py         | 27 ++++++++++++---
 ast/src/builtin.rs     | 28 +++++++++++++++-
 ast/src/gen/generic.rs | 76 ++++++++++++++++++++++++++++++++++++------
 ast/src/impls.rs       |  2 +-
 6 files changed, 117 insertions(+), 18 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a8ce804c4b..4ed152b6e6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ anyhow = "1.0.45"
 cfg-if = "1.0"
 insta = "1.14.0"
 itertools = "0.10.3"
+is-macro = "0.2.2"
 log = "0.4.16"
 num-complex = "0.4.0"
 num-bigint = "0.4.3"
diff --git a/ast/Cargo.toml b/ast/Cargo.toml
index 78e90e1def..82827c7025 100644
--- a/ast/Cargo.toml
+++ b/ast/Cargo.toml
@@ -19,4 +19,5 @@ visitor = []
 rustpython-parser-core = { workspace = true }
 rustpython-literal = { workspace = true, optional = true }
 
+is-macro = { workspace = true }
 num-bigint = { workspace = true }
diff --git a/ast/asdl_rs.py b/ast/asdl_rs.py
index a6f4f86301..822badc699 100755
--- a/ast/asdl_rs.py
+++ b/ast/asdl_rs.py
@@ -6,6 +6,7 @@
 import sys
 import json
 import textwrap
+import re
 
 from argparse import ArgumentParser
 from pathlib import Path
@@ -16,26 +17,34 @@ import asdl
 TABSIZE = 4
 AUTO_GEN_MESSAGE = "// File automatically generated by {}.\n\n"
 
-builtin_type_mapping = {
+BUILTIN_TYPE_NAMES = {
     "identifier": "Identifier",
     "string": "String",
     "int": "Int",
     "constant": "Constant",
 }
-assert builtin_type_mapping.keys() == asdl.builtin_types
+assert BUILTIN_TYPE_NAMES.keys() == asdl.builtin_types
 
-builtin_int_mapping = {
+BUILTIN_INT_NAMES = {
     "simple": "bool",
     "is_async": "bool",
 }
 
+RUST_KEYWORDS = {"if", "while", "for", "return", "match", "try", "await", "yield"}
+
+
+def rust_field_name(name):
+    name = rust_type_name(name)
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+
+
 def rust_type_name(name):
     """Return a string for the C name of the type.
 
     This function special cases the default types provided by asdl.
     """
     if name in asdl.builtin_types:
-        builtin =  builtin_type_mapping[name]
+        builtin = BUILTIN_TYPE_NAMES[name]
         return builtin
     elif name.islower():
         return "".join(part.capitalize() for part in name.split("_"))
@@ -271,6 +280,7 @@ class StructVisitor(EmitVisitor):
     def simple_sum(self, sum, name, depth):
         rust_name = rust_type_name(name)
         self.emit_attrs(depth)
+        self.emit("#[derive(is_macro::Is)]", depth)
         self.emit(f"pub enum {rust_name} {{", depth)
         for variant in sum.types:
             self.emit(f"{variant.name},", depth + 1)
@@ -291,8 +301,15 @@ class StructVisitor(EmitVisitor):
 
         generics, generics_applied = self.apply_generics(name, "U = ()", "U")
         self.emit_attrs(depth)
+        self.emit("#[derive(is_macro::Is)]", depth)
         self.emit(f"pub enum {rust_name}{suffix}{generics} {{", depth)
+        needs_escape = any(rust_field_name(t.name) in RUST_KEYWORDS for t in sum.types)
         for t in sum.types:
+            if needs_escape:
+                self.emit(
+                    f'#[is(name = "{rust_field_name(t.name)}_{rust_name.lower()}")]',
+                    depth + 1,
+                )
             if t.fields:
                 (t_generics_applied,) = self.apply_generics(t.name, "U")
                 self.emit(
@@ -361,7 +378,7 @@ class StructVisitor(EmitVisitor):
         if field.seq:
             typ = f"Vec<{typ}>"
         if typ == "Int":
-            typ = builtin_int_mapping.get(field.name, typ)
+            typ = BUILTIN_INT_NAMES.get(field.name, typ)
         name = rust_field(field.name)
         self.emit(f"{vis}{name}: {typ},", depth)
 
diff --git a/ast/src/builtin.rs b/ast/src/builtin.rs
index 3f8f1b0c00..726ed4991c 100644
--- a/ast/src/builtin.rs
+++ b/ast/src/builtin.rs
@@ -114,7 +114,7 @@ impl std::cmp::PartialEq<usize> for Int {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Constant {
     None,
     Bool(bool),
@@ -127,6 +127,21 @@ pub enum Constant {
     Ellipsis,
 }
 
+impl Constant {
+    pub fn is_true(self) -> bool {
+        self.bool().map_or(false, |b| b)
+    }
+    pub fn is_false(self) -> bool {
+        self.bool().map_or(false, |b| !b)
+    }
+    pub fn complex(self) -> Option<(f64, f64)> {
+        match self {
+            Constant::Complex { real, imag } => Some((real, imag)),
+            _ => None,
+        }
+    }
+}
+
 impl From<String> for Constant {
     fn from(s: String) -> Constant {
         Self::Str(s)
@@ -247,3 +262,14 @@ impl<T, U> std::ops::Deref for Attributed<T, U> {
         &self.node
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_is_macro() {
+        let none = Constant::None;
+        assert!(none.is_none());
+        assert!(!none.is_bool());
+    }
+}
diff --git a/ast/src/gen/generic.rs b/ast/src/gen/generic.rs
index 908d317027..f79eff4bb1 100644
--- a/ast/src/gen/generic.rs
+++ b/ast/src/gen/generic.rs
@@ -46,7 +46,7 @@ impl<U> From<ModFunctionType<U>> for Mod<U> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Mod<U = ()> {
     Module(ModModule<U>),
     Interactive(ModInteractive<U>),
@@ -366,34 +366,61 @@ impl<U> From<StmtExpr<U>> for StmtKind<U> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum StmtKind<U = ()> {
+    #[is(name = "function_def_stmt")]
     FunctionDef(StmtFunctionDef<U>),
+    #[is(name = "async_function_def_stmt")]
     AsyncFunctionDef(StmtAsyncFunctionDef<U>),
+    #[is(name = "class_def_stmt")]
     ClassDef(StmtClassDef<U>),
+    #[is(name = "return_stmt")]
     Return(StmtReturn<U>),
+    #[is(name = "delete_stmt")]
     Delete(StmtDelete<U>),
+    #[is(name = "assign_stmt")]
     Assign(StmtAssign<U>),
+    #[is(name = "aug_assign_stmt")]
     AugAssign(StmtAugAssign<U>),
+    #[is(name = "ann_assign_stmt")]
     AnnAssign(StmtAnnAssign<U>),
+    #[is(name = "for_stmt")]
     For(StmtFor<U>),
+    #[is(name = "async_for_stmt")]
     AsyncFor(StmtAsyncFor<U>),
+    #[is(name = "while_stmt")]
     While(StmtWhile<U>),
+    #[is(name = "if_stmt")]
     If(StmtIf<U>),
+    #[is(name = "with_stmt")]
     With(StmtWith<U>),
+    #[is(name = "async_with_stmt")]
     AsyncWith(StmtAsyncWith<U>),
+    #[is(name = "match_stmt")]
     Match(StmtMatch<U>),
+    #[is(name = "raise_stmt")]
     Raise(StmtRaise<U>),
+    #[is(name = "try_stmt")]
     Try(StmtTry<U>),
+    #[is(name = "try_star_stmt")]
     TryStar(StmtTryStar<U>),
+    #[is(name = "assert_stmt")]
     Assert(StmtAssert<U>),
+    #[is(name = "import_stmt")]
     Import(StmtImport<U>),
+    #[is(name = "import_from_stmt")]
     ImportFrom(StmtImportFrom<U>),
+    #[is(name = "global_stmt")]
     Global(StmtGlobal),
+    #[is(name = "nonlocal_stmt")]
     Nonlocal(StmtNonlocal),
+    #[is(name = "expr_stmt")]
     Expr(StmtExpr<U>),
+    #[is(name = "pass_stmt")]
     Pass,
+    #[is(name = "break_stmt")]
     Break,
+    #[is(name = "continue_stmt")]
     Continue,
 }
 pub type Stmt<U = ()> = Attributed<StmtKind<U>, U>;
@@ -726,52 +753,79 @@ impl<U> From<ExprSlice<U>> for ExprKind<U> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum ExprKind<U = ()> {
+    #[is(name = "bool_op_expr")]
     BoolOp(ExprBoolOp<U>),
+    #[is(name = "named_expr_expr")]
     NamedExpr(ExprNamedExpr<U>),
+    #[is(name = "bin_op_expr")]
     BinOp(ExprBinOp<U>),
+    #[is(name = "unary_op_expr")]
     UnaryOp(ExprUnaryOp<U>),
+    #[is(name = "lambda_expr")]
     Lambda(ExprLambda<U>),
+    #[is(name = "if_exp_expr")]
     IfExp(ExprIfExp<U>),
+    #[is(name = "dict_expr")]
     Dict(ExprDict<U>),
+    #[is(name = "set_expr")]
     Set(ExprSet<U>),
+    #[is(name = "list_comp_expr")]
     ListComp(ExprListComp<U>),
+    #[is(name = "set_comp_expr")]
     SetComp(ExprSetComp<U>),
+    #[is(name = "dict_comp_expr")]
     DictComp(ExprDictComp<U>),
+    #[is(name = "generator_exp_expr")]
     GeneratorExp(ExprGeneratorExp<U>),
+    #[is(name = "await_expr")]
     Await(ExprAwait<U>),
+    #[is(name = "yield_expr")]
     Yield(ExprYield<U>),
+    #[is(name = "yield_from_expr")]
     YieldFrom(ExprYieldFrom<U>),
+    #[is(name = "compare_expr")]
     Compare(ExprCompare<U>),
+    #[is(name = "call_expr")]
     Call(ExprCall<U>),
+    #[is(name = "formatted_value_expr")]
     FormattedValue(ExprFormattedValue<U>),
+    #[is(name = "joined_str_expr")]
     JoinedStr(ExprJoinedStr<U>),
+    #[is(name = "constant_expr")]
     Constant(ExprConstant),
+    #[is(name = "attribute_expr")]
     Attribute(ExprAttribute<U>),
+    #[is(name = "subscript_expr")]
     Subscript(ExprSubscript<U>),
+    #[is(name = "starred_expr")]
     Starred(ExprStarred<U>),
+    #[is(name = "name_expr")]
     Name(ExprName),
+    #[is(name = "list_expr")]
     List(ExprList<U>),
+    #[is(name = "tuple_expr")]
     Tuple(ExprTuple<U>),
+    #[is(name = "slice_expr")]
     Slice(ExprSlice<U>),
 }
 pub type Expr<U = ()> = Attributed<ExprKind<U>, U>;
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum ExprContext {
     Load,
     Store,
     Del,
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Boolop {
     And,
     Or,
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Operator {
     Add,
     Sub,
@@ -788,7 +842,7 @@ pub enum Operator {
     FloorDiv,
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Unaryop {
     Invert,
     Not,
@@ -796,7 +850,7 @@ pub enum Unaryop {
     USub,
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum Cmpop {
     Eq,
     NotEq,
@@ -831,7 +885,7 @@ impl<U> From<ExcepthandlerExceptHandler<U>> for ExcepthandlerKind<U> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum ExcepthandlerKind<U = ()> {
     ExceptHandler(ExcepthandlerExceptHandler<U>),
 }
@@ -977,7 +1031,7 @@ impl<U> From<PatternMatchOr<U>> for PatternKind<U> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum PatternKind<U = ()> {
     MatchValue(PatternMatchValue<U>),
     MatchSingleton(PatternMatchSingleton),
@@ -1002,7 +1056,7 @@ impl From<TypeIgnoreTypeIgnore> for TypeIgnore {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, is_macro::Is)]
 pub enum TypeIgnore {
     TypeIgnore(TypeIgnoreTypeIgnore),
 }
diff --git a/ast/src/impls.rs b/ast/src/impls.rs
index e3e827a715..db85612626 100644
--- a/ast/src/impls.rs
+++ b/ast/src/impls.rs
@@ -2,7 +2,7 @@ use crate::{Constant, ExprKind};
 
 impl<U> ExprKind<U> {
     /// Returns a short name for the node suitable for use in error messages.
-    pub fn name(&self) -> &'static str {
+    pub fn python_name(&self) -> &'static str {
         match self {
             ExprKind::BoolOp { .. } | ExprKind::BinOp { .. } | ExprKind::UnaryOp { .. } => {
                 "operator"