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