Add requirements to db write

d2warehouse/item.py

@@ -204,6 +204,37 @@ class Item:
         base = lookup_basetype(self.code)
         return base["width"], base["height"]
 
+    def requirements(self) -> dict:
+        base = lookup_basetype(self.code)
+        # TODO: How is class requirements determined on basetypes?
+        reqs = {
+            "lvl": base["req_lvl"],
+            "dex": base["req_dex"],
+            "str": base["req_str"],
+            "class": None,
+        }
+        # TODO: Can implicit mod end up increasing requirements?
+        if self.quality == Quality.UNIQUE:
+            reqs["lvl"] = lookup_unique(self.unique_id)["req_lvl"]
+        elif self.quality == Quality.SET:
+            reqs["lvl"] = lookup_set_item(self.set_id)["req_lvl"]
+        elif self.is_runeword:
+            # TODO: What affects runeword level? Only the sockets?
+            pass
+        elif self.quality in [Quality.MAGIC, Quality.RARE, Quality.CRAFTED]:
+            for m, id in [("suffixes", id) for id in self.suffixes] + [
+                ("prefixes", id) for id in self.prefixes
+            ]:
+                affix = lookup_affix(id, m)
+                reqs["lvl"] = max(reqs["lvl"], affix["req_lvl"])
+                if affix["req_class"]:
+                    reqs["class"] = affix["req_class"]
+        if self.sockets:
+            for socket in self.sockets:
+                socket_reqs = socket.requirements()
+                reqs["lvl"] = max(reqs["lvl"], socket_reqs["lvl"])
+        return reqs
+
     def write_to_db(self, socketed_into=None, commit=True) -> int:
         name = lookup_basetype(self.code)["name"]
         # FIXME: handle magic & rare names
@@ -218,13 +249,14 @@ class Item:
             lookup_set_item(self.set_id)["set"] if self.quality == Quality.SET else None
         )
 
+        req = self.requirements()
         db = get_db()
         cur = db.cursor()
         cur.execute(
             """INSERT INTO item (itembase_name, socketed_into, raw_data, raw_version,
             is_identified, is_socketed, is_beginner, is_simple, is_ethereal,
-            is_personalized, is_runeword, code)
+            is_personalized, is_runeword, code, req_lvl)
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
             (
                 lookup_basetype(self.code)["name"],
                 socketed_into,
@@ -238,38 +270,44 @@ class Item:
                 self.is_personalized,
                 self.is_runeword,
                 self.code,
+                req["lvl"],
             ),
         )
         item_id = cur.lastrowid
 
-        cur.execute(
+        if not self.is_simple:
-            """INSERT INTO item_extra (item_id, item_name,
+            cur.execute(
-            set_name, uid, lvl, quality, graphic, implicit, low_quality, set_id,
+                """INSERT INTO item_extra (item_id, item_name,
-            unique_id, nameword1, nameword2, runeword_id, personal_name,
+                set_name, uid, lvl, quality, graphic, implicit, low_quality, set_id,
-            defense, durability, max_durability, quantity)
+                unique_id, nameword1, nameword2, runeword_id, personal_name,
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                defense, durability, max_durability, quantity, req_str, req_dex,
-            (
+                req_class)
-                item_id,
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
-                name,
+                (
-                set_name,
+                    item_id,
-                self.uid,
+                    name,
-                self.lvl,
+                    set_name,
-                int(self.quality) if self.quality else None,
+                    self.uid,
-                self.graphic,
+                    self.lvl,
-                self.implicit,
+                    int(self.quality) if self.quality else None,
-                int(self.low_quality) if self.low_quality else None,
+                    self.graphic,
-                self.set_id,
+                    self.implicit,
-                self.unique_id,
+                    int(self.low_quality) if self.low_quality else None,
-                self.nameword1,
+                    self.set_id,
-                self.nameword2,
+                    self.unique_id,
-                self.runeword_id,
+                    self.nameword1,
-                self.personal_name,
+                    self.nameword2,
-                self.defense,
+                    self.runeword_id,
-                self.durability,
+                    self.personal_name,
-                self.max_durability,
+                    self.defense,
-                self.quantity,
+                    self.durability,
-            ),
+                    self.max_durability,
-        )
+                    self.quantity,
+                    req["str"],
+                    req["dex"],
+                    req["class"],
+                ),
+            )
 
         if self.quality in [Quality.MAGIC, Quality.RARE, Quality.CRAFTED]:
             for prefix, id in [(False, id) for id in self.suffixes] + [
@@ -313,7 +351,7 @@ class Item:
             uid, lvl, quality, graphic, implicit, low_quality, set_id, unique_id,
             nameword1, nameword2, runeword_id, personal_name, defense, durability,
             max_durability, quantity
-            FROM item INNER JOIN item_extra ON id = item_id WHERE id = ?""",
+            FROM item LEFT JOIN item_extra ON id = item_id WHERE id = ?""",
             (id,),
         ).fetchone()
         if row["raw_version"] != STASH_TAB_VERSION:
@@ -431,3 +469,11 @@ def lookup_runeword(id: int) -> dict:
         with open(os.path.join(_data_path, "runewords.json")) as f:
             _runeword_map = json.load(f)
     return _runeword_map[str(id)]
+
+
+def lookup_affix(id: int, prefix: bool) -> dict:
+    global _affix_map
+    if _affix_map is None:
+        with open(os.path.join(_data_path, "affixes.json")) as f:
+            _affix_map = json.load(f)
+    return _affix_map["prefixes" if prefix else "suffixes"][str(id)]

d2warehouse/schema.sql

@@ -14,6 +14,7 @@ CREATE TABLE item (
     --        items have two names: the base name and item_extra.item_name.
     itembase_name TEXT NOT NULL,
     socketed_into INTEGER DEFAULT NULL,
+    req_lvl INTEGER DEFAULT 0,
 
     -- The following fields match the fields of the item object in item.py
     raw_data BLOB NOT NULL,
@@ -38,6 +39,9 @@ CREATE TABLE item_extra (
     item_id INTEGER PRIMARY KEY,
     item_name TEXT DEFAULT NULL,
     set_name TEXT DEFAULT NULL,
+    req_str INTEGER DEFAULT 0,
+    req_dex INTEGER DEFAULT 0,
+    req_class TEXT DEFAULT NULL,
 
     -- The following fields match the fields of the item object in item.py
     uid INTEGER DEFAULT NULL,

d2warehouse/tests/test_db.py

@@ -2,7 +2,7 @@ import os
 import tempfile
 import unittest
 
-from d2warehouse.db import close_db, create_db, set_db_path
+from d2warehouse.db import close_db, create_db, get_db, set_db_path
 from d2warehouse.item import Item
 from d2warehouse.parser import parse_item
 
@@ -29,3 +29,14 @@ class DbTest(unittest.TestCase):
         db_id = item.write_to_db()
         loaded_item = Item.load_from_db(db_id)
         self.assertEqual(item, loaded_item)
+
+        # Check that requirement was written properly
+        db = get_db()
+        reqs = db.execute(
+            "SELECT req_lvl, req_str, req_dex, req_class FROM item JOIN item_extra ON id = item_id WHERE id = 1"
+        ).fetchone()
+        expected_reqs = item.requirements()
+        self.assertEqual(reqs["req_lvl"], expected_reqs["lvl"])
+        self.assertEqual(reqs["req_str"], expected_reqs["str"])
+        self.assertEqual(reqs["req_dex"], expected_reqs["dex"])
+        self.assertEqual(reqs["req_class"], expected_reqs["class"])