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.SPELL, ], 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) inventory_map = association_proxy("_inventory_map", "id", creator=inventory_map_creator) def add(self, item): if item.item_type not in inventory_type_map[self.inventory_type]: return False self.inventory_map.append(InventoryMap(inventory_id=self.id, item_id=item.id)) def __iter__(self): yield from [mapping.item for mapping in 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) equipped: Mapped[bool] = mapped_column(default=False) count: Mapped[int] = mapped_column(nullable=False, default=1)