lib.py

import itertools
import json


class TaxonomicNode:
    __slots__ = ("name", "index", "root", "_children")

    def __init__(self, name, index, root):
        self.name = name
        self.index = index
        self.root = root
        self._children = {}

    def add(self, name):
        added = 0
        if not name:
            return added

        first, rest = name[0], name[1:]
        if first not in self._children:
            self._children[first] = TaxonomicNode(first, self.root.size, self.root)
            self.root.size += 1

        self._children[first].add(rest)

    def children(self, name):
        if not name:
            return set((child.name, child.index) for child in self._children.values())

        first, rest = name[0], name[1:]
        if first not in self._children:
            return set()

        return self._children[first].children(rest)

    def __iter__(self):
        yield self.name, self.index

        for child in self._children.values():
            for name, index in child:
                yield f"{self.name} {name}", index

    @classmethod
    def from_dict(cls, dct, root):
        node = cls(dct["name"], dct["index"], root)
        node._children = {
            child["name"]: cls.from_dict(child, root) for child in dct["children"]
        }
        return node


class TaxonomicTree:
    """
    Efficient structure for finding taxonomic names and their descendants.
    Also returns an integer index i for each possible name.
    """

    def __init__(self):
        self.kingdoms = {}
        self.size = 0

    def add(self, name: list[str]):
        if not name:
            return

        first, rest = name[0], name[1:]
        if first not in self.kingdoms:
            self.kingdoms[first] = TaxonomicNode(first, self.size, self)
            self.size += 1

        self.kingdoms[first].add(rest)

    def children(self, name=None):
        if not name:
            return set(
                (kingdom.name, kingdom.index) for kingdom in self.kingdoms.values()
            )

        first, rest = name[0], name[1:]
        if first not in self.kingdoms:
            return set()

        return self.kingdoms[first].children(rest)

    def __iter__(self):
        for kingdom in self.kingdoms.values():
            yield from kingdom

    def __len__(self):
        return self.size

    @classmethod
    def from_dict(cls, dct):
        tree = cls()
        tree.kingdoms = {
            kingdom["name"]: TaxonomicNode.from_dict(kingdom, tree)
            for kingdom in dct["kingdoms"]
        }
        tree.size = dct["size"]
        return tree


class TaxonomicJsonEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, TaxonomicNode):
            return {
                "name": obj.name,
                "index": obj.index,
                "children": list(obj._children.values()),
            }
        elif isinstance(obj, TaxonomicTree):
            return {
                "kingdoms": list(obj.kingdoms.values()),
                "size": obj.size,
            }
        else:
            super().default(self, obj)


def batched(iterable, n):
    # batched('ABCDEFG', 3) --> ABC DEF G
    if n < 1:
        raise ValueError("n must be at least one")
    it = iter(iterable)
    while batch := tuple(itertools.islice(it, n)):
        yield zip(*batch)