bta-proxy/bta_proxy/entitydata.py

from typing import Any
from bta_proxy.datainputstream import AsyncDataInputStream, SyncDataInputStream
from enum import Enum
from dataclasses import dataclass

from bta_proxy.itemstack import ItemStack

from logging import getLogger

logger = getLogger(__name__)

class DataItemType(Enum):
    BYTE = 0
    SHORT = 1
    INTEGER = 2
    FLOAT = 3
    STRING = 4
    ITEMSTACK = 5
    CHUNK_COORDINATES = 6

@dataclass
class DataItem:
    type: DataItemType
    id: int
    value: Any

class EntityData:
    @classmethod
    async def read_from(cls, dis: AsyncDataInputStream) -> list[DataItem]:
        items = []
        while (data := await dis.read()) not in (127, 255):
            logger.debug(f"Read byte: {data} (type={(data & 0xE0) >> 5}, id={data & 0x1F})")
            item_type = DataItemType((data & 0xE0) >> 5)
            item_id: int = data & 0x1F
            match item_type:
                case DataItemType.BYTE:
                    items.append(DataItem(item_type, item_id, await dis.read()))
                case DataItemType.SHORT:
                    items.append(DataItem(item_type, item_id, await dis.read_short()))
                case DataItemType.FLOAT:
                    items.append(DataItem(item_type, item_id, await dis.read_float()))
                case DataItemType.STRING:
                    items.append(DataItem(item_type, item_id, await dis.read_string()))
                case DataItemType.ITEMSTACK:
                    items.append(DataItem(item_type, item_id, await ItemStack.read_from(dis)))
                case DataItemType.CHUNK_COORDINATES:
                    x = await dis.read_float()
                    y = await dis.read_float()
                    z = await dis.read_float()
                    items.append(DataItem(item_type, item_id, (x, y, z)))
        return items

    @classmethod
    def read_from_sync(cls, dis: SyncDataInputStream) -> list[DataItem]:
        items = []
        while (data := dis.read()) not in (127, 255):
            item_type = DataItemType((data & 0xE0) >> 5)
            item_id: int = data & 0x1F
            match item_type:
                case DataItemType.BYTE:
                    items.append(DataItem(item_type, item_id, dis.read()))
                case DataItemType.SHORT:
                    items.append(DataItem(item_type, item_id, dis.read_short()))
                case DataItemType.FLOAT:
                    items.append(DataItem(item_type, item_id, dis.read_float()))
                case DataItemType.STRING:
                    items.append(DataItem(item_type, item_id, dis.read_string()))
                case DataItemType.ITEMSTACK:
                    items.append(DataItem(item_type, item_id, ItemStack.read_from_sync(dis)))
                case DataItemType.CHUNK_COORDINATES:
                    x = dis.read_float()
                    y = dis.read_float()
                    z = dis.read_float()
                    items.append(DataItem(item_type, item_id, (x, y, z)))
        return items