From e24e06bb1371ac583851f7b3dd9c87471e3ab607 Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Sun, 28 Feb 2021 19:07:19 -0500
Subject: [PATCH 1/2] Added ast_utils + example

---
 examples/advanced/create_ast.py | 105 ++++++++++++++++++++++++++++++++
 lark/ast_utils.py               |  51 ++++++++++++++++
 2 files changed, 156 insertions(+)
 create mode 100644 examples/advanced/create_ast.py
 create mode 100644 lark/ast_utils.py

diff --git a/examples/advanced/create_ast.py b/examples/advanced/create_ast.py
new file mode 100644
index 0000000..e20b7e6
--- /dev/null
+++ b/examples/advanced/create_ast.py
@@ -0,0 +1,105 @@
+"""
+    This example demonstrates how to transform a parse-tree into an AST using `lark.ast_utils`.
+
+    This example only works with Python 3.
+"""
+
+import sys
+from typing import List
+from dataclasses import dataclass
+
+from lark import Lark, ast_utils, Transformer, v_args
+
+this_module = sys.modules[__name__]
+
+
+#
+#   Define AST
+#
+class _Ast(ast_utils.Ast):
+    pass
+
+class _Statement(_Ast):
+    pass
+
+@dataclass
+class Value(_Ast):
+    value: object
+
+@dataclass
+class Name(_Ast):
+    name: str
+
+@dataclass
+class CodeBlock(_Ast, ast_utils.AsList):
+    statements: List[_Statement]
+
+@dataclass
+class If(_Statement):
+    cond: Value
+    then: CodeBlock
+
+@dataclass
+class SetVar(_Statement):
+    name: str
+    value: Value
+
+@dataclass
+class Print(_Statement):
+    value: Value
+
+
+class ToAst(Transformer):
+    def STRING(self, s):
+        # Remove quotation marks
+        return s[1:-1]
+
+    def DEC_NUMBER(self, n):
+        return int(n)
+
+    @v_args(inline=True)
+    def start(self, x):
+        return x
+
+#
+#   Define Parser
+#
+
+parser = Lark("""
+    start: code_block
+
+    code_block: statement+
+
+    ?statement: if | set_var | print
+
+    if: "if" value "{" code_block "}"
+    set_var: NAME "=" value ";"
+    print: "print" value ";"
+
+    value: name | STRING | DEC_NUMBER
+    name: NAME
+
+    %import python (NAME, STRING, DEC_NUMBER)
+    %import common.WS
+    %ignore WS
+    """,
+    parser="lalr",
+)
+
+transformer = ast_utils.create_transformer(this_module, ToAst())
+
+def parse(text):
+    return transformer.transform(parser.parse(text))
+
+#
+#   Test
+#
+
+if __name__ == '__main__':
+    print(parse("""
+        a = 1;
+        if a {
+            print "a is 1";
+            a = 2;
+        }
+    """))
diff --git a/lark/ast_utils.py b/lark/ast_utils.py
new file mode 100644
index 0000000..27b00a3
--- /dev/null
+++ b/lark/ast_utils.py
@@ -0,0 +1,51 @@
+"""
+    Module of utilities for transforming a lark.Tree into a custom Abstract Syntax Tree
+"""
+
+import inspect, re
+
+from lark import Transformer, v_args
+
+class Ast(object):
+    """Abstract class
+
+    Subclasses will be collected by `create_transformer()`
+    """
+    pass
+
+class AsList(object):
+    """Abstract class
+
+    Subclasses will be instanciated with the parse results as a single list, instead of as arguments.
+    """
+
+def camel_to_snake(name):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()
+
+def _call(func, _data, children, _meta):
+    return func(*children)
+
+inline = v_args(wrapper=_call)
+
+def create_transformer(ast_module, transformer=None):
+    """Collects `Ast` subclasses from the given module, and creates a Lark transformer that builds the AST.
+
+    For each class, we create a corresponding rule in the transformer, with a matching name.
+    CamelCase names will be converted into snake_case. Example: "CodeBlock" -> "code_block".
+
+    Parameters:
+        ast_module - A Python module containing all the subclasses of `ast_utils.Ast`
+                     Classes starting with an underscore (`_`) will be skipped.
+        transformer (Optional[Transformer]) - An initial transformer. Its attributes may be overwritten.
+    """
+    t = transformer or Transformer()
+
+    for name, obj in inspect.getmembers(ast_module):
+        if not name.startswith('_') and inspect.isclass(obj):
+            if issubclass(obj, Ast):
+                if not issubclass(obj, AsList):
+                    obj = inline(obj).__get__(t)
+
+                setattr(t, camel_to_snake(name), obj)
+
+    return t
\ No newline at end of file

From e759718c3d437129155032fb874485c9c4d01064 Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Mon, 1 Mar 2021 08:03:08 -0500
Subject: [PATCH 2/2] Added pyi

---
 lark-stubs/ast_utils.pyi | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 lark-stubs/ast_utils.pyi

diff --git a/lark-stubs/ast_utils.pyi b/lark-stubs/ast_utils.pyi
new file mode 100644
index 0000000..28246cf
--- /dev/null
+++ b/lark-stubs/ast_utils.pyi
@@ -0,0 +1,17 @@
+import types
+from typing import Optional
+
+from .visitors import Transformer
+
+class Ast(object):
+    pass
+
+class AsList(object):
+    pass
+
+
+def create_transformer(
+        ast_module: types.ModuleType,
+        transformer: Optional[Transformer]=None
+) -> Transformer:
+    ...
\ No newline at end of file