from nodes import *
from tokens import python_lexer

@with_lexer(python_lexer)
grammar python_grammar {
    name <- Id(@Identifier)
    number <- NumberLit(@Number)
    string <- StringLit(@String)
    cat_string <- ConcatStringLit(string list+(string))
    nl <- NL(@Newline)
    @main_rule
    main_rule <- FileNode(
        list+(pick(discard(list*(nl)) stmt discard(list*(nl)))) @Termination
    )
    decorator <- Decorator("@" dotted_name ?pick("(" arg_list ")") @Newline)
    decorators <- list+(decorator)
    decorated <- Decorated(decorators or(class_def | func_def))
    func_def <- FuncDef("def" name parameters ":" suite)
    parameters <- pick("(" ?varargslist ")")
    varargslist <- Params(
        list*(
            SingleParam(
                VarArgsFlag("*") KwArgsFlag("**") fpdef ?pick("=" test)
            ),
            ","
        )
    )
    fpdef <- or(name | pick("(" name_list ")"))
    name_list <- pick(list+(name, ",") ?",")
    stmt <- or(simple_stmt | compound_stmt)
    simple_stmt <- pick(
        or(small_stmt | pick(list+(small_stmt, ";") ?";"))
        @Newline
    )
    small_stmt <- or(
        | expr_stmt
        | print_stmt
        | del_stmt
        | pass_stmt
        | flow_stmt
        | import_stmt
        | global_stmt
        | exec_stmt
        | assert_stmt
    )
    expr_stmt <- or(
        | AugAssignStmt(
            test_list
            Op(
                or(
                    | "+="
                    | "-="
                    | "*="
                    | "/="
                    | "%="
                    | "&="
                    | "|="
                    | "^="
                    | "<<="
                    | ">>="
                    | "**="
                    | "//="
                )
            )
            or(yield_expr | test_list)
        )
        | AssignStmt(test_list list+(pick("=" or(yield_expr | test_list))))
        | test_list
    )
    print_stmt <- or(
        | PrintStmt("print" test_list)
        | StreamPrintStmt("print" ">>" test "," test_list)
    )
    del_stmt <- DelStmt("del" expr_list)
    pass_stmt <- PassStmt("pass")
    flow_stmt <- or(
        | break_stmt
        | continue_stmt
        | return_stmt
        | raise_stmt
        | yield_stmt
    )
    break_stmt <- BreakStmt("break")
    continue_stmt <- ContinueStmt("continue")
    return_stmt <- ReturnStmt("return" ?test_list)
    yield_stmt <- yield_expr
    raise_stmt <- RaiseStmt("raise" ?test_list)
    import_stmt <- or(import_name | import_from)
    import_name <- ImportName("import" dotted_as_names)
    dot <- Dot(".")
    import_from <- ImportFrom(
        "from"
        or(dotted_name | RelName(list+(dot) ?dotted_name))
        "import"
        or(| ImportStar("*") | pick("(" import_as_names ")") | import_as_names)
    )
    as_name <- AsNameNode(name "as" name)
    dotted_as_name <- AsNameNode(dotted_name "as" name)
    import_as_names <- pick(list+(or(as_name | name), ",") ?",")
    dotted_as_names <- pick(list+(or(dotted_as_name | dotted_name), ",") ?",")
    dotted_name <- or(DottedName(dotted_name "." name) | name)
    global_stmt <- GlobalStmt("global" name_list)
    exec_stmt <- ExecStmt("exec" expr ?pick("in" test_list))
    assert_stmt <- AssertStmt("assert" test ?pick("," test))
    compound_stmt <- or(
        | if_stmt
        | while_stmt
        | for_stmt
        | try_stmt
        | with_stmt
        | func_def
        | class_def
        | decorated
    )
    else_part <- ElsePart("else" ":" suite)
    if_stmt <- IfStmt(
        "if"
        test
        ":"
        suite
        list*(pick("elif" ElifBranch(test ":" suite)))
        ?else_part
    )

    while_stmt <- WhileStmt("while" test ":" suite ?else_part)

    for_stmt <- ForStmt("for" expr_list "in" test_list ":" suite ?else_part)

    try_stmt <- TryStmt(
        "try"
        ":"
        suite
        list*(
            ExceptPart("except" ?AsNameNode(test ?pick("as" test)) ":" suite)
        )
        ?else_part
        ?pick("finally" ":" suite)
    )

    with_stmt <- WithStmt("with" list+(with_item, ",") ":" suite)

    with_item <- AsNameNode(test ?pick("as" expr))

    suite <- or(
        | pick(
            discard(list*(nl))
            @Indent
            list+(pick(discard(list*(nl)) stmt discard(list*(nl))))
            @Dedent
        )
        | simple_stmt
    )
    test <- or(| lambdef | IfExpr(or_test "if" or_test "else" test) | or_test)
    or_test <- or(OrOp(or_test "or" and_test) | and_test)
    and_test <- or(AndOp(and_test "and" not_test) | not_test)
    not_test <- or(NotOp("not" not_test) | comparison)
    comparison <- or(
        | CompOp(
            comparison
            or(
                | CompOpKind.Lt("<")
                | CompOpKind.Gt(">")
                | CompOpKind.Eq("==")
                | CompOpKind.Gte(">=")
                | CompOpKind.Lte("<=")
                | CompOpKind.Diamond("<>")
                | CompOpKind.Noteq("!=")
                | CompOpKind.In("in")
                | CompOpKind.Notin("not" "in")
                | CompOpKind.Isnot("is" "not")
                | CompOpKind.Is("is")
            )
            expr
        )
        | expr
    )
    expr <- or(OrExpr(expr "|" xor_expr) | xor_expr)
    xor_expr <- or(XorExpr(xor_expr "^" and_expr) | and_expr)
    and_expr <- or(| AndExpr(and_expr "&" shift_expr) | shift_expr)
    shift_expr <- or(
        | ShiftExpr(shift_expr Op(or("<<" | ">>")) arith_expr)
        | arith_expr
    )
    arith_expr <- or(| ArithExpr(arith_expr Op(or("+" | "-")) term) | term)
    term <- or(| Term(term Op(or("*" | "/" | "%" | "//")) factor) | factor)
    factor <- or(Factor(Op(or("+" | "-" | "~")) factor) | power)
    power <- or(Power(atom_expr "**" factor) | atom_expr)
    atom_expr <- or(
        | DottedName(atom_expr "." name)
        | CallExpr(atom_expr "(" arg_list ")")
        | SubscriptExpr(atom_expr "[" subscript_list "]")
        | atom
    )

    dict_assoc <- DictAssoc(test ":" test)
    yield_expr <- YieldExpr("yield" ?test_list)
    atom <- or(
        | pick("(" yield_expr ")")
        | ListGen("(" test list_for ")")
        | TupleLit("(" ?test_list ")")
        | ListComp("[" test list_for "]")
        | ListLit("[" empty_test_list "]")
        | SetComp("{" test comp_for "}")
        | set_lit
        | DictComp("{" dict_assoc comp_for "}")
        | DictLit("{" pick(list+(dict_assoc, ",") ?",") "}")
        | InlineEval("`" test_list "`")
        | name
        | number
        | cat_string
        | string
    )
    set_lit <- SetLit("{" empty_test_list "}")
    lambdef <- LambdaDef("lambda" varargslist ":" test)
    subscript_list <- pick(list+(subscript, ",") ?",")
    subscript <- or(
        | EllipsisExpr("." "." ".")
        | ExtSliceExpr(?test ":" ?test ":" ?test)
        | SliceExpr(?test ":" ?test)
        | test
    )
    expr_list <- pick(list+(expr, ",") ?",")
    test_list <- pick(list+(test, ",") ?",")
    empty_test_list <- pick(list*(test, ",") ?",")
    class_def <- ClassDef("class" name ?pick("(" ?test_list ")") ":" suite)
    arg_list <- pick(
        list*(
            or(
                | ArgGen(test comp_for)
                | ArgAssoc(?pick(test "=") test)
                | VarArgs("*" test)
                | KwArgs("**" test)
            ),
            ","
        )
        ?","
    )
    list_iter <- or(list_for | list_if)
    list_for <- CompForL("for" expr_list "in" test_list ?list_iter)
    list_if <- CompIf("if" test ?list_iter)
    comp_iter <- or(comp_for | comp_if)
    comp_for <- CompFor("for" expr_list "in" or_test ?comp_iter)
    comp_if <- CompIf("if" test ?comp_iter)
}