feat: impl string-like schema attribute parser and ast printer (kcl-l…

…ang#849) Signed-off-by: peefy <[email protected]>
Peefy · Nov 7, 2023 · b24a298 · b24a298
1 parent 7f146e2
commit b24a298
Show file tree

Hide file tree

Showing 15 changed files with 482 additions and 13 deletions.
diff --git a/kclvm/ast/src/token.rs b/kclvm/ast/src/token.rs
@@ -334,6 +334,20 @@ impl Token {
         self.run_on_ident(|id| id.name == kw)
     }
 
+    /// Whether the token is a string literal token.
+    pub fn is_string_lit(&self) -> bool {
+        match self.kind {
+            TokenKind::Literal(lit) => {
+                if let LitKind::Str { .. } = lit.kind {
+                    true
+                } else {
+                    false
+                }
+            }
+            _ => false,
+        }
+    }
+
     fn run_on_ident(&self, pred: impl FnOnce(Ident) -> bool) -> bool {
         match self.ident() {
             Some(id) => pred(id),

diff --git a/kclvm/ast_pretty/src/node.rs b/kclvm/ast_pretty/src/node.rs
@@ -345,7 +345,7 @@ impl<'p, 'ctx> MutSelfTypedResultWalker<'ctx> for Printer<'p> {
         if !schema_attr.decorators.is_empty() {
             self.write_newline();
         }
-        self.write(&schema_attr.name.node);
+        self.write_attribute(&schema_attr.name);
         if schema_attr.is_optional {
             self.write("?");
         }
@@ -878,6 +878,16 @@ impl<'p> Printer<'p> {
             }
         }
     }
+
+    fn write_attribute(&mut self, attr: &ast::NodeRef<String>) {
+        let re = fancy_regex::Regex::new(IDENTIFIER_REGEX).unwrap();
+        let need_quote = !re.is_match(&attr.node).unwrap();
+        if need_quote {
+            self.write(&format!("{:?}", attr.node));
+        } else {
+            self.write(&attr.node);
+        };
+    }
 }
 
 impl<'p> Printer<'p> {

diff --git a/kclvm/ast_pretty/src/test_data/codelayout.input b/kclvm/ast_pretty/src/test_data/codelayout.input
@@ -7,6 +7,8 @@ import           math   as  alias_math
 schema Person (  Base):
   name:str
   age:int
+  "attr": str
+  "attr-x":  str
   check   :
         age>0 if    age ,   "age must > 0"
 person = Person{

diff --git a/kclvm/ast_pretty/src/test_data/codelayout.output b/kclvm/ast_pretty/src/test_data/codelayout.output
@@ -6,6 +6,8 @@ import math as alias_math
 schema Person(Base):
     name: str
     age: int
+    attr: str
+    "attr-x": str
 
     check:
         age > 0 if age, "age must > 0"

diff --git a/kclvm/parser/src/parser/stmt.rs b/kclvm/parser/src/parser/stmt.rs
@@ -909,8 +909,16 @@ impl<'a> Parser<'a> {
     fn parse_schema_body(&mut self) -> SchemaStmt {
         self.bump_token(TokenKind::Indent);
 
-        // doc string
-        let body_doc = self.parse_doc();
+        // doc string when it is not a string-like attribute statement.
+        let body_doc = if let Some(peek) = self.cursor.peek() {
+            if matches!(peek.kind, TokenKind::Colon) {
+                None
+            } else {
+                self.parse_doc()
+            }
+        } else {
+            self.parse_doc()
+        };
 
         // mixin
         let body_mixins = if self.token.is_keyword(kw::Mixin) {
@@ -942,6 +950,20 @@ impl<'a> Parser<'a> {
                 body_body.push(self.parse_if_stmt());
                 continue;
             }
+            // schema_attribute_stmt: string COLON type_annotation
+            else if self.token.is_string_lit() {
+                if let Some(peek) = self.cursor.peek() {
+                    if let TokenKind::Colon = peek.kind {
+                        let token = self.token;
+                        let attr = self.parse_schema_attribute();
+                        body_body.push(node_ref!(
+                            Stmt::SchemaAttr(attr),
+                            self.token_span_pos(token, self.prev_token)
+                        ));
+                        continue;
+                    }
+                }
+            }
             // schema_attribute_stmt
             else if let TokenKind::At = self.token.kind {
                 let token = self.token;
@@ -1133,7 +1155,7 @@ impl<'a> Parser<'a> {
 
     /// Syntax:
     /// schema_attribute_stmt: attribute_stmt NEWLINE
-    /// attribute_stmt: [decorators] identifier [QUESTION] COLON type [(ASSIGN|COMP_OR) test]
+    /// attribute_stmt: [decorators] (identifier | string) [QUESTION] COLON type [(ASSIGN|COMP_OR) test]
     fn parse_schema_attribute(&mut self) -> SchemaAttr {
         let doc = "".to_string();
 
@@ -1146,21 +1168,28 @@ impl<'a> Parser<'a> {
             Vec::new()
         };
 
-        let name_expr = self.parse_identifier();
-
-        let name_pos = name_expr.pos();
-        let name = name_expr.node;
-        let name = node_ref!(name.get_names().join("."), name_pos);
+        // Parse schema identifier-like or string-like attributes
+        let name = if let Some(name) = self.parse_string_attribute() {
+            name
+        } else {
+            let name_expr = self.parse_identifier();
+            let name_pos = name_expr.pos();
+            let name = name_expr.node;
+            node_ref!(name.get_names().join("."), name_pos)
+        };
 
+        // Parse attribute optional annotation `?`
         let is_optional = if let TokenKind::Question = self.token.kind {
             self.bump_token(TokenKind::Question);
             true
         } else {
             false
         };
 
+        // Bump the schema attribute annotation token `:`
         self.bump_token(TokenKind::Colon);
 
+        // Parse the schema attribute type annotation.
         let typ = self.parse_type_annotation();
         let type_str = node_ref!(typ.node.to_string(), typ.pos());
 
@@ -1479,6 +1508,23 @@ impl<'a> Parser<'a> {
         )
     }
 
+    pub(crate) fn parse_string_attribute(&mut self) -> Option<NodeRef<String>> {
+        match self.token.kind {
+            TokenKind::Literal(lit) => {
+                if let LitKind::Str { .. } = lit.kind {
+                    let str_expr = self.parse_str_expr(lit);
+                    match &str_expr.node {
+                        Expr::StringLit(str) => Some(node_ref!(str.value.clone(), str_expr.pos())),
+                        _ => None,
+                    }
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    }
+
     pub(crate) fn parse_joined_string(
         &mut self,
         s: &StringLit,

diff --git a/kclvm/parser/src/tests/error_recovery.rs b/kclvm/parser/src/tests/error_recovery.rs
@@ -274,6 +274,19 @@ parse_module_snapshot! { schema_stmt_recovery_31, r#"
 schema A:
     [str]: str
     [str]: int"#}
+parse_module_snapshot! { schema_stmt_recovery_32, r#"
+schema A:
+    "attr": str"#}
+parse_module_snapshot! { schema_stmt_recovery_33, r#"
+schema A:
+    """Schema Doc"""
+    "attr": str"#}
+parse_module_snapshot! { schema_stmt_recovery_34, r#"
+schema A:
+    "attr: str"#}
+parse_module_snapshot! { schema_stmt_recovery_35, r#"
+schema A:
+    "attr":"#}
 parse_module_snapshot! { rule_stmt_recovery_0, r#"rule"#}
 parse_module_snapshot! { rule_stmt_recovery_1, r#"rule A"#}
 parse_module_snapshot! { rule_stmt_recovery_2, r#"rule A["#}

diff --git a/...ser/src/tests/snapshots/kclvm_parser__tests__error_recovery__schema_stmt_recovery_32.snap b/...ser/src/tests/snapshots/kclvm_parser__tests__error_recovery__schema_stmt_recovery_32.snap
@@ -0,0 +1,88 @@
+---
+source: parser/src/tests/error_recovery.rs
+assertion_line: 277
+expression: "crate::tests::parsing_module_string(r#\"\nschema A:\n    \"attr\": str\"#)"
+---
+Module {
+    filename: "",
+    pkg: "",
+    doc: None,
+    name: "",
+    body: [
+        Node {
+            node: Schema(
+                SchemaStmt {
+                    doc: None,
+                    name: Node {
+                        node: "A",
+                        filename: "",
+                        line: 2,
+                        column: 7,
+                        end_line: 2,
+                        end_column: 8,
+                    },
+                    parent_name: None,
+                    for_host_name: None,
+                    is_mixin: false,
+                    is_protocol: false,
+                    args: None,
+                    mixins: [],
+                    body: [
+                        Node {
+                            node: SchemaAttr(
+                                SchemaAttr {
+                                    doc: "",
+                                    name: Node {
+                                        node: "attr",
+                                        filename: "",
+                                        line: 3,
+                                        column: 4,
+                                        end_line: 3,
+                                        end_column: 10,
+                                    },
+                                    type_str: Node {
+                                        node: "str",
+                                        filename: "",
+                                        line: 3,
+                                        column: 12,
+                                        end_line: 3,
+                                        end_column: 15,
+                                    },
+                                    op: None,
+                                    value: None,
+                                    is_optional: false,
+                                    decorators: [],
+                                    ty: Node {
+                                        node: Basic(
+                                            Str,
+                                        ),
+                                        filename: "",
+                                        line: 3,
+                                        column: 12,
+                                        end_line: 3,
+                                        end_column: 15,
+                                    },
+                                },
+                            ),
+                            filename: "",
+                            line: 3,
+                            column: 4,
+                            end_line: 3,
+                            end_column: 15,
+                        },
+                    ],
+                    decorators: [],
+                    checks: [],
+                    index_signature: None,
+                },
+            ),
+            filename: "",
+            line: 2,
+            column: 0,
+            end_line: 3,
+            end_column: 15,
+        },
+    ],
+    comments: [],
+}
+
diff --git a/...ser/src/tests/snapshots/kclvm_parser__tests__error_recovery__schema_stmt_recovery_33.snap b/...ser/src/tests/snapshots/kclvm_parser__tests__error_recovery__schema_stmt_recovery_33.snap
@@ -0,0 +1,97 @@
+---
+source: parser/src/tests/error_recovery.rs
+assertion_line: 280
+expression: "crate::tests::parsing_module_string(r#\"\nschema A:\n    \"\"\"Schema Doc\"\"\"\n    \"attr\": str\"#)"
+---
+Module {
+    filename: "",
+    pkg: "",
+    doc: None,
+    name: "",
+    body: [
+        Node {
+            node: Schema(
+                SchemaStmt {
+                    doc: Some(
+                        Node {
+                            node: "\"\"\"Schema Doc\"\"\"",
+                            filename: "",
+                            line: 3,
+                            column: 4,
+                            end_line: 3,
+                            end_column: 20,
+                        },
+                    ),
+                    name: Node {
+                        node: "A",
+                        filename: "",
+                        line: 2,
+                        column: 7,
+                        end_line: 2,
+                        end_column: 8,
+                    },
+                    parent_name: None,
+                    for_host_name: None,
+                    is_mixin: false,
+                    is_protocol: false,
+                    args: None,
+                    mixins: [],
+                    body: [
+                        Node {
+                            node: SchemaAttr(
+                                SchemaAttr {
+                                    doc: "",
+                                    name: Node {
+                                        node: "attr",
+                                        filename: "",
+                                        line: 4,
+                                        column: 4,
+                                        end_line: 4,
+                                        end_column: 10,
+                                    },
+                                    type_str: Node {
+                                        node: "str",
+                                        filename: "",
+                                        line: 4,
+                                        column: 12,
+                                        end_line: 4,
+                                        end_column: 15,
+                                    },
+                                    op: None,
+                                    value: None,
+                                    is_optional: false,
+                                    decorators: [],
+                                    ty: Node {
+                                        node: Basic(
+                                            Str,
+                                        ),
+                                        filename: "",
+                                        line: 4,
+                                        column: 12,
+                                        end_line: 4,
+                                        end_column: 15,
+                                    },
+                                },
+                            ),
+                            filename: "",
+                            line: 4,
+                            column: 4,
+                            end_line: 4,
+                            end_column: 15,
+                        },
+                    ],
+                    decorators: [],
+                    checks: [],
+                    index_signature: None,
+                },
+            ),
+            filename: "",
+            line: 2,
+            column: 0,
+            end_line: 4,
+            end_column: 15,
+        },
+    ],
+    comments: [],
+}
+