pixelflut.py

import socket
import struct
import logging

from PIL import Image

class PixelClient():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    buffer = []
    log = logging.getLogger("pixelclient")

    def __init__(self, host: str, port=5004):
        self.host = host
        self.port = port

    def _send(self, data: bytes):
        """ Send data to the server"""
        self.sock.sendto(data, (self.host, self.port))

    def flush(self):
        """ Send everything that is inside the buffer to the server"""
        buffer_string = b"".join(self.buffer)
        data = self.SetVersionBit(1) + self.SetRGBAMode(0) + buffer_string
        self.log.debug(f"Sending {len(data)} bytes")
        self._send(data)
        self.buffer = []

    def raw_image(self, image: Image):
        """ Takes an PIL image and load the image data into an array"""
        pixels = image.load()
        width, height = image.size

        # Read all the pixels into a buffer ready to be send
        for x in range(width):
            for y in range(height):
                # TODO add support for transparency
                self.buffer.append(bytes(self.RGBPixel(x, y, *pixels[x, y])))
                print(len(self.buffer))

    @staticmethod
    def RGBPixel(x: int, y: int, r: int, g: int, b: int, a=None):
        """Generates the packed data for a pixel"""
        if a:
            return struct.pack("<2H4B", x, y, r, g, b, a)
        return struct.pack("<2H3B", x, y, r, g, b)

    @staticmethod
    def SetRGBAMode(mode: int) -> bytes:
        """
            Generate the rgb/rgba bit
            0 for RGB and 1 for RGBA
        """
        return struct.pack("<?", mode)

    @staticmethod
    def SetVersionBit(protocol=1) -> bytes:
        """Generate the Version bit"""
        return struct.pack("<B", protocol)