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] # TODO: replace with a full enum @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 SpellRangeType(enum.StrEnum): Special = "special" Point = "point" Line = "line" Cube = "cube" Cone = "cone" Emanation = "emanation" Radius = "radius" Sphere = "sphere" Hemisphere = "hemisphere" Cylinder = "cylinder" Self = "self" Sight = "sight" Unlimited = "unlimited" UnlimitedSamePlane = "plane" Touch = "touch" class SpellRange(typing.NamedTuple): type: SpellRangeType distance_type: str # TODO Replace with a full enum distance_value: typing.Optional[int] @classmethod def from_json(cls, json_data) -> typing.Self: return cls( type=SpellRangeType(json_data["type"]), distance_type=json_data["distance"]["type"], distance_value=( int(json_data["distance"]["amount"]) if "amount" in json_data["distance"] else None ), ) 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] range: SpellRange # remaining: 'affectsCreatureType', 'alias', 'basicRules2024', 'components', 'conditionImmune', 'conditionInflict', 'damageImmune', 'damageInflict', 'damageResist', 'damageVulnerable', 'entriesHigherLevel', 'hasFluffImages', 'level', 'meta', 'miscTags', 'page', '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}") try: range = cls.range_from_json(json_data) except Exception as e: raise Exception(f"Unable to parse range 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, range=range, ) @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"]] @classmethod def range_from_json(cls, json_data) -> typing.List[SpellRange]: return SpellRange.from_json(json_data["range"]) 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")