Files
ai-char-bot/repos/char_repo.py
2026-02-03 16:21:15 +03:00

44 lines
1.5 KiB
Python

from typing import List
from bson import ObjectId
from motor.motor_asyncio import AsyncIOMotorClient
from models.Character import Character
class CharacterRepo:
def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
self.collection = client[db_name]["characters"]
async def add_character(self, character: Character) -> Character:
op = await self.collection.insert_one(character.model_dump())
character.id = op.inserted_id
return character
async def get_character(self, character_id: str, with_image_data: bool = False) -> Character | None:
args = {}
if not with_image_data:
args["character_image_data"] = 0
res = await self.collection.find_one({"_id": ObjectId(character_id)}, args)
if res is None:
return None
else:
res["id"] = str(res.pop("_id"))
return Character(**res)
async def get_all_characters(self) -> List[Character]:
docs = await self.collection.find({}, {"character_image_data": 0}).to_list(None)
characters = []
for doc in docs:
# Конвертируем ObjectId в строку и кладем в поле id
doc["id"] = str(doc.pop("_id"))
# Создаем объект
characters.append(Character(**doc))
return characters
async def update_char(self, char_id: str, character: Character) -> None:
await self.collection.update_one({"_id": ObjectId(char_id)}, {"$set": character.model_dump()})