dost.py (cda51fd) - dost

dost.py

#!/usr/bin/python3
import argparse
from collections import defaultdict
import re
import sys
from typing import Generator

def cost(token_child: tuple[str, "Trie"]) -> tuple:
    token, child = token_child
    weight = child.deep_weight
    return child.deep_usage / weight, -weight, child.min_depth, len(token), token

class Trie:
    def __init__(self, weight: int = 0):
        self.children: defaultdict[str, Trie] = defaultdict(Trie)
        self.weight: int = weight
        self.deep_weight: int = weight
        self.usage: int = 0
        self.deep_usage: int = 0
        self.min_depth: int = 0 if weight else 999999

def __str__(self) -> str:
        s = f"[{self.usage}/{self.weight}]"
        if self.children:
            s += f"({" ".join('"' + k + '"' + str(v) for (k, v) in self.children.items())})"
        return s

def add(self, word: list[str], weight: int = 1):
        self.deep_weight += weight
        self.min_depth = min(self.min_depth, len(word))
        if word:
            self.children[word[0]].add(word[1:], weight)
        else:
            self.weight += weight

def best(self) -> tuple[list[str], "Trie"]:
        if self.weight > self.usage:
            return [], self
        elif self.children:
            token, child = min(self.children.items(), key=cost)
            suffix, leaf = child.best()
            if not leaf:
                return None, None
            return [token] + suffix, leaf
        else:
            return None, None

def use(self, word: list[str]):
        self.deep_usage += 1
        if word:
            self.children[word[0]].use(word[1:])
        else:
            assert self.usage < self.weight
            self.usage += 1

def traverse(self) -> Generator[tuple[list[str], "Trie"], None, None]:
        while True:
            tokens, leaf = self.best()
            if not leaf:
                assert self.deep_usage == self.deep_weight
                break
            yield tokens
            self.use(tokens)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="dost", description="Diverse-order sampling by trie",
        epilog="Tends to run faster on shorter tokens"
    )
    parser.add_argument("files", default="-", nargs="*")
    parser.add_argument("-i", "--in-place", action="store_true")
    parser.add_argument(
        "-s", "--separators", default=" \t",
        help="characters that separate tokens in an input item (default space and tab)"
    )
    parser.add_argument(
        "-t", "--token",
        help="structural regexp for a token of an input item, overrides -s"
    )
    args = parser.parse_args()
    if args.in_place and len(args.files) != 1:
        sys.exit(1)
    RE_FLAGS = re.MULTILINE | re.DOTALL
    if args.token:
        token_pattern = args.token
    else:
        token_pattern = f"[{args.separators}]*[^{args.separators}\r\n]+"
    entry_pattern = re.compile(f"({token_pattern})+", RE_FLAGS)
    token_pattern = re.compile(token_pattern, RE_FLAGS)
    t = Trie()
    for filename in args.files:
        with sys.stdin if filename == "-" else open(filename, "r") as handle:
            for word in re.finditer(entry_pattern, handle.read()):
                t.add(re.findall(token_pattern, word.group(0)))
    output_handle = sys.stdout
    if args.in_place and args.files[0] != "-":
        output_handle = open(args.files[0], "w")
    for word in t.traverse():
        output_handle.write("".join(word) + "\n")
    output_handle.close()

dost

Fix 6993264