<commit_before>from ..grammar_parsing.parse_grammar import parse_grammar_file
from ..grammar_parsing.production import *
from ..grammar_parsing.production_part import *

from collections import defaultdict
from typing import Dict, List, Optional, Tuple


# TODO: Parameters most be processed prior to line generation so lifting can occur.
# TODO: Need to handle special tokens such as NAME, CLASS, and NUMBER.
    
    
# Three levels:
#
#   1. Rules that are superclasses of other rules (through aliasing).
#   2. Rules that are superclasses only of productions.
#   3. Productions.
#
# Classes in (1) must be written first, then (2), then (3).

Param = Tuple[str, Optional[str]]


class Arg:
    def __init__(self, name: Optional[str], type: Optional[str], wrappers=None):
        self.name = name
        self.type = type
        if wrappers:
            self.wrappers = wrappers
        else:
            self.wrappers = []


class ASTNodeGenerator:
    BASE_AST_CLASS_NAME = 'AST'

    def __init__(self, filename: str):
        self.tree = None
        self.parsed_rules = parse_grammar_file(filename)
        self.build_rule_tree(self.parsed_rules)
        for rule, production_list in self.parsed_rules.items():
            if len(production_list) == 1:
                self.make_ast_node_class_from_single_production(rule, production_list[0])
            else:
                self.make_ast_node_classes_from_production_list(rule, production_list)

    def generate_text(self) -> str:
        # Add the header.
        lines = [
            "# This module was automatically generated.",
            "",
            "from typing import List, Optional",
            "",
            "",
        ]
        # Set the depths of the nodes to accommodate parameters.
        from itertools import chain
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            node.depth = max(chain([0], (other.depth for other in chain(node.parents, map(lambda n: self.tree[n], node.params))))) + 1
            queue += node.children
        # Organize nodes by depth in the tree.
        tiers = defaultdict(list)
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            tiers[node.depth].append(node)
            queue += node.children
        # Iterating by tier, add nodes' lines.
        for depth in sorted(tiers.keys()):
            tier = tiers[depth]
            for node in tier:
                lines += node.lines
                lines += ["", ""]
        # Remove extra blank line at the end and return.
        del(lines[-1])
        return '\n'.join(lines)

    class ClassTree:
        def __init__(self):
            self._nodes = {}
            root = ASTNodeGenerator.ClassTreeNode('AST')
            root.lines = [
                'class AST:',
                '    pass'
            ]
            self._nodes['AST'] = root
            self.root = root

        def __getitem__(self, item):
            return self._nodes[item]

        def __setitem__(self, key, value):
            if key in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            self._nodes[key] = value

        def add(self, child_name: str, parent_name: str):<commit_msg>Use constant for root class name<commit_after>from ..grammar_parsing.parse_grammar import parse_grammar_file
from ..grammar_parsing.production import *
from ..grammar_parsing.production_part import *

from collections import defaultdict
from typing import Dict, List, Optional, Tuple


# TODO: Parameters most be processed prior to line generation so lifting can occur.
# TODO: Need to handle special tokens such as NAME, CLASS, and NUMBER.
    
    
# Three levels:
#
#   1. Rules that are superclasses of other rules (through aliasing).
#   2. Rules that are superclasses only of productions.
#   3. Productions.
#
# Classes in (1) must be written first, then (2), then (3).

Param = Tuple[str, Optional[str]]


class Arg:
    def __init__(self, name: Optional[str], type: Optional[str], wrappers=None):
        self.name = name
        self.type = type
        if wrappers:
            self.wrappers = wrappers
        else:
            self.wrappers = []


class ASTNodeGenerator:
    BASE_AST_CLASS_NAME = 'AST'

    def __init__(self, filename: str):
        self.tree = None
        self.parsed_rules = parse_grammar_file(filename)
        self.build_rule_tree(self.parsed_rules)
        for rule, production_list in self.parsed_rules.items():
            if len(production_list) == 1:
                self.make_ast_node_class_from_single_production(rule, production_list[0])
            else:
                self.make_ast_node_classes_from_production_list(rule, production_list)

    def generate_text(self) -> str:
        # Add the header.
        lines = [
            "# This module was automatically generated.",
            "",
            "from typing import List, Optional",
            "",
            "",
        ]
        # Set the depths of the nodes to accommodate parameters.
        from itertools import chain
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            node.depth = max(chain([0], (other.depth for other in chain(node.parents, map(lambda n: self.tree[n], node.params))))) + 1
            queue += node.children
        # Organize nodes by depth in the tree.
        tiers = defaultdict(list)
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            tiers[node.depth].append(node)
            queue += node.children
        # Iterating by tier, add nodes' lines.
        for depth in sorted(tiers.keys()):
            tier = tiers[depth]
            for node in tier:
                lines += node.lines
                lines += ["", ""]
        # Remove extra blank line at the end and return.
        del(lines[-1])
        return '\n'.join(lines)

    class ClassTree:
        def __init__(self):
            self._nodes = {}
            root = ASTNodeGenerator.ClassTreeNode(ASTNodeGenerator.BASE_AST_CLASS_NAME)
            root.lines = [
                'class ' + root.name + ':',
                '    pass'
            ]
            self._nodes[ASTNodeGenerator.BASE_AST_CLASS_NAME] = root
            self.root = root

        def __getitem__(self, item):
            return self._nodes[item]

        def __setitem__(self, key, value):
            if key in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            self._nodes[key] = value

        def add(self, child_name: str, parent_name: str):
{'fraction_non_alphanumeric': 0.09504443440277151, 'fraction_numerical': 0.00391625244765778, 'mean_word_length': 2.014071720381298, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '1202737', 'n_tokens_mistral': 1867, 'n_tokens_neox': 1826, 'n_words': 542}