92 lines
3.3 KiB
Python

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 Spell(typing.NamedTuple):
name: str
source: str
description: str
ability_check: typing.Set[Ability]
# remaining: 'affectsCreatureType', 'alias', 'areaTags', 'basicRules2024', 'components', 'conditionImmune', 'conditionInflict', 'damageImmune', 'damageInflict', 'damageResist', 'damageVulnerable', 'duration', '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}")
return cls(
name=name,
source=source,
description=description,
ability_check=ability_check,
)
@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"]])
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] = []
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")