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 Spell(typing.NamedTuple): name: str source: str description: str ability_check: typing.Set[Ability] duration: typing.List[Duration] # remaining: 'affectsCreatureType', 'alias', 'areaTags', 'basicRules2024', 'components', 'conditionImmune', 'conditionInflict', 'damageImmune', 'damageInflict', 'damageResist', 'damageVulnerable', 'entriesHigherLevel', 'hasFluffImages', 'level', 'meta', 'miscTags', 'page', 'range', 'savingThrow', 'scalingLevelDice', 'school', 'spellAttack', '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}") return cls( name=name, source=source, description=description, ability_check=ability_check, duration=duration, ) @classmethod def ability_check_from_json(cls, json_data) -> typing.Set[Ability]: if "abilityCheck" not in json_data: return set() return {Ability(c) for c in json_data["abilityCheck"]} @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")