wplace-export.py

import argparse
import asyncio
import io
import os
import re
from datetime import datetime

import aiohttp
from PIL import Image
from tqdm import tqdm

MAX_CONCURRENT_REQUESTS = 5
RATE_LIMIT_SLEEP_TIME = 5
TILE_URL = lambda x, y: "https://backend.wplace.live/files/s0/tiles/" + str(x) + "/" + str(y) + ".png"
TILE_W, TILE_H = 1000, 1000

parser = argparse.ArgumentParser(
    prog='WPlace Export Tool',
    description='Exports a designated area of wplace into an image')

parser.add_argument('start_x')
parser.add_argument('start_y')
parser.add_argument('end_x')
parser.add_argument('end_y')

parser.add_argument('-s', '--save-chunks',
                    action='store_true',
                    help="Also save an image of each chunk")
parser.add_argument('-w', '--watch',
                    metavar="SECONDS",
                    help="Exports the area every X seconds")
parser.add_argument('-f', '--fill',
                    metavar="COLOR (#RRGGBBAA)",
                    help="Specify color to use as background (Defaults to transparent)")

args = parser.parse_args()

def create_directories():
    if not os.path.exists("chunks"):
        os.makedirs("chunks")

    if not os.path.exists("exports"):
        os.makedirs("exports")

def parse_chunk_coordinates():
    try:
        return int(args.start_x), int(args.end_x), int(args.start_y), int(args.end_y)
    except ValueError:
        print("Invalid chunk coordinates (Use -h for help)")
        exit(1)

def create_and_save_final_image(tiles, fill):
    time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

    xs = [coord[0] for coord in tiles.keys()]
    ys = [coord[1] for coord in tiles.keys()]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)

    width = (max_x - min_x + 1) * TILE_W
    height = (max_y - min_y + 1) * TILE_H
    final_img = Image.new("RGBA", (width, height), fill)

    for (x, y), img in tiles.items():
        px = (x - min_x) * TILE_W
        py = (y - min_y) * TILE_H
        final_img.paste(img, (px, py))

    final_img.save(f"exports/export_{time}.png")

def parse_color(val):
    if not re.fullmatch(r"#(?:[0-9a-fA-F]{3,4}|[0-9a-fA-F]{6}|[0-9a-fA-F]{8})", val):
        raise ValueError(f"Invalid hex color format: {val}")

    hex_color = val.lstrip("#")

    if len(hex_color) in (3, 4):
        hex_color = "".join(c * 2 for c in hex_color)

    if len(hex_color) == 6:
        hex_color += "FF"

    r = int(hex_color[0:2], 16)
    g = int(hex_color[2:4], 16)
    b = int(hex_color[4:6], 16)
    a = int(hex_color[6:8], 16)

    return r, g, b, a

def parse_watch(val):
    try:
        return int(val)
    except ValueError:
        print("Invalid watch value (Use -h for help)")
        exit(1)

async def get_chunk(x, y, session, sem, fill, bar):
    time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

    retry = True

    async with sem:
        while retry:
            async with session.get(TILE_URL(x, y)) as response:
                if response.status == 404:
                    chunk_image = Image.new("RGBA", (TILE_W, TILE_H), fill)
                    retry = False
                elif response.status == 429:
                    tqdm.write(f"({x}, {y}) Rate limited. Waiting {RATE_LIMIT_SLEEP_TIME} seconds")
                    await asyncio.sleep(RATE_LIMIT_SLEEP_TIME)
                else:
                    image_data = await response.content.read()

                    img_rgba = Image.open(io.BytesIO(image_data)).convert("RGBA")

                    chunk_image = Image.new("RGBA", img_rgba.size, fill)
                    chunk_image.paste(img_rgba, mask=img_rgba.getchannel("A"))
                    retry = False

            if retry:
                continue

            if args.save_chunks is True:
                chunk_image.save(f"chunks/{x}_{y}_{time}.png")

            await asyncio.sleep(1)
            bar.update(1)
            return chunk_image

async def main():
    sem = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
    tiles = {}
    tasks = {}
    fill = parse_color(args.fill) if args.fill is not None else (0,0,0,0)
    watch = parse_watch(args.watch) if args.watch is not None else 0

    start_x, end_x, start_y, end_y = parse_chunk_coordinates()

    create_directories()

    total_chunks = (
            (max(start_x, end_x) - min(start_x, end_x) + 1) *
            (max(start_y, end_y) - min(start_y, end_y) + 1)
    )

    bar = tqdm(total=total_chunks)
    bar.set_description("Obtaining chunks")

    while True:
        async with aiohttp.ClientSession() as session:
            async with asyncio.TaskGroup() as group:
                for x in range(min(start_x, end_x), max(start_x, end_x) + 1):
                    for y in range(min(start_y, end_y), max(start_y, end_y) + 1):
                        tasks[(x, y)] = group.create_task(get_chunk(x, y, session, sem, fill, bar))

        bar.close()
        print("Done")
        print("Creating final image... ", end="")

        results = {coords: task.result() for coords, task in tasks.items()}
        for coords, value in results.items():
            tiles[coords] = value

        create_and_save_final_image(tiles, fill)
        print("Done")

        if watch == 0:
            break
        else:
            print(f"Waiting {watch} seconds")
            await asyncio.sleep(watch)


try:
    asyncio.run(main())
except KeyboardInterrupt:
    print("\nStopping")