dnd_5etools_utils/dnd5etools/db/spells.py

import enum
import json
import pathlib
import typing
import dnd5etools.db


class Ability(enum.StrEnum):
    Strength = "strength"
    Dexterity = "dexterity"
    Constitution = "constitution"
    Intelligence = "intelligence"
    Wisdom = "wisdom"
    Charisma = "charisma"


class DurationType(enum.StrEnum):
    Timed = "timed"
    Instant = "instant"
    Permanent = "permanent"
    Special = "special"


class DurationTimeType(enum.StrEnum):
    Round = "round"
    Minute = "minute"
    Hour = "hour"
    Day = "day"


class Duration(typing.NamedTuple):
    type: DurationType
    time_type: typing.Optional[DurationTimeType]
    time_value: typing.Optional[int]
    concentration: bool
    end_condition: typing.Set[str]

    @classmethod
    def from_json(cls, json_data) -> typing.Self:
        dt = DurationType(json_data["type"])
        return cls(
            type=dt,
            time_type=(
                DurationTimeType(json_data["duration"]["type"])
                if dt == DurationType.Timed
                else None
            ),
            time_value=(
                int(json_data["duration"]["amount"])
                if dt == DurationType.Timed
                else None
            ),
            concentration=json_data.get("concentration", False),
            end_condition=set(json_data.get("end", [])),
        )


class SpellAttackType(enum.StrEnum):
    Melee = "M"
    Ranged = "R"
    Other = "O"


class SpellAreaType(enum.StrEnum):
    SingleTarget = "ST"
    MultipleTargets = "MT"
    Cube = "C"
    Cone = "N"
    Cylinder = "Y"
    Sphere = "S"
    Circle = "R"
    Square = "Q"
    Line = "L"
    Hemisphere = "H"
    Wall = "W"


class Spell(typing.NamedTuple):
    name: str
    source: str
    description: str
    ability_check: typing.Set[Ability]
    duration: typing.List[Duration]
    attack_type: typing.Set[SpellAttackType]
    area_type: typing.Set[SpellAreaType]
    # remaining: 'affectsCreatureType', 'alias', 'basicRules2024', 'components', 'conditionImmune', 'conditionInflict', 'damageImmune', 'damageInflict', 'damageResist', 'damageVulnerable', 'entriesHigherLevel', 'hasFluffImages', 'level', 'meta', 'miscTags', 'page', 'range', 'savingThrow', 'scalingLevelDice', 'school', 'srd52', 'time'

    @classmethod
    def from_json(cls, json_data) -> typing.Self:
        name = json_data["name"]
        source = json_data["source"]
        try:
            ability_check = cls.ability_check_from_json(json_data)
        except Exception as e:
            raise Exception(f"Unable to parse abilityCheck for {name}: {e}")
        try:
            description = cls.description_from_json(json_data)
        except Exception as e:
            raise Exception(f"Unable to parse description for {name}: {e}")
        try:
            duration = cls.duration_from_json(json_data)
        except Exception as e:
            raise Exception(f"Unable to parse duration for {name}: {e}")
        try:
            attack_type = cls.attack_type_from_json(json_data)
        except Exception as e:
            raise Exception(f"Unable to parse spellAttack for {name}: {e}")
        try:
            area_type = cls.area_type_from_json(json_data)
        except Exception as e:
            raise Exception(f"Unable to parse areaTags for {name}: {e}")
        return cls(
            name=name,
            source=source,
            description=description,
            ability_check=ability_check,
            duration=duration,
            attack_type=attack_type,
            area_type=area_type,
        )

    @classmethod
    def ability_check_from_json(cls, json_data) -> typing.Set[Ability]:
        return cls.json_str_enum_list(json_data, "abilityCheck", Ability)

    @classmethod
    def attack_type_from_json(cls, json_data) -> typing.Set[SpellAttackType]:
        return cls.json_str_enum_list(json_data, "spellAttack", SpellAttackType)

    @classmethod
    def area_type_from_json(cls, json_data) -> typing.Set[SpellAreaType]:
        return cls.json_str_enum_list(json_data, "areaTags", SpellAreaType)

    @classmethod
    def json_str_enum_list(
        cls, json_data, key: str, enum_cls: typing.Type[enum.StrEnum]
    ) -> typing.Set[enum.StrEnum]:
        if key not in json_data:
            return set()
        return {enum_cls(c) for c in json_data[key]}

    @classmethod
    def description_from_json(cls, json_data) -> str:
        return " ".join([str(e) for e in json_data["entries"]])

    @classmethod
    def duration_from_json(cls, json_data) -> typing.List[Duration]:
        return [Duration.from_json(d) for d in json_data["duration"]]


class SpellList:
    def __init__(self, code: str, filepath: pathlib.Path) -> None:
        self.code = code
        self.filepath = filepath
        self.spells: typing.List[Spell] = self.parse_spell_list(self.filepath)

    @staticmethod
    def parse_spell_list(filepath: pathlib.Path) -> typing.List[Spell]:
        try:
            with filepath.open("rt") as spell_list_file:
                spell_list_data = json.load(spell_list_file)
        except Exception as e:
            raise Exception(f"Failed to read spells file {filepath}: {e}")
        spells: typing.List[Spell] = []
        dbg_duration_type = set()
        for spell_data in spell_list_data["spell"]:
            if "_copy" in spell_data:
                continue  # Ignore spells that are just tweaks or reprints
            try:
                spells.append(Spell.from_json(spell_data))
            except Exception as e:
                print(f"[WARN] Failed to parse spell in {filepath} -- {e}")
                continue
        return spells


class SpellsDb:
    def __init__(self, data_dir: pathlib.Path) -> None:
        self.db_index = dnd5etools.db.DbIndex(data_dir, "spells")
        self.spell_list_cache: typing.dict[str, SpellList] = {}

    def get_spell_list(self, code: str) -> SpellList:
        if code in self.spell_list_cache:
            return self.spell_list_cache[code]
        if code in self.db_index.source_index:
            spell_list = SpellList(code, self.db_index.source_index[code])
            self.spell_list_cache[code] = spell_list
            return spell_list
        raise Exception(f"No spell list matching code {code} found in index")