Add basic functionality to add and remove items from stash tabs

d2warehouse/fileformat.py

@@ -0,0 +1,19 @@
+# Copyright 2023 <omicron.me@protonmail.com>
+# Copyright 2023 <andreasruden91@gmail.com>
+#
+# This file is part of d2warehouse.
+#
+# d2warehouse is free software: you can redistribute it and/or modify it under the
+# terms of the GNU General Public License as published by the Free Software
+# Foundation, either version 3 of the License, or (at your option) any later
+# version.
+#
+# d2warehouse is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# Mercator. If not, see <https://www.gnu.org/licenses/>.
+
+STASH_TAB_MAGIC = b"\x55\xAA\x55\xAA"
+ITEM_DATA_MAGIC = b"JM"

d2warehouse/item.py

@@ -19,6 +19,7 @@ import os
 import re
 from typing import Optional
 from bitarray import bitarray
+from bitarray.util import int2ba
 from dataclasses import dataclass
 from enum import Enum
 
@@ -181,9 +182,25 @@ class Item:
             bits = bitarray(endian="little")
             bits.frombytes(self.raw_data)
             print(" " * indent, txtbits(bits))
+            print(" " * indent, self.raw_data.hex())
         print("")
         print("")
 
+    def set_position(self, x: int, y: int) -> None:
+        if x < 0 or x > 15 or y < 0 or y > 15:
+            raise ValueError("position needs to be in range [0, 16)")
+        self.pos_x = x
+        self.pos_y = y
+        bits = bitarray(endian="little")
+        bits.frombytes(self.raw_data)
+        bits[42:46] = int2ba(x, 4, endian="little")
+        bits[46:50] = int2ba(y, 4, endian="little")
+        self.raw_data = bits.tobytes()
+
+    def size(self) -> tuple[int, int]:
+        base = lookup_basetype(self.code)
+        return base["width"], base["height"]
+
 
 def lookup_basetype(code: str) -> dict:
     global _basetype_map

d2warehouse/parser.py

@@ -27,9 +27,7 @@ from d2warehouse.item import (
     lookup_stat,
 )
 import d2warehouse.huffman as huffman
-
-STASH_TAB_MAGIC = b"\x55\xAA\x55\xAA"
-ITEM_DATA_MAGIC = b"JM"
+from d2warehouse.fileformat import STASH_TAB_MAGIC, ITEM_DATA_MAGIC
 
 
 class ParseError(RuntimeError):

d2warehouse/stash.py

@@ -15,15 +15,102 @@
 # You should have received a copy of the GNU General Public License along with
 # Mercator. If not, see <https://www.gnu.org/licenses/>.
 from d2warehouse.item import Item
+from d2warehouse.fileformat import STASH_TAB_MAGIC, ITEM_DATA_MAGIC
+from struct import pack
+
+
+class StashError(RuntimeError):
+    pass
+
+
+class StashFullError(StashError):
+    pass
 
 
 class Stash:
     def __init__(self) -> None:
         self.tabs: list[StashTab] = []
 
+    def raw(self) -> bytes:
+        return b"".join(tab.raw() for tab in self.tabs)
+
 
 class StashTab:
     def __init__(self) -> None:
         self.gold: int = 0
+        self.version: int = 99
         self.item_data: bytes = b""
         self.items: list[Item] = []
+
+    def raw(self) -> bytes:
+        """Get the computed raw representation of the stash"""
+        item_raw = b"".join(item.raw_data for item in self.items)
+        raw_length = len(item_raw) + 0x44
+        return (
+            STASH_TAB_MAGIC
+            + pack("I", 1)
+            + pack("I", self.version)
+            + pack("I", self.gold)
+            + pack("I", raw_length)
+            + b"\x00" * 44
+            + ITEM_DATA_MAGIC
+            + pack("H", len(self.items))
+            + item_raw
+        )
+
+    def remove(self, item: Item) -> None:
+        for idx, current_item in enumerate(self.items):
+            if item is current_item:
+                self.items.pop(idx)
+                return
+        raise StashError("Item not found in stash")
+
+    def add(self, item: Item) -> None:
+        width, height = item.size()
+        dest = self._find_destination(width, height)
+        if not dest:
+            raise StashFullError(
+                "Could not locate an open spot in the stash to add the item"
+            )
+        x, y = dest
+        item.set_position(x, y)
+        self.items.append(item)
+
+    @staticmethod
+    def _is_free_location(
+        slots: list[list[bool]], x: int, y: int, width: int, height: int
+    ) -> bool:
+        """Check if a given location can fit an item of a given width and height"""
+        if x + width > len(slots):
+            return False
+        if y + height > len(slots[0]):
+            return False
+        for i in range(width):
+            for j in range(height):
+                if slots[x + i][y + j]:
+                    return False
+        return True
+
+    def _find_destination(self, width: int, height: int) -> tuple[int, int] | None:
+        """Find a free destination to store an item of given width and height"""
+        slots = self._build_slot_matrix()
+        stash_width = len(slots)
+        stash_height = len(slots[0])
+        for y in range(stash_width):
+            for x in range(stash_height):
+                if self._is_free_location(slots, x, y, width, height):
+                    return x, y
+        return None
+
+    def _build_slot_matrix(self, rows: int = 10, cols: int = 10) -> list[list[bool]]:
+        """Build a matrix of filled and empty slots for the stash tab."""
+        slots = []
+        for i in range(cols):
+            slots.append([False] * rows)
+        for item in self.items:
+            x, y = item.pos_x, item.pos_y
+            width, height = item.size()
+            for n in range(width):
+                for m in range(height):
+                    slots[x + n][y + m] = True
+        return slots

d2warehouse/tests/test_stash.py

@@ -0,0 +1,86 @@
+import unittest
+from d2warehouse.parser import parse_stash, parse_item
+from d2warehouse.stash import StashFullError
+
+
+class StashTest(unittest.TestCase):
+    @staticmethod
+    def simple_stash():
+        # tab 0 contains a rune, a buckler and arrows
+        # tab 1 contains a gem
+        # tab 3 contains nothing
+        # All tabs contain some gold
+        data = bytes.fromhex(
+            "55aa55aa0100000063000000a40100007400000000000000000000000000"
+            "000000000000000000000000000000000000000000000000000000000000"
+            "000000004a4d03001000a0000500f47cfb0010008000050454a1083f57a9"
+            "190c01041818fc0710008000050cf46deec23225b20c4340ff3f55aa55aa"
+            "0100000063000000391b00004d0000000000000000000000000000000000"
+            "000000000000000000000000000000000000000000000000000000000000"
+            "4a4d01001000a000050474891555aa55aa01000000630000003905000044"
+            "000000000000000000000000000000000000000000000000000000000000"
+            "00000000000000000000000000000000004a4d0000"
+        )
+        stash = parse_stash(data)
+        return data, stash
+
+    def dupe(self, item):
+        data = item.raw_data
+        data, duped = parse_item(data)
+        self.assertEqual(data, b"")
+        return duped
+
+    def test_raw_equals(self):
+        data, stash = self.simple_stash()
+        self.assertEqual(len(stash.tabs), 3)
+        self.assertEqual([len(tab.items) for tab in stash.tabs], [3, 1, 0])
+        gold = tuple(tab.gold for tab in stash.tabs)
+        self.assertEqual(gold, (420, 6969, 1337))
+        self.assertEqual(data, stash.raw())
+
+    def test_remove(self):
+        data, stash = self.simple_stash()
+        rune, buckler, arrows = stash.tabs[0].items[:3]
+        gem = stash.tabs[1].items[0]
+
+        stash.tabs[0].remove(buckler)
+        self.assertEqual(stash.tabs[0].items, [rune, arrows])
+        stash.tabs[0].remove(rune)
+        self.assertEqual(stash.tabs[0].items, [arrows])
+        stash.tabs[0].remove(arrows)
+        self.assertEqual(stash.tabs[0].items, [])
+        stash.tabs[1].remove(gem)
+        self.assertEqual(stash.tabs[1].items, [])
+
+        new_stash = parse_stash(stash.raw())
+        self.assertEqual(len(new_stash.tabs[0].items), 0)
+        self.assertEqual(len(new_stash.tabs[1].items), 0)
+        self.assertEqual(len(new_stash.tabs[2].items), 0)
+
+    def test_add(self):
+        data, stash = self.simple_stash()
+        rune, buckler, arrows = stash.tabs[0].items[:3]
+        stash.tabs[0].remove(buckler)
+        stash.tabs[1].add(buckler)
+        self.assertEqual(buckler.pos_x, 2)
+        self.assertEqual(buckler.pos_y, 0)
+
+        new_stash = parse_stash(stash.raw())
+        self.assertEqual(len(new_stash.tabs[0].items), 2)
+        self.assertEqual(len(new_stash.tabs[1].items), 2)
+
+    def test_overflow(self):
+        data, stash = self.simple_stash()
+        rune, buckler, arrows = stash.tabs[0].items[:3]
+        bucklers = [self.dupe(buckler) for _ in range(25)]
+
+        for buckler in bucklers:
+            stash.tabs[2].add(buckler)
+
+        with self.assertRaises(StashFullError):
+            stash.tabs[2].add(self.dupe(buckler))
+
+        new_stash = parse_stash(stash.raw())
+        self.assertEqual(len(new_stash.tabs[0].items), 3)
+        self.assertEqual(len(new_stash.tabs[1].items), 1)
+        self.assertEqual(len(new_stash.tabs[2].items), 25)