mirror of
https://github.com/Xevion/v6-place.git
synced 2025-12-05 23:16:48 -06:00
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
import asyncio
|
|
import io
|
|
import os
|
|
from typing import Optional, List, Union
|
|
|
|
import websockets
|
|
from PIL import Image
|
|
|
|
from place.constants import Environment
|
|
from place.differencing import get_pixel_differences
|
|
from place.network import upload_pixels
|
|
from place.pixel_types import Pixel
|
|
|
|
width, height = int(os.getenv(Environment.CANVAS_HEIGHT)), int(os.getenv(Environment.CANVAS_HEIGHT))
|
|
total_pixels = width * height
|
|
minimum_tolerance = 5 / total_pixels
|
|
|
|
|
|
class PlaceClient:
|
|
"""
|
|
Defines a stateful client that manages a 'source of truth' with the image created by incremental changes to the Websocket.
|
|
"""
|
|
|
|
def __init__(self, connection) -> None:
|
|
self.connection: websockets.WebSocketClientProtocol = connection
|
|
|
|
# A lock used to manage the 'source of truth' image while performing read intensive operations.
|
|
self.source_lock = asyncio.Lock()
|
|
|
|
# The 'source of truth' image describing what is currently on the canvas.
|
|
self.source: Image = Image.new("RGBA", (width, height), (255, 0, 0, 0))
|
|
|
|
# The current targeted 'output' image.
|
|
self.current_target: Optional[Image] = None
|
|
|
|
def lock(self) -> asyncio.Lock:
|
|
return self.source_lock
|
|
|
|
async def get_differences(self) -> List[Pixel]:
|
|
"""
|
|
:return: A list of pixels that must be placed on the canvas to meet the currently set task.
|
|
"""
|
|
if self.current_target is not None:
|
|
async with self.lock():
|
|
return get_pixel_differences(self.source, self.current_target)
|
|
return []
|
|
|
|
async def complete(self, tolerance: Union[int, float] = 15, sleep: float = 0.25) -> None:
|
|
pixel_tolerance = tolerance if type(tolerance) == int else total_pixels * tolerance
|
|
|
|
if self.current_target is None:
|
|
return
|
|
|
|
pixels = await self.get_differences()
|
|
while len(pixels) > pixel_tolerance:
|
|
# Upload all the differences
|
|
upload_pixels(pixels)
|
|
|
|
# Wait a bit for the client to catch up. Is this super necessary?
|
|
await asyncio.sleep(sleep)
|
|
|
|
# Recalculate the difference
|
|
pixels = await self.get_differences()
|
|
|
|
@classmethod
|
|
async def connect(cls, address: str):
|
|
"""A factory method for connecting to the websocket and instantiating the client."""
|
|
connection = await websockets.connect(address)
|
|
client = cls(connection)
|
|
if connection.open:
|
|
message = await connection.recv()
|
|
if type(message) != bytes:
|
|
raise RuntimeError("Fatal: Initial message from websocket was not 'bytes'")
|
|
|
|
img = Image.open(io.BytesIO(message))
|
|
client.source.paste(img)
|
|
|
|
return client
|
|
|
|
async def receive(self):
|
|
"""
|
|
Receiving all server messages and handling them
|
|
"""
|
|
while True:
|
|
try:
|
|
message = await self.connection.recv()
|
|
if type(message) == bytes:
|
|
img = Image.open(io.BytesIO(message))
|
|
|
|
async with self.lock():
|
|
self.source.paste(img, (0, 0), img)
|
|
|
|
except websockets.ConnectionClosed:
|
|
print('Connection with server closed')
|
|
break
|