txlyre 2 дней назад
Родитель
Сommit
36fe2d1507
2 измененных файлов с 119 добавлено и 17 удалено
  1. 117 16
      commands.py
  2. 2 1
      requirements.txt

+ 117 - 16
commands.py

@@ -16,6 +16,8 @@ from cairosvg import svg2png
 
 from PIL import Image
 
+import magic
+
 from actions import (
     find_action,
     create_action,
@@ -683,9 +685,9 @@ def chess_game_stats(chess, id):
 
 def chess_game_over(chess, id, e):
     return [
-      chess_render(chess, id, False),
-      f"Конец игры: {str(e)}",
-      chess_game_stats(chess, id)
+        chess_render(chess, id, False),
+        f"Конец игры: {str(e)}",
+        chess_game_stats(chess, id),
     ]
 
 
@@ -758,7 +760,7 @@ async def chess_move_handler(chess, id, move):
 
     draw = chess.draw(id)
     if draw:
-        reply.append(f"Ничья: {str(draw)}")   
+        reply.append(f"Ничья: {str(draw)}")
 
     return reply
 
@@ -789,6 +791,7 @@ async def chess_undo_handler(chess, id):
 
     return reply
 
+
 async def chess_skip_handler(chess, id):
     is_group = isinstance(id, str)
 
@@ -839,7 +842,7 @@ async def chess_pass_handler(chess, id):
     return reply
 
 
-async def chess_board_handler(chess, id):     
+async def chess_board_handler(chess, id):
     try:
         e = chess.game_over(id)
         if e is not None:
@@ -927,20 +930,34 @@ async def chess_anim_handler(chess, id, arg):
         return ["Ходов ещё не сделано."]
 
     frames = list(map(lambda f: Image.open(BytesIO(svg2png(f))), frames))
-    
+
     buffer = BytesIO()
     buffer.name = "board.gif"
-    frames[0].save(fp=buffer, format="GIF", append_images=frames[1:], save_all=True, duration=500, loop=0)
+    frames[0].save(
+        fp=buffer,
+        format="GIF",
+        append_images=frames[1:],
+        save_all=True,
+        duration=500,
+        loop=0,
+    )
     buffer.seek(0)
 
     return [buffer]
 
+
 async def chess_fen_handler(chess, id):
-  return [chess.fen(id)]
+    return [chess.fen(id)]
+
 
 CHESS_COMMANDS = {
     "start": (chess_start_handler, "Начать новую игру", 0),
-    "from": (chess_from_handler, "Начать новую игру с доской в указанном состоянии", 1, True),
+    "from": (
+        chess_from_handler,
+        "Начать новую игру с доской в указанном состоянии",
+        1,
+        True,
+    ),
     "end": (chess_stop_handler, "Завершить игру", 0),
     "move": (chess_move_handler, "Сделать ход", 1),
     "undo": (chess_undo_handler, "Отменить ход", 0),
@@ -949,12 +966,14 @@ CHESS_COMMANDS = {
     "board": (chess_board_handler, "Показать доску", 0),
     "moves": (chess_moves_handler, "Показать историю ходов и состояние игры", 0),
     "set": (chess_set_handler, "Установить состояние доски", 1, True),
-
     "create": (chess_start_handler, "Создать общую доску", 0),
-    "createfrom": (chess_from_handler, "Создать общую доску в указанном состоянии", 1, True),
-
+    "createfrom": (
+        chess_from_handler,
+        "Создать общую доску в указанном состоянии",
+        1,
+        True,
+    ),
     "fen": (chess_fen_handler, "Получить представление доски в FEN", 0),
-
     "anim": (chess_anim_handler, "Создать анимацию последних N ходов", 0, True),
 }
 
@@ -1019,7 +1038,16 @@ async def chess_handler(bot, event, command):
         args = command.args[1:]
 
     peer_id = str(get_peer_id(event.peer_id))
-    result = await subcommand[0](bot.chess, peer_id if name in ("create", "createfrom") or (peer_id in bot.chess.sessions and name not in ("start", "from")) else event.sender.id, *args)
+    result = await subcommand[0](
+        bot.chess,
+        (
+            peer_id
+            if name in ("create", "createfrom")
+            or (peer_id in bot.chess.sessions and name not in ("start", "from"))
+            else event.sender.id
+        ),
+        *args,
+    )
 
     for reply in result:
         if isinstance(reply, BytesIO):
@@ -1038,12 +1066,82 @@ async def chess_handler(bot, event, command):
         else:
             await event.reply(reply)
 
-    if name in ("move", "skip", "undo", "pass") and peer_id in bot.chess.sessions and bot.chess.game_over(peer_id) is None:
+    if (
+        name in ("move", "skip", "undo", "pass")
+        and peer_id in bot.chess.sessions
+        and bot.chess.game_over(peer_id) is None
+    ):
         await event.reply(f"Ход {'белых' if bot.chess.turn(peer_id) else 'чёрных'}.")
 
 
 _chess_handler = Handler(chess_handler, is_public=True)
 
+
+async def wandscript(image, program):
+    try:
+        async with ClientSession() as sess:
+            async with sess.post(
+                "https://wandscript.txlyre.website/api/wand",
+                data={"image": BytesIO(image), "program": program},
+            ) as resp:
+                return resp.status != 200, await resp.read()
+    except Exception:
+        return None, None
+
+    return None, None
+
+
+mime_table = {
+    "image/jpeg": ".jpg",
+    "image/png": ".png",
+    "image/webp": ".webp",
+    "image/gif": ".gif",
+}
+
+
+async def wand_handler(bot, event, command):
+    if not event.document and not event.photo:
+        await event.reply("Требуется приложить фото или документ для обработки.")
+
+        return
+
+    is_document = event.document is not None
+
+    image = await bot.download_media(event.photo or event.document, file=bytes)
+    program = command.args_string
+
+    if program:
+        err, image = await wandscript(image, program)
+
+        if err is None:
+            await event.reply(
+                "Произошла неизвестная ошибка при попытке обращения к API… :("
+            )
+
+            return
+
+        if err:
+            await event.reply(f"Ошибка API: {image.decode('UTF-8')}")
+
+            return
+
+    mime_type = magic.from_buffer(image, mime=True)
+    ext = mime_table.get(mime_type)
+    if ext:
+        image = BytesIO(image)
+        image.name = f"photo{ext}"
+
+    await bot.send_file(
+        event.peer_id,
+        file=image,
+        reply_to=event,
+        mime_type=mime_type,
+        force_document=is_document,
+    )
+
+
+_wand_handler = Handler(wand_handler, is_public=True)
+
 COMMANDS = {
     "newadmin": Handler(newadmin_handler, is_restricted=True),
     "deladmin": Handler(deladmin_handler, is_restricted=True),
@@ -1062,5 +1160,8 @@ COMMANDS = {
     "run": Handler(run_handler, is_public=True),
     "say": Handler(say_handler, is_public=True),
     "roll": Handler(roll_handler, is_public=True),
-    "chess": _chess_handler, "c": _chess_handler,
+    "chess": _chess_handler,
+    "c": _chess_handler,
+    "wand": _wand_handler,
+    "w": _wand_handler,
 }

+ 2 - 1
requirements.txt

@@ -8,4 +8,5 @@ emoji
 markovify
 ujson
 spacy
-chess
+chess
+python-magic