PGN exportHEADmaster

1 files changed, 48 insertions, 14 deletions

diff --git a/server.py b/server.py
index c7abb20..340a09a 100644
--- a/server.py
+++ b/server.py
@@ -1,20 +1,24 @@
 import asyncio
+import datetime
 import json
 import os
 import pathlib
 import typing
-import weakref
+from collections import namedtuple
 
 import aiohttp
 import aiohttp_jinja2
 import chess
 import chess.engine
+import chess.pgn
 import jinja2
 from aiohttp import web
 
+RunningGame = namedtuple('RunningGame', 'websocket pgn')
+
 app = web.Application()
-app['websockets'] = weakref.WeakSet()
-app['chessgames'] = weakref.WeakSet()
+app['games'] = set()
+app['game_count'] = 0
 
 basepath = pathlib.Path(__file__).parent.absolute()
 print(f"Loading templates from {basepath}")
@@ -29,8 +33,23 @@ async def index(request: web.Request):
 async def handle_socket(request: web.Request):
     ws = web.WebSocketResponse()
     await ws.prepare(request)
-    request.app['websockets'].add(ws)
+
+    pgn = chess.pgn.GameBuilder()
+    app['game_count'] += 1
+    chess.pgn.Game()
+    pgn.begin_game()
+    pgn.begin_headers()
+    pgn.visit_header("White", "Player")
+    pgn.visit_header("Black", "Giri, Anish")
+    pgn.visit_header("Date", datetime.date.today().strftime('%Y.%m.%d'))
+    pgn.visit_header("Round", "-")
+    pgn.visit_header("Site", "https://chess.nea.moe")
+    pgn.end_headers()
+    running_game = RunningGame(ws, pgn)
+
+    request.app['games'].add(running_game)
     engine: typing.Optional[chess.engine.UciProtocol] = None
+
     try:
         board = chess.Board()
         transport, engine = await chess.engine.popen_uci('stockfish')
@@ -42,6 +61,18 @@ async def handle_socket(request: web.Request):
 
         await send_to_user(dict(event="ready", player_color='white'))
 
+        async def check_outcome():
+            outcome = board.outcome(claim_draw=True)
+            if outcome:
+                pgn.visit_result(outcome.result())
+                pgn.visit_header("Termination", "normal")
+                pgn.end_game()
+                pgn_str = pgn.result().accept(chess.pgn.StringExporter())
+                await send_to_user(dict(event="pgn", pgn=pgn_str))
+                await send_to_user(dict(event="game_over", result=outcome.result()))
+
+            return bool(outcome)
+
         async for msg in ws:
             msg: aiohttp.WSMessage
             if msg.type == aiohttp.WSMsgType.TEXT:
@@ -57,6 +88,7 @@ async def handle_socket(request: web.Request):
                     if user_move not in board.legal_moves:
                         await send_to_user(dict(event="reject_move"))
                         continue
+                    pgn.visit_move(board, user_move)
                     board.push(user_move)
                     await send_to_user(dict(event="accept_move", lastmove=user_move.uci()))
                     candidates: typing.List[chess.engine.InfoDict] = await engine.analyse(
@@ -67,26 +99,24 @@ async def handle_socket(request: web.Request):
                         numscore = score.relative.score(mate_score=100000)
                         return abs(numscore)
 
-                    if board.is_game_over():
-                        await send_to_user(dict(event="game_over", result=board.result()))
+                    if await check_outcome():
                         break
 
                     most_drawy_move: chess.engine.InfoDict = min(candidates, key=appraise)
                     my_move: chess.Move = (most_drawy_move['pv'][0])
+                    pgn.visit_move(board, my_move)
                     board.push(my_move)
                     await send_to_user(dict(
                         event="computer_moved",
                         lastmove=my_move.uci(),
                     ))
-                    if board.is_game_over(claim_draw=True):
-                        await send_to_user(dict(event="game_over", result=board.result(claim_draw=True)))
+                    if await check_outcome():
                         break
-
     finally:
         if not ws.closed:
             await ws.close()
         print("Cleaning up websocket")
-        request.app['websockets'].discard(ws)
+        request.app['games'].discard(running_game)
         if engine is not None:
             asyncio.create_task(engine.quit())
     return ws
@@ -94,8 +124,12 @@ async def handle_socket(request: web.Request):
 
 async def on_shutdown(app):
     print("On shutdown called")
-    for ws in set(app['websockets']):
-        ws: web.WebSocketResponse
+    for game in set(app['games']):
+        game: RunningGame
+        ws: web.WebSocketResponse = game.websocket
+        game.pgn.end_game()
+        pgn_str = game.pgn.result().accept(chess.pgn.StringExporter())
+        await ws.send_json(dict(event="pgn", pgn=pgn_str))
         await ws.close(code=aiohttp.WSCloseCode.GOING_AWAY,
                        message=b'Server shutdown')
         print("Closing websocket")
@@ -104,8 +138,8 @@ async def on_shutdown(app):
 async def handle_status(request: web.Request):
     return web.json_response(dict(
         all_tasks=len(asyncio.all_tasks(asyncio.get_running_loop())),
-        websockets=len(request.app['websockets']),
-        games=len(request.app['chessgames']),
+        games=len(request.app['games']),
+        total_recent_games=app['game_count'],
     ))