From b1aad7b0f4a30e3566a144dfab923dd3eae793fa Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Sun, 31 Dec 2023 11:17:11 -0800
Subject: [PATCH] add ability to generate items of specific rarity

---
 dnd_item/sources/rarity.yaml |  3 +--
 dnd_item/types.py            | 32 ++++++++++++++++++++------------
 tests/test_types.py          | 26 ++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/dnd_item/sources/rarity.yaml b/dnd_item/sources/rarity.yaml
index bb85e86..7acbf44 100644
--- a/dnd_item/sources/rarity.yaml
+++ b/dnd_item/sources/rarity.yaml
@@ -39,8 +39,7 @@ metadata:
       uncommon: 0.05
       rare: 0.5
       very rare: 0.5
-      #legendary: 0.3
-      legendary: 0.0
+      legendary: 0.3
 common:
   - 0
 uncommon:
diff --git a/dnd_item/types.py b/dnd_item/types.py
index da5f0f9..a58a49e 100644
--- a/dnd_item/types.py
+++ b/dnd_item/types.py
@@ -280,15 +280,17 @@ class ItemGenerator:
             def __init__(self):
                 super().__init__(
                     bases=WeightedSet(
-                        (dict(name='{type} stick', type='wooden', ...), 0.3),
-                        (dict(name='{type} stick', type='lead', ...), 1.0),
-                        (dict(name='{type} stick', type='silver', ...), 0.5),
-                        (dict(name='{type} stick', type='glass', ...), 0.1),
+                        (dict(name='{type} stick', type='wooden'), 0.3),
+                        (dict(name='{type} stick', type='lead'), 1.0),
+                        (dict(name='{type} stick', type='silver'), 0.5),
+                        (dict(name='{type} stick', type='glass'), 0.1),
                     ),
-                    raritytypes.RARITY,
+                    rarity.types.RARITY,
                     properties_by_rarity=types.PROPERTIES_BY_RARITY,
                 )
 
+                # ...Add callbacks for string templates in PROPERTIES_BY_RARITY...
+
     Generating Random Items
 
     Given an ItemGenerator class, use the ItemGenerator.random() method to
@@ -300,7 +302,7 @@ class ItemGenerator:
 
     Example:
 
-        >>> stick = SharpStickGenerator().random(count=1, challenge_rating=17)
+        >>> stick = SharpStickGenerator().random(count=1, rarity='legendary')
         >>> stick[0].name
         Silver Stick
         >>> stick[0].rarity
@@ -369,7 +371,7 @@ class ItemGenerator:
 
         return set(getreqs(item))
 
-    def random_properties(self) -> dict:
+    def random_properties(self, rarity: str = '') -> dict:
         """
         Create a dictionary of item attributes appropriate for a given rarity.
         Dictionaries generated by this method are used as arguments to the
@@ -395,14 +397,17 @@ class ItemGenerator:
         This will result in:
 
                 >>> item['extras']['owner']
-                Jules Ultandottir
+                Jules Ultardottir
         """
 
         # select a random base
         item = self.bases.random()
 
         # select a random rarity
-        item["rarity"] = self.rarity.random()
+        if rarity:
+            item["rarity"] = self.rarity.source.as_dict()[rarity]
+        else:
+            item["rarity"] = self.rarity.random()
 
         # select a number of properties appropriate to the rarity
         num_properties = self._property_count_by_rarity(item["rarity"]["rarity"])
@@ -431,12 +436,12 @@ class ItemGenerator:
             try:
                 item[requirement] = getattr(self, f"get_{requirement}")(**item)
             except AttributeError:
-                logging.error("{item['name']} requires {self.__class__.__name__} to have a get_{requirement}() method.")
+                logging.error(f"{item['name']} requires {self.__class__.__name__} to define get_{requirement}().")
                 raise
 
         return item
 
-    def random(self, count: int = 1, challenge_rating: int = 0) -> list:
+    def random(self, count: int = 1, challenge_rating: int = 0, rarity: str = '') -> list:
         """
         Generate one or more random Item instances by selecting random values
         from the available data sources, appropriate to the specified challenge
@@ -446,6 +451,9 @@ class ItemGenerator:
         an encounter for an adventuring party of four. This will prevent
         lower-level encounters from generating legendary weapons and so on. If
         challenge_rating is 0, a rarity is chosen at random.
+
+        If a rarity is supplied, ensure we return items of that rarity; in this
+        case challenge_rating is ignored.
         """
         if challenge_rating in range(1, 5):
             frequency = "1-4"
@@ -461,7 +469,7 @@ class ItemGenerator:
 
         items = []
         for _ in range(count):
-            items.append(self.item_class.from_dict(self.random_properties()))
+            items.append(self.item_class.from_dict(self.random_properties(rarity=rarity)))
         return items
 
 
diff --git a/tests/test_types.py b/tests/test_types.py
index 7bb9cec..e68fc84 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -61,3 +61,29 @@ def test_Item_overrides():
     ten_foot_pole = types.Item.from_dict(attrs)
     assert ten_foot_pole.name == '7ft. Pole'
     assert ten_foot_pole.description == 'Broken. The end of this 10ft. pole has been snapped off.'
+
+
+def test_ItemGenerator_subclass():
+
+    class SharpStickGenerator(types.ItemGenerator):
+        def __init__(self):
+            super().__init__(
+                bases=types.WeightedSet(
+                    (dict(name='{type} stick', type='wooden'), 0.3),
+                    (dict(name='{type} stick', type='lead'), 1.0),
+                    (dict(name='{type} stick', type='silver'), 0.5),
+                    (dict(name='{type} stick', type='glass'), 0.1),
+                ),
+                rarity=types.RARITY,
+                properties_by_rarity=types.PROPERTIES_BY_RARITY,
+            )
+
+    stick = SharpStickGenerator().random(count=1, rarity='common')
+    assert stick[0].name in [
+        'wooden stick',
+        'lead stick',
+        'silver stick',
+        'glass stick',
+    ]
+    assert stick[0].rarity.rarity == 'common'
+    assert stick[0].description == ''