cobblemon-survival-pack/tools/fieldguide/loot.py

"""Loot tables flattened to (item, approximate chance, count, notes).

Chances are estimates for a reader (expected rolls x weight share x
random_chance conditions); conditions that are not random become notes.
"""
from .datasource import norm_id


def number(v, default=1.0):
    """A loot number provider -> (min, max)."""
    if v is None:
        return default, default
    if isinstance(v, (int, float)):
        return float(v), float(v)
    t = v.get("type", "minecraft:uniform")
    if "min" in v or "max" in v:
        lo, hi = number(v.get("min"), 0)[0], number(v.get("max"), 0)[1]
        return lo, hi
    if t.endswith("binomial"):
        n = number(v.get("n"))[1]
        return 0.0, n
    if t.endswith("constant"):
        return float(v.get("value", 1)), float(v.get("value", 1))
    return default, default


def condition_effect(conds):
    """(chance factor, notes) of a list of loot conditions."""
    factor, notes = 1.0, []
    for c in conds or []:
        t = c.get("condition", "")
        if t.endswith("random_chance"):
            ch = c.get("chance", 1)
            factor *= ch if isinstance(ch, (int, float)) else number(ch)[1]
        elif t.endswith("random_chance_with_enchanted_bonus"):
            ch = c.get("unenchanted_chance", 0)
            factor *= ch if isinstance(ch, (int, float)) else 0.05
            notes.append("Looting raises it")
        elif t.endswith("random_chance_with_looting"):
            factor *= c.get("chance", 1)
        elif t.endswith("table_bonus"):
            factor *= (c.get("chances") or [1])[0]
        elif t.endswith("killed_by_player"):
            notes.append("player kill")
        elif t.endswith("match_tool"):
            pred = c.get("predicate", {})
            enchant = str(pred.get("predicates", {}).get("minecraft:enchantments", "")) + str(pred.get("enchantments", ""))
            if "silk_touch" in enchant:
                notes.append("Silk Touch")
            elif pred.get("items"):
                items = pred["items"] if isinstance(pred["items"], list) else [pred["items"]]
                notes.append("with " + ", ".join(i.split(":")[-1].replace("_", " ") for i in items))
            else:
                notes.append("specific tool")
        elif t.endswith("block_state_property"):
            props = c.get("properties", {})
            if "age" in props or "stage" in props:
                notes.append("when fully grown" if str(props.get("age", props.get("stage"))) not in ("0",) else "when young")
        elif t.endswith("inverted"):
            _, inner_n = condition_effect([c.get("term", {})])
            if "Silk Touch" in inner_n:
                notes.append("without Silk Touch")
        elif t.endswith("entity_properties") and "on_fire" in str(c):
            notes.append("cooked if on fire")
        elif t.endswith("any_of") or t.endswith("alternative"):
            pass
    return factor, notes


class Loot:
    def __init__(self, stack, item_tags):
        self.stack, self.item_tags = stack, item_tags
        self.ids = stack.ids("loot_table")
        self.cache = {}

    def table(self, tid):
        return self.stack.json(self.ids[tid]) if tid in self.ids else None

    def flatten(self, tid, _depth=0):
        """[{item, p, count:[lo,hi], notes}] for a loot table id or inline table."""
        if isinstance(tid, str):
            tid = norm_id(tid)
            if tid in self.cache:
                return self.cache[tid]
            data = self.table(tid)
        else:
            data = tid
        out = []
        if data and _depth < 8:
            for pool in data.get("pools", []):
                lo, hi = number(pool.get("rolls", 1))
                rolls = (lo + hi) / 2  # bonus_rolls only apply with Luck
                pf, pnotes = condition_effect(pool.get("conditions"))
                entries = list(self._entries(pool.get("entries", [])))
                total = sum(w for _, w in entries) or 1
                for e, w in entries:
                    ef, enotes = condition_effect(e.get("conditions"))
                    share = rolls * pf * ef * w / total
                    for d in self._entry_items(e, _depth):
                        out.append({"item": d["item"], "p": min(1.0, share * d["p"]),
                                    "count": d["count"], "notes": sorted(set(pnotes + enotes + d["notes"]))})
        if isinstance(tid, str):
            self.cache[tid] = out
        return out

    def _entries(self, entries):
        """(entry, weight) pairs; composite entries are opened up."""
        for e in entries:
            t = e.get("type", "")
            if t.endswith(("alternatives", "group", "sequence")):
                for child, w in self._entries(e.get("children", [])):
                    yield child, w
            else:
                yield e, e.get("weight", 1)

    def _entry_items(self, e, depth):
        t = e.get("type", "")
        lo, hi = 1.0, 1.0
        notes = []
        for f in e.get("functions", []):
            fn = f.get("function", "")
            if fn.endswith("set_count"):
                lo, hi = number(f.get("count"))
            elif fn.endswith(("furnace_smelt",)):
                notes.append("cooked if on fire")
        count = [lo, hi]
        if t.endswith(":item"):
            return [{"item": norm_id(e["name"]), "p": 1.0, "count": count, "notes": notes}]
        if t.endswith(":tag"):
            items = sorted(self.item_tags.resolve("#" + e["name"]))
            if e.get("expand"):
                return [{"item": i, "p": 1.0 / max(1, len(items)), "count": count, "notes": notes} for i in items]
            return [{"item": i, "p": 1.0, "count": count, "notes": notes + ["all items of a tag"]} for i in items]
        if t.endswith("loot_table"):
            ref = e.get("value", e.get("name"))
            nested = self.flatten(ref, depth + 1)
            return [dict(d, notes=d["notes"] + notes) for d in nested]
        return []