<commit_before>import gen

# illegal: first term might start with X or be empty, second term starts with X
try:
    g = gen.Grammar()
    X = g.token("X")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)
        e.match(X)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass

# same thing, but with a term in the middle
try:
    g = gen.Grammar()
    X = g.token("X")
    Y = g.token("Y")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)
        with e.zero_or_more():
            e.match(Y)
        e.match(X)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass

# same thing, but in different scopes
try:
    g = gen.Grammar()
    X = g.token("X")
    Y = g.token("Y")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)

    @g.expr
    def b(e):
        with e.zero_or_more():
            e.match(Y)

    @g.expr
    def c(e):
        e.match(X)

    @g.expr
    def d(e):
        e.match(a)
        e.match(b)
        e.match(c)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass


# Unit tests
g = gen.Grammar()
W = g.token("W")
X = g.token("X")
Y = g.token("Y")
Z = g.token("Z")

@g.expr
def a(e):
    with e.maybe():
        e.match(W)
    with e.maybe():
        e.match(X)
    with e.maybe():
        e.match(Y)
    with e.maybe():
        e.match(Z)

@g.expr
def b(e):
    with e.maybe():
        e.match(W)
    with e.maybe():
        e.match(X)
    e.match(Y)
    with e.maybe():
        e.match(Z)

assert a.seq.get_first() == {"W", "X", "Y", "Z", None}
assert a.seq.get_disallowed_after() == {"W", "X", "Y", "Z", None}

assert b.seq.get_first() == {"W", "X", "Y"}
assert b.seq.get_disallowed_after() == {"Z"}


# testing MetaTokenizer
text = r"{{{ hi {{hello}} bye }}}"
tokens = [t.type for t in gen.Tokenizer().iter(text)]
assert tokens == [gen.CODE, gen.EOF], tokens


# test extract_fallbacks (tree expansion)
grammar_text = """
%fallback LETTER A B C;
%fallback WORD LETTER;
asdf = WORD LETTER A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
exp = {
    "LETTER": {"A", "B", "C"},
    "WORD": {"LETTER", "A", "B", "C"},
    "A": set(),
    "B": set(),
    "C": set(),
}
fallbackmap = gen.extract_fallbacks(parsed_doc.fallbacks, g, None)
assert fallbackmap == exp, '\ngot: ' + str(got) + '\nexp: ' + str(exp)
fallbacks = gen.Fallbacks(fallbackmap)
got = fallbacks.fallback_set("WORD", ["WORD"])
assert got == {"LETTER", "A", "B", "C"}, got
got = fallbacks.fallback_set("WORD", ["WORD", "A"])
assert got == {"LETTER", "B", "C"}, got
got = fallbacks.fallback_set("A", ["WORD", "A"])
assert got == set(), got
got = fallbacks.fallback_set("WORD", ["WORD", "LETTER"])
assert got == set(), got
got = fallbacks.fallback_set("LETTER", ["WORD", "LETTER"])
assert got == {"A", "B", "C"}, got

# test extract_fallbacks (multi-parent rejection)
grammar_text = """
%fallback LETTER A B C;
%fallback WORD LETTER C;
asdf = WORD LETTER A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
try:
    raise ValueError(gen.extract_fallbacks(parsed_doc.fallbacks, g, None))
except gen.RenderedError as e:
    assert "fallback to multiple other types" in str(e), e

# test extract_fallbacks (multi-parent rejection)
grammar_text = """
%fallback A B;
%fallback B C;
%fallback C A;
asdf = A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
try:
    raise ValueError(gen.extract_fallbacks(parsed_doc.fallbacks, g, None))<commit_msg>gen.py: add test for all_fallbacks<commit_after>import gen

# illegal: first term might start with X or be empty, second term starts with X
try:
    g = gen.Grammar()
    X = g.token("X")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)
        e.match(X)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass

# same thing, but with a term in the middle
try:
    g = gen.Grammar()
    X = g.token("X")
    Y = g.token("Y")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)
        with e.zero_or_more():
            e.match(Y)
        e.match(X)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass

# same thing, but in different scopes
try:
    g = gen.Grammar()
    X = g.token("X")
    Y = g.token("Y")

    @g.expr
    def a(e):
        with e.maybe():
            e.match(X)

    @g.expr
    def b(e):
        with e.zero_or_more():
            e.match(Y)

    @g.expr
    def c(e):
        e.match(X)

    @g.expr
    def d(e):
        e.match(a)
        e.match(b)
        e.match(c)

    g.check()
    1/0
except gen.FirstFollow as e:
    pass


# Unit tests
g = gen.Grammar()
W = g.token("W")
X = g.token("X")
Y = g.token("Y")
Z = g.token("Z")

@g.expr
def a(e):
    with e.maybe():
        e.match(W)
    with e.maybe():
        e.match(X)
    with e.maybe():
        e.match(Y)
    with e.maybe():
        e.match(Z)

@g.expr
def b(e):
    with e.maybe():
        e.match(W)
    with e.maybe():
        e.match(X)
    e.match(Y)
    with e.maybe():
        e.match(Z)

assert a.seq.get_first() == {"W", "X", "Y", "Z", None}
assert a.seq.get_disallowed_after() == {"W", "X", "Y", "Z", None}

assert b.seq.get_first() == {"W", "X", "Y"}
assert b.seq.get_disallowed_after() == {"Z"}


# testing MetaTokenizer
text = r"{{{ hi {{hello}} bye }}}"
tokens = [t.type for t in gen.Tokenizer().iter(text)]
assert tokens == [gen.CODE, gen.EOF], tokens


# test extract_fallbacks (tree expansion)
grammar_text = """
%fallback LETTER A B C;
%fallback WORD LETTER;
asdf = WORD LETTER A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
exp = {
    "LETTER": {"A", "B", "C"},
    "WORD": {"LETTER", "A", "B", "C"},
    "A": set(), "B": set(), "C": set(),
}
fallbackmap = gen.extract_fallbacks(parsed_doc.fallbacks, g, None)
assert fallbackmap == exp, '\ngot: ' + str(got) + '\nexp: ' + str(exp)
fallbacks = gen.Fallbacks(fallbackmap)
got = fallbacks.fallback_set("WORD", ["WORD"])
assert got == {"LETTER", "A", "B", "C"}, got
got = fallbacks.fallback_set("WORD", ["WORD", "A"])
assert got == {"LETTER", "B", "C"}, got
got = fallbacks.fallback_set("A", ["WORD", "A"])
assert got == set(), got
got = fallbacks.fallback_set("WORD", ["WORD", "LETTER"])
assert got == set(), got
got = fallbacks.fallback_set("LETTER", ["WORD", "LETTER"])
assert got == {"A", "B", "C"}, got

# test all_fallbacks (tree expansion)
grammar_text = """
%fallback LETTER A B C;
%fallback NUM N1 N2 N3;
%fallback WORD LETTER NUM;
thing1 = LETTER | N2;
thing2 = NUM | A;
thing3 = WORD | C;
asdf = (thing1 | thing2 | thing3) B N1 N3;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
fallbackmap = gen.extract_fallbacks(parsed_doc.fallbacks, g, None)
exp = {
    "LETTER": {"A", "B", "C"},
    "NUM": {"N1", "N2", "N3"},
    "WORD": {"LETTER", "A", "B", "C", "NUM", "N1", "N2", "N3"},
    "A": set(), "B": set(), "C": set(), "N1": set(), "N2": set(), "N3": set(),
}
assert fallbackmap == exp, fallbackmap
fallbacks = gen.Fallbacks(fallbackmap)
# first look at thing1, thing2, and thing3 directly
thing1 = g.exprs["thing1"].seq.get_first()
thing2 = g.exprs["thing2"].seq.get_first()
thing3 = g.exprs["thing3"].seq.get_first()
got = fallbacks.all_fallbacks(thing1)
assert got == {"A", "B", "C"}, got
got = fallbacks.all_fallbacks(thing2)
assert got == {"N1", "N2", "N3"}, got
got = fallbacks.all_fallbacks(thing3)
assert got == {"LETTER", "A", "B", "NUM", "N1", "N2", "N3"}, got
# now imagine we are looking at the branch statement in `asdf`
exclude = g.exprs["asdf"].seq.get_first()
got = fallbacks.all_fallbacks(thing1, exclude)
assert got == {"B"}, got
got = fallbacks.all_fallbacks(thing2, exclude)
assert got == {"N1", "N3"}, got
got = fallbacks.all_fallbacks(thing3, exclude)
assert got == set(), got

# test extract_fallbacks (multi-parent rejection)
grammar_text = """
%fallback LETTER A B C;
%fallback WORD LETTER C;
asdf = WORD LETTER A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
try:
    raise ValueError(gen.extract_fallbacks(parsed_doc.fallbacks, g, None))
except gen.RenderedError as e:
    assert "fallback to multiple other types" in str(e), e

# test extract_fallbacks (multi-parent rejection)
grammar_text = """
%fallback A B;
%fallback B C;
%fallback C A;
asdf = A B C;
"""
parsed_doc = gen.parse_doc(grammar_text)
g = parsed_doc.build_grammar(grammar_text)
try:
    raise ValueError(gen.extract_fallbacks(parsed_doc.fallbacks, g, None))
{'fraction_non_alphanumeric': 0.18534935261868657, 'fraction_numerical': 0.006532135775107897, 'mean_word_length': 3.0965121834687053, 'pattern_counts': {'":': 19, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5350451', 'n_tokens_mistral': 3369, 'n_tokens_neox': 3339, 'n_words': 875}