tabletop-frog/src/ttfrog/db/schema/inventory.py

109 lines
3.4 KiB
Python

from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import Mapped, mapped_column, relationship
from ttfrog.db.base import BaseObject, EnumField
from ttfrog.db.schema.item import Item, ItemType
class InventoryType(EnumField):
EQUIPMENT = "EQUIPMENT"
SPELL = "SPELL"
inventory_type_map = {
InventoryType.EQUIPMENT: [
ItemType.ITEM,
ItemType.SCROLL,
],
InventoryType.SPELL: [ItemType.SPELL],
}
def inventory_map_creator(fields):
if isinstance(fields, InventoryMap):
return fields
return InventoryMap(**fields)
class Inventory(BaseObject):
__tablename__ = "inventory"
__table_args__ = (UniqueConstraint("character_id", "inventory_type"),)
id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
character_id: Mapped[int] = mapped_column(ForeignKey("character.id"))
inventory_type: Mapped[InventoryType] = mapped_column(nullable=False)
_inventory_map = relationship("InventoryMap", lazy="immediate", uselist=True, cascade="all,delete,delete-orphan")
inventory_map = association_proxy("_inventory_map", "id", creator=inventory_map_creator)
def get(self, item):
return [mapping for mapping in self._inventory_map if mapping.item == item]
def add(self, item):
if item.item_type not in inventory_type_map[self.inventory_type]:
return False
mapping = InventoryMap(inventory_id=self.id, item_id=item.id)
if item.consumable:
mapping.count = item.count
self.inventory_map.append(mapping)
return mapping
def remove(self, mapping):
if mapping.id not in self.inventory_map:
return False
self.inventory_map.remove(mapping.id)
return True
def equip(self, mapping):
if mapping.equipped:
return False
mapping.equipped = True
return True
def unequip(self, mapping):
if not mapping.equipped:
return False
mapping.equipped = False
return True
def __contains__(self, obj):
for mapping in self._inventory_map:
if mapping.item == obj:
return True
return False
def __iter__(self):
yield from self._inventory_map
class InventoryMap(BaseObject):
__tablename__ = "inventory_map"
id: Mapped[int] = mapped_column(init=False, primary_key=True, autoincrement=True)
inventory_id: Mapped[int] = mapped_column(ForeignKey("inventory.id"))
item_id: Mapped[int] = mapped_column(ForeignKey("item.id"))
item: Mapped["Item"] = relationship(uselist=False, lazy="immediate", viewonly=True, init=False)
inventory: Mapped["Inventory"] = relationship(uselist=False, viewonly=True, init=False)
equipped: Mapped[bool] = mapped_column(default=False)
count: Mapped[int] = mapped_column(nullable=False, default=1)
always_prepared: Mapped[bool] = mapped_column(default=False)
@property
def prepared(self):
if self.item.item_type == ItemType.SPELL:
return self.equipped or self.always_prepared
def use(self, count=1):
if count < 0:
return False
if not self.item.consumable:
return False
if self.count < count:
return False
self.count -= count
if self.count == 0:
self.inventory.remove(self)
return 0
return self.count