<commit_before>#! /usr/bin/env python3

"""
The tree module extract, from a table, a tree limited to a given depth.
"""

from collections import defaultdict
from wagoner.utils import random_weighted_choice

from wagoner.word import weighted_choices


def extract_tree(table, length, prefix=0, flatten=False):
    """
    Extract from table a tree for word length, taking only prefixes of prefix
    length (if greater than 0) into account to compute successors.

    :param table: the table to extract the tree from;
    :param length: the length of words generated by the extracted tree;
                   greater or equal to 1;
    :param prefix: if greater than 0, the length of the prefixes used for
                   computing successors;
    :param flatten: whether to flatten the table or not;
    :return: the tree corresponding to words of length from table.
    """
    # Build the expanded tree with necessary suffix and length
    tree = defaultdict(dict)  # The tree
    pending = {(">", 0)}  # The nodes to expand
    while pending:
        suffix, size = pending.pop()
        choices = weighted_choices(suffix, table, flatten=flatten)
        if size < length:
            # The word length is not reached yet, expand
            for successor, weight in choices.items():
                expanded = suffix + successor
                if prefix > 0:
                    expanded = expanded[-prefix:]
                new_node = (expanded, size + 1)
                tree[(suffix, size)][new_node] = weight
                pending.add(new_node)
        else:
            # The word length is reached, only add < if present
            if "<" in choices:
                tree[(suffix, size)][("<", size + 1)] = 1
            else:
                tree[(suffix, size)] = dict()
    return trim_tree(tree)


def trim_tree(tree):
    """
    Remove the dead branches of tree, that is, the resulting tree accepts the
    same language as the original one (that is, the same words that end with
    the < character), but parts of the tree that lead to nothing are removed.

    :param tree: the tree;
    :return: the tree without dead branches.
    """
    # Remove empty nodes
    new_tree = {k: v for k, v in tree.items() if v}
    # Remove missing successors
    new_tree = {k: {successor: weight for successor, weight in v.items()
                    if successor in new_tree or successor[0] == "<"}
                for k, v in new_tree.items()}
    while tree != new_tree:
        tree = new_tree
        # Remove empty nodes
        new_tree = {k: v for k, v in tree.items() if v}
        # Remove missing successors
        new_tree = {k: {successor: weight for successor, weight in v.items()
                        if successor in new_tree or successor[0] == "<"}
                    for k, v in new_tree.items()}
    return new_tree


def random_word(tree):
    """
    Return a random word from tree. The length of the word depends on the tree.

    :param tree: the tree to generate a word from;
    :return: a random word from tree.
    """
    word = ""
    current = (">", 0)
    while current[0] != "<":
        choices = tree[current]
        choice = random_weighted_choice(choices)
        current = choice
        word += current[0][-1]
    return word[:-1]<commit_msg>Add commandline call to wagoner.tree: extract the tree and generate words.<commit_after>#! /usr/bin/env python3

"""
The tree module extract, from a table, a tree limited to a given depth.
"""

import argparse
import pickle
from collections import defaultdict
from wagoner.utils import random_weighted_choice, nonzero_natural, natural

from wagoner.word import weighted_choices


def extract_tree(table, length, prefix=0, flatten=False):
    """
    Extract from table a tree for word length, taking only prefixes of prefix
    length (if greater than 0) into account to compute successors.

    :param table: the table to extract the tree from;
    :param length: the length of words generated by the extracted tree;
                   greater or equal to 1;
    :param prefix: if greater than 0, the length of the prefixes used for
                   computing successors;
    :param flatten: whether to flatten the table or not;
    :return: the tree corresponding to words of length from table.
    """
    # Build the expanded tree with necessary suffix and length
    tree = defaultdict(dict)  # The tree
    pending = {(">", 0)}  # The nodes to expand
    while pending:
        suffix, size = pending.pop()
        choices = weighted_choices(suffix, table, flatten=flatten)
        if size < length:
            # The word length is not reached yet, expand
            for successor, weight in choices.items():
                expanded = suffix + successor
                if prefix > 0:
                    expanded = expanded[-prefix:]
                new_node = (expanded, size + 1)
                tree[(suffix, size)][new_node] = weight
                pending.add(new_node)
        else:
            # The word length is reached, only add < if present
            if "<" in choices:
                tree[(suffix, size)][("<", size + 1)] = 1
            else:
                tree[(suffix, size)] = dict()
    return trim_tree(tree)


def trim_tree(tree):
    """
    Remove the dead branches of tree, that is, the resulting tree accepts the
    same language as the original one (that is, the same words that end with
    the < character), but parts of the tree that lead to nothing are removed.

    :param tree: the tree;
    :return: the tree without dead branches.
    """
    # Remove empty nodes
    new_tree = {k: v for k, v in tree.items() if v}
    # Remove missing successors
    new_tree = {k: {successor: weight for successor, weight in v.items()
                    if successor in new_tree or successor[0] == "<"}
                for k, v in new_tree.items()}
    while tree != new_tree:
        tree = new_tree
        # Remove empty nodes
        new_tree = {k: v for k, v in tree.items() if v}
        # Remove missing successors
        new_tree = {k: {successor: weight for successor, weight in v.items()
                        if successor in new_tree or successor[0] == "<"}
                    for k, v in new_tree.items()}
    return new_tree


def random_word(tree):
    """
    Return a random word from tree. The length of the word depends on the tree.

    :param tree: the tree to generate a word from;
    :return: a random word from tree.
    """
    word = ""
    current = (">", 0)
    while current[0] != "<":
        choices = tree[current]
        choice = random_weighted_choice(choices)
        current = choice
        word += current[0][-1]
    return word[:-1]


def process_arguments():
    """
    Process the command line arguments. The arguments are:
     * the table to generate random from;
     * -l (or --length) for the length of generated words (default: 10);
     * -p (or --prefix) for the maximum of prefixes to consider (default: 0);
     * -c (or --count) for the number of words to generate (default: 10);
     * -f (or --flatten) if the table must be flattened before generation.
    """
    parser = argparse.ArgumentParser(description="Generate random words from "
                                                 "the given table")
    parser.add_argument("table", type=argparse.FileType('rb'),
                        help="the table")
    parser.add_argument("--length", "-l", type=nonzero_natural, default=10,
                        dest="length", help="the length of generated words "
                                            "(default: 10)")
    parser.add_argument("--prefix", "-p", type=natural, default=0,
                        dest="prefix", help="if not 0, the maximum length of "
                                            "prefixes to consider when "
                                            "choosing the next character "
                                            "(default: 0)")
    parser.add_argument("--count", "-c", type=natural, default=10,
                        dest="count", help="the number of words to generate "
                                           "(default: 10)")
    parser.add_argument("--flatten", "-f", action="store_true", default=False,
                        dest="flatten", help="flatten the table")
    return parser.parse_args()

if __name__ == "__main__":
    args = process_arguments()
    table = pickle.load(args.table)
    tree = extract_tree(table, args.length, prefix=args.prefix,
                        flatten=args.flatten)
    for _ in range(args.count):
        print(random_word(tree))
{'fraction_non_alphanumeric': 0.09445989800649049, 'fraction_numerical': 0.005795085767269356, 'mean_word_length': 2.093940480458946, 'pattern_counts': {'":': 3, '<': 19, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '756452', 'n_tokens_mistral': 2338, 'n_tokens_neox': 2236, 'n_words': 918}