Merge pull request #22 from tassaron/limit-spell-lists

Limit spell lists

dnd_character/character.py

@@ -9,6 +9,7 @@
 from .SRD import SRD, SRD_class_levels
 from .equipment import _Item, Item
 from .experience import Experience, experience_at_level, level_at_experience
+from .spellcasting import SpellList
 from .dice import sum_rolls
 from .features import get_class_features_data
 
@@ -24,10 +25,9 @@ class InvalidParameterError(Exception):
 
 class Character:
     """
-    Character Object deals with all aspects of the player character including
-    name, age, gender, description, biography, level, wealth, and all
-    player ability scores.  All can be omitted to create a blank, level 1
-    player.
+    Character object deals with all aspects of a player character including
+    name, class features, level/experience, wealth, and all ability scores.
+    All can be omitted to create a blank, level 1 character.
     """
 
     def __init__(
@@ -107,7 +107,6 @@ def __init__(
                 intelligence (int):  character's starting intelligence
                 charisma     (int):  character's starting charisma
         """
-
         # Decorative attrs that don't affect program logic
         self.uid: UUID = (
             uuid4() if uid is None else uid if isinstance(uid, UUID) else UUID(uid)
@@ -174,15 +173,9 @@ def __init__(
         self.proficiencies = proficiencies if proficiencies is not None else {}
         self.saving_throws = saving_throws if saving_throws is not None else []
         self.spellcasting_stat = spellcasting_stat
-        self._cantrips_known: list["_SPELL"] = (
-            cantrips_known if cantrips_known is not None else []
-        )
-        self._spells_known: list["_SPELL"] = (
-            spells_known if spells_known is not None else []
-        )
-        self._spells_prepared: list["_SPELL"] = (
-            spells_prepared if spells_prepared is not None else []
-        )
+        self._cantrips_known: SpellList["_SPELL"] = SpellList(cantrips_known)
+        self._spells_known: SpellList["_SPELL"] = SpellList(spells_known)
+        self._spells_prepared: SpellList["_SPELL"] = SpellList(spells_prepared)
         self.set_spell_slots(spell_slots)
 
         # Experience points
@@ -416,28 +409,46 @@ def class_features_data(self):
         return self._class_features_data
 
     @property
-    def cantrips_known(self) -> list["_SPELL"]:
+    def cantrips_known(self) -> SpellList["_SPELL"]:
         return self._cantrips_known
 
     @cantrips_known.setter
     def cantrips_known(self, new_val) -> None:
-        self._cantrips_known = new_val
+        if len(new_val) > self._cantrips_known.maximum:
+            raise ValueError(
+                f"Too many spells in list (max {self._cantrips_known.maximum})"
+            )
+        self._cantrips_known = (
+            new_val if isinstance(new_val, SpellList) else SpellList(initial=new_val)
+        )
 
     @property
-    def spells_known(self) -> list["_SPELL"]:
+    def spells_known(self) -> SpellList["_SPELL"]:
         return self._spells_known
 
     @spells_known.setter
     def spells_known(self, new_val) -> None:
-        self._spells_known = new_val
+        if len(new_val) > self._spells_known.maximum:
+            raise ValueError(
+                f"Too many spells in list (max {self._spells_known.maximum})"
+            )
+        self._spells_known = (
+            new_val if isinstance(new_val, SpellList) else SpellList(initial=new_val)
+        )
 
     @property
-    def spells_prepared(self) -> list["_SPELL"]:
+    def spells_prepared(self) -> SpellList["_SPELL"]:
         return self._spells_prepared
 
     @spells_prepared.setter
     def spells_prepared(self, new_val) -> None:
-        self._spells_prepared = new_val
+        if len(new_val) > self._spells_prepared.maximum:
+            raise ValueError(
+                f"Too many spells in list (max {self._spells_prepared.maximum})"
+            )
+        self._spells_prepared = (
+            new_val if isinstance(new_val, SpellList) else SpellList(initial=new_val)
+        )
 
     @property
     def inventory(self) -> list[_Item]:
@@ -550,9 +561,11 @@ class name, hit dice, level progression data, proficiencies, saving throws,
             self.class_index = new_class.index
             self.hd = new_class.hit_die
             self._class_levels = SRD_class_levels[self.class_index]
+            # Set spellcasting stat to the full name of an ability score
+            ability = {"wis": "wisdom", "cha": "charisma", "int": "intelligence"}
             if new_class.spellcasting:
-                self.spellcasting_stat = new_class.spellcasting["spellcasting_ability"][
-                    "index"
+                self.spellcasting_stat = ability[
+                    new_class.spellcasting["spellcasting_ability"]["index"]
                 ]
             else:
                 self.spellcasting_stat = None
@@ -640,7 +653,7 @@ def apply_class_level(self) -> None:
         e.g., adds new class features, spell slots
         Called by `level.setter` and `classs.setter`
         """
-        if self.level > 20:
+        if not self._class_levels or self.level > 20:
             return
         for data in self._class_levels:
             if data["level"] > self.level:
@@ -651,12 +664,9 @@ def apply_class_level(self) -> None:
             self.prof_bonus = data.get("prof_bonus", self.prof_bonus)
             for feat in data["features"]:
                 self.class_features[feat["index"]] = SRD(feat["url"])
-            while len(self.class_features_enabled) < len(self.class_features):
-                self.class_features_enabled.append(True)
 
-            # Fetch new spell slots
-            spell_slots = data.get("spellcasting", self.spell_slots)
-            self.set_spell_slots(spell_slots)
+        while len(self.class_features_enabled) < len(self.class_features):
+            self.class_features_enabled.append(True)
 
         # During level up some class specific values change. example: rage damage bonus 2 -> 4
         # Class specific counters do not reset! example: available inspirations
@@ -674,10 +684,39 @@ def apply_class_level(self) -> None:
                     for k, v in new_cfd.items()
                 }
 
-    def set_spell_slots(self, new_spell_slots: dict[str, int]) -> dict[str, int]:
+        # Fetch new spell slots
+        spell_slots = (
+            self._class_levels[self.level - 1]
+            .get("spellcasting", self.spell_slots)
+            .copy()
+        )
+        spell_slots.pop("cantrips_known", None)
+        spell_slots.pop("spells_known", None)
+        self.update_spell_lists()
+        self.set_spell_slots(spell_slots)
+
+    def update_spell_lists(self) -> None:
+        """Set maximum of spells_known, spells_prepared, cantrips_known"""
+        spell_slots = (
+            self._class_levels[self.level - 1]
+            .get("spellcasting", self.spell_slots)
+            .copy()
+        )
+        # Set maximums of _spells_known and _cantrips_known
+        if "cantrips_known" in spell_slots:
+            self._cantrips_known.maximum = spell_slots.pop("cantrips_known")
+        if "spells_known" in spell_slots:
+            self._spells_known.maximum = spell_slots.pop("spells_known")
+        # Calculate new maximum of spells_prepared
+        if self.spellcasting_stat is not None:
+            self._spells_prepared.maximum = max(
+                self.get_ability_modifier(self.__dict__[self.spellcasting_stat])
+                + self.level,
+                1,
+            )
+
+    def set_spell_slots(self, new_spell_slots: Optional[dict[str, int]]) -> None:
         default_spell_slots = {
-            "cantrips_known": 0,
-            "spells_known": 0,
             "spell_slots_level_1": 0,
             "spell_slots_level_2": 0,
             "spell_slots_level_3": 0,

dnd_character/spellcasting.py

@@ -1,8 +1,10 @@
+import logging
 from typing import Union, Optional
 from dataclasses import dataclass, asdict
 from .SRD import SRD, SRD_endpoints, SRD_classes
 
 
+LOG = logging.getLogger(__package__)
 SRD_spells = {
     spell["index"]: SRD(spell["url"])
     for spell in SRD(SRD_endpoints["spells"])["results"]
@@ -50,6 +52,31 @@ def __iter__(self):
 }
 
 
+class SpellList(list):
+    """A list with a maximum size, for storing spells and cantrips"""
+
+    def __init__(self, initial: Optional[list["_SPELL"]]) -> None:
+        initial = initial if initial is not None else []
+        self._maximum: int = len(initial)
+        super().__init__(initial)
+
+    @property
+    def maximum(self) -> int:
+        return self._maximum
+
+    @maximum.setter
+    def maximum(self, new_val: int) -> None:
+        if len(self) > new_val:
+            LOG.error("Too many spells in spell list to lower its maximum.")
+            return
+        self._maximum = new_val
+
+    def append(self, new_val: "_SPELL") -> None:
+        if len(self) + 1 > self.maximum:
+            raise ValueError(f"Too many spells in list (max {self.maximum})")
+        super().append(new_val)
+
+
 spell_names_by_level = {
     i: [key for key, val in SRD_spells.items() if val["level"] == i] for i in range(10)
 }

tests/test_spellcasting.py

@@ -1,6 +1,7 @@
+import pytest
 from ast import literal_eval
 from dnd_character.character import Character
-from dnd_character.classes import Bard, Wizard, Ranger
+from dnd_character.classes import CLASSES, Bard, Wizard, Ranger
 from dnd_character.spellcasting import spells_for_class_level, SPELLS, _SPELL
 
 
@@ -26,8 +27,6 @@ def test_cantrips_wizard():
 
 def test_spell_slots_bard():
     assert Bard().spell_slots == {
-        "cantrips_known": 2,
-        "spells_known": 4,
         "spell_slots_level_1": 2,
         "spell_slots_level_2": 0,
         "spell_slots_level_3": 0,
@@ -43,8 +42,6 @@ def test_spell_slots_bard():
 def test_spell_slots_wizard():
     # wizard's class data is missing `spells_known`, which should receive the default of 0
     assert Wizard().spell_slots == {
-        "cantrips_known": 3,
-        "spells_known": 0,
         "spell_slots_level_1": 2,
         "spell_slots_level_2": 0,
         "spell_slots_level_3": 0,
@@ -60,8 +57,6 @@ def test_spell_slots_wizard():
 def test_spell_slots_ranger():
     # ranger's class data is missing `cantrips_known` and all `spell_slots` > 5
     assert Ranger().spell_slots == {
-        "cantrips_known": 0,
-        "spells_known": 0,
         "spell_slots_level_1": 0,
         "spell_slots_level_2": 0,
         "spell_slots_level_3": 0,
@@ -104,3 +99,50 @@ def test_spells_prepared_serializes_in_character():
         wiz.spells_prepared[0].index
         == literal_eval(str(dict(wiz)))["spells_prepared"][0]["index"]
     )
+
+
+@pytest.mark.parametrize(
+    ("class_index", "level", "expected"),
+    (("bard", 2, 5), ("warlock", 10, 10), ("ranger", 4, 3)),
+)
+def test_spells_known_maximum(class_index, level, expected):
+    c = Character(classs=CLASSES[class_index], level=level)
+    assert c.spells_known.maximum == expected
+
+
+def test_spells_prepared_maximum_cleric():
+    c = Character(classs=CLASSES["cleric"], wisdom=14)
+    assert c.spells_prepared.maximum == 3
+
+
+def test_spells_prepared_maximum_wizard():
+    c = Character(classs=CLASSES["wizard"], intelligence=14)
+    assert c.spells_prepared.maximum == 3
+
+
+def test_spells_exceeded_assignment():
+    c = Character(classs=CLASSES["bard"])
+    c.spells_known = [SPELLS["thunderwave"]] * 4
+    assert c.spells_known.maximum == 4
+    assert len(c.spells_known) == 4
+    with pytest.raises(ValueError, match=" spells "):
+        c.spells_known += [SPELLS["identify"]]
+
+
+def test_spells_exceeded_append():
+    c = Character(classs=CLASSES["bard"])
+    c.spells_known = [SPELLS["thunderwave"]] * 4
+    assert c.spells_known.maximum == 4
+    assert len(c.spells_known) == 4
+    with pytest.raises(ValueError, match=" spells "):
+        c.spells_known.append(SPELLS["identify"])
+
+
+def test_spells_exceeded_at_init():
+    c = Character(classs=CLASSES["bard"], spells_known=[SPELLS["thunderwave"]] * 4)
+    c_as_dict = dict(c)
+    # add illegal spell while character is a dictionary
+    c_as_dict["spells_known"].append(c_as_dict["spells_known"][0].copy())
+    new_c = Character(**c_as_dict)
+    # illegal spell increases maximum in this edge case - desireable? I think so
+    assert len(new_c.spells_known) == new_c.spells_known.maximum