Python reconstructor list comprehension error

[ellieyhcheng]

Describe the bug

Using the official Lark Python parser and using the PythonReconstructor from the examples, list comprehension produces AssertionError when the code is valid Python code.

E.g. The inputs "[a for a in range(10)]" produces the error

AssertionError: (NonTerminal('comprehension{test_or_star_expr}'), Tree(Token('RULE', 'comprehension'), [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'a')])]), Tree(Token('RULE', 'comp_fors'), [Tree(Token('RULE', 'comp_for'), [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'a')])]), Tree('funccall', [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'range')])]), Tree(Token('RULE', 'arguments'), [Tree(Token('RULE', 'number'), [Token('DEC_NUMBER', '10')])])])])])]))

To Reproduce

The reproduction test script is the code from the example but given the input "[a for a in range(10)]"

This is the full script:

from lark import Token, Lark
from lark.reconstruct import Reconstructor
from lark.indenter import PythonIndenter

# Official Python grammar by Lark
python_parser3 = Lark.open_from_package(
    "lark",
    "python.lark",
    ["grammars"],
    parser="lalr",
    postlex=PythonIndenter(),
    start="file_input",
    maybe_placeholders=False,  # Necessary for reconstructor
)

SPACE_AFTER = set(',+-*/~@<>="|:')
SPACE_BEFORE = (SPACE_AFTER - set(",:")) | set("'")


def special(sym):
    return Token("SPECIAL", sym.name)


def postproc(items):
    stack = ["\n"]
    actions = []
    last_was_whitespace = True
    for item in items:
        if isinstance(item, Token) and item.type == "SPECIAL":
            actions.append(item.value)
        else:
            if actions:
                assert actions[0] == "_NEWLINE" and "_NEWLINE" not in actions[1:], actions

                for a in actions[1:]:
                    if a == "_INDENT":
                        stack.append(stack[-1] + " " * 4)
                    else:
                        assert a == "_DEDENT"
                        stack.pop()
                actions.clear()
                yield stack[-1]
                last_was_whitespace = True
            if not last_was_whitespace:
                if item[0] in SPACE_BEFORE:
                    yield " "
            yield item
            last_was_whitespace = item[-1].isspace()
            if not last_was_whitespace:
                if item[-1] in SPACE_AFTER:
                    yield " "
                    last_was_whitespace = True
    yield "\n"


class PythonReconstructor:
    def __init__(self, parser):
        self._recons = Reconstructor(parser, {"_NEWLINE": special, "_DEDENT": special, "_INDENT": special})

    def reconstruct(self, tree):
        return self._recons.reconstruct(tree, postproc)


python_reconstructor = PythonReconstructor(python_parser3)

self_contents = "[a for a in range(10)]"

tree = python_parser3.parse(self_contents + "\n")
output = python_reconstructor.reconstruct(tree)

tree_new = python_parser3.parse(output)
print(tree.pretty())
print(tree_new.pretty())
# assert tree.pretty() == tree_new.pretty()
assert tree == tree_new

Running the script produces the described error. Here's the full trace:

Traceback (most recent call last):
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/visitors.py", line 115, in _call_userfunc
    f = getattr(self, tree.data)
        ^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'WriteTokensTransformer' object has no attribute 'list_comprehension'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/ellieyhc/Documents/test/test.py", line 69, in <module>
    output = python_reconstructor.reconstruct(tree)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/test.py", line 61, in reconstruct
    return self._recons.reconstruct(tree, postproc)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/reconstruct.py", line 102, in reconstruct
    for item in x:
                ^
  File "/Users/ellieyhc/Documents/test/test.py", line 28, in postproc
    for item in items:
                ^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/reconstruct.py", line 92, in _reconstruct
    yield from self._reconstruct(item)
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/reconstruct.py", line 92, in _reconstruct
    yield from self._reconstruct(item)
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/reconstruct.py", line 88, in _reconstruct
    res = self.write_tokens.transform(unreduced_tree)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/visitors.py", line 286, in transform
    return self._transform_tree(tree)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/visitors.py", line 280, in _transform_tree
    return self._call_userfunc(tree)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/visitors.py", line 117, in _call_userfunc
    return self.__default__(tree.data, children, tree.meta)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ellieyhc/Documents/test/venv/lib/python3.12/site-packages/lark/reconstruct.py", line 58, in __default__
    assert NonTerminal(x.data) == sym, (sym, x)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: (NonTerminal('comprehension{test_or_star_expr}'), Tree(Token('RULE', 'comprehension'), [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'a')])]), Tree(Token('RULE', 'comp_fors'), [Tree(Token('RULE', 'comp_for'), [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'a')])]), Tree('funccall', [Tree('var', [Tree(Token('RULE', 'name'), [Token('NAME', 'range')])]), Tree(Token('RULE', 'arguments'), [Tree(Token('RULE', 'number'), [Token('DEC_NUMBER', '10')])])])])])]))