Merge pull request #16 from youknowone/ast-int

Give identifier and int ast types

Author: youknowone

ast/asdl_rs.py

@@ -17,21 +17,26 @@
 AUTO_GEN_MESSAGE = "// File automatically generated by {}.\n\n"
 
 builtin_type_mapping = {
-    "identifier": "Ident",
+    "identifier": "Identifier",
     "string": "String",
-    "int": "u32",
+    "int": "Int",
     "constant": "Constant",
 }
 assert builtin_type_mapping.keys() == asdl.builtin_types
 
+builtin_int_mapping = {
+    "simple": "bool",
+    "is_async": "bool",
+}
 
 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:
-        return builtin_type_mapping[name]
+        builtin =  builtin_type_mapping[name]
+        return builtin
     elif name.islower():
         return "".join(part.capitalize() for part in name.split("_"))
     else:
@@ -355,6 +360,8 @@ def visitField(self, field, parent, vis, depth, constructor=None):
             typ = f"Option<{typ}>"
         if field.seq:
             typ = f"Vec<{typ}>"
+        if typ == "Int":
+            typ = builtin_int_mapping.get(field.name, typ)
         name = rust_field(field.name)
         self.emit(f"{vis}{name}: {typ},", depth)
 

ast/src/attributed.rs

@@ -0,0 +1,186 @@
+//! `builtin_types` in asdl.py and Attributed
+
+use num_bigint::BigInt;
+use rustpython_parser_core::text_size::{TextRange, TextSize};
+
+pub type String = std::string::String;
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct Identifier(String);
+
+impl Identifier {
+    #[inline]
+    pub fn new(s: impl Into<String>) -> Self {
+        Self(s.into())
+    }
+    #[inline]
+    pub fn as_str(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
+impl std::string::ToString for Identifier {
+    #[inline]
+    fn to_string(&self) -> String {
+        self.0.clone()
+    }
+}
+
+impl From<Identifier> for String {
+    #[inline]
+    fn from(id: Identifier) -> String {
+        id.0
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Int(u32);
+
+impl Int {
+    pub fn new(i: u32) -> Self {
+        Self(i)
+    }
+    pub fn new_bool(i: bool) -> Self {
+        Self(i as u32)
+    }
+    pub fn to_u32(&self) -> u32 {
+        self.0
+    }
+    pub fn to_bool(&self) -> bool {
+        self.0 > 0
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum Constant {
+    None,
+    Bool(bool),
+    Str(String),
+    Bytes(Vec<u8>),
+    Int(BigInt),
+    Tuple(Vec<Constant>),
+    Float(f64),
+    Complex { real: f64, imag: f64 },
+    Ellipsis,
+}
+
+impl From<String> for Constant {
+    fn from(s: String) -> Constant {
+        Self::Str(s)
+    }
+}
+impl From<Vec<u8>> for Constant {
+    fn from(b: Vec<u8>) -> Constant {
+        Self::Bytes(b)
+    }
+}
+impl From<bool> for Constant {
+    fn from(b: bool) -> Constant {
+        Self::Bool(b)
+    }
+}
+impl From<BigInt> for Constant {
+    fn from(i: BigInt) -> Constant {
+        Self::Int(i)
+    }
+}
+
+#[cfg(feature = "rustpython-literal")]
+impl std::fmt::Display for Constant {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Constant::None => f.pad("None"),
+            Constant::Bool(b) => f.pad(if *b { "True" } else { "False" }),
+            Constant::Str(s) => rustpython_literal::escape::UnicodeEscape::new_repr(s.as_str())
+                .str_repr()
+                .write(f),
+            Constant::Bytes(b) => {
+                let escape = rustpython_literal::escape::AsciiEscape::new_repr(b);
+                let repr = escape.bytes_repr().to_string().unwrap();
+                f.pad(&repr)
+            }
+            Constant::Int(i) => i.fmt(f),
+            Constant::Tuple(tup) => {
+                if let [elt] = &**tup {
+                    write!(f, "({elt},)")
+                } else {
+                    f.write_str("(")?;
+                    for (i, elt) in tup.iter().enumerate() {
+                        if i != 0 {
+                            f.write_str(", ")?;
+                        }
+                        elt.fmt(f)?;
+                    }
+                    f.write_str(")")
+                }
+            }
+            Constant::Float(fp) => f.pad(&rustpython_literal::float::to_string(*fp)),
+            Constant::Complex { real, imag } => {
+                if *real == 0.0 {
+                    write!(f, "{imag}j")
+                } else {
+                    write!(f, "({real}{imag:+}j)")
+                }
+            }
+            Constant::Ellipsis => f.pad("..."),
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct Attributed<T, U = ()> {
+    pub range: TextRange,
+    pub custom: U,
+    pub node: T,
+}
+
+impl<T, U> Attributed<T, U> {
+    /// Returns the node
+    #[inline]
+    pub fn node(&self) -> &T {
+        &self.node
+    }
+
+    /// Returns the `range` of the node. The range offsets are absolute to the start of the document.
+    #[inline]
+    pub const fn range(&self) -> TextRange {
+        self.range
+    }
+
+    /// Returns the absolute start position of the node from the beginning of the document.
+    #[inline]
+    pub const fn start(&self) -> TextSize {
+        self.range.start()
+    }
+
+    /// Returns the absolute position at which the node ends in the source document.
+    #[inline]
+    pub const fn end(&self) -> TextSize {
+        self.range.end()
+    }
+}
+
+impl<T> Attributed<T, ()> {
+    /// Creates a new node that spans the position specified by `range`.
+    pub fn new(range: impl Into<TextRange>, node: T) -> Self {
+        Self {
+            range: range.into(),
+            custom: (),
+            node,
+        }
+    }
+
+    /// Consumes self and returns the node.
+    #[inline]
+    pub fn into_node(self) -> T {
+        self.node
+    }
+}
+
+impl<T, U> std::ops::Deref for Attributed<T, U> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.node
+    }
+}

ast/src/builtin.rs

@@ -1,4 +1,4 @@
-use crate::{constant, fold::Fold};
+use crate::{builtin, fold::Fold};
 
 pub trait Foldable<T, U> {
     type Mapped;
@@ -62,4 +62,10 @@ macro_rules! simple_fold {
     };
 }
 
-simple_fold!(u32, String, bool, constant::Constant);
+simple_fold!(
+    builtin::Int,
+    builtin::String,
+    builtin::Identifier,
+    bool,
+    builtin::Constant
+);