hxapy_read_write.py

import array
import struct
from . import hxapy_header as hxa


# *** Read functions


def read_u8(f):
    return f.read(1)[0]


def read_u32(f):
    return struct.unpack("<I", f.read(4))[0]


def read_name(f):
    len = read_u8(f)
    return f.read(len).decode()


def read_array(f, typecode, length):
    arr = array.array(typecode)
    arr.fromfile(f, length)
    return arr


def read_array1(f, typecode, length):
    """Reads array, but if there's only one element, returns that instead."""
    arr = read_array(f, typecode, length)
    return arr if len(arr) != 1 else arr[0]


def read_meta(f):
    meta = {}
    meta["name"] = read_name(f)
    meta["type"] = hxa.HXAMetaDataType(read_u8(f))
    length = read_u32(f)

    if meta["type"] == hxa.HXAMetaDataType.INT64:
        data = read_array1(f, "Q", length)
    elif meta["type"] == hxa.HXAMetaDataType.DOUBLE:
        data = read_array1(f, "d", length)
    elif meta["type"] == hxa.HXAMetaDataType.NODE:
        data = read_array1(f, "I", length)
    elif meta["type"] == hxa.HXAMetaDataType.TEXT:
        data = f.read(length).decode()
    elif meta["type"] == hxa.HXAMetaDataType.BINARY:
        data = f.read(length)
    elif meta["type"] == hxa.HXAMetaDataType.META:
        data = [read_meta(f) for _ in range(length)]

    meta["data"] = data

    return meta


def read_layer(f, count):
    layer = {}
    layer["name"] = read_name(f)
    layer["components"] = read_u8(f)
    layer["type"] = hxa.HXALayerDataType(read_u8(f))

    length = count * layer["components"]

    if layer["type"] == hxa.HXALayerDataType.UINT8:
        data = read_array(f, "B", length)
    elif layer["type"] == hxa.HXALayerDataType.INT32:
        data = read_array(f, "i", length)
    elif layer["type"] == hxa.HXALayerDataType.FLOAT:
        data = read_array(f, "f", length)
    elif layer["type"] == hxa.HXALayerDataType.DOUBLE:
        data = read_array(f, "d", length)

    layer["data"] = data

    return layer


def read_layerstack(f, count):
    stack = {}
    num_layers = read_u32(f)
    stack["layers"] = [read_layer(f, count) for _ in range(num_layers)]
    return stack


def read_node(f):
    node = {}
    node["type"] = hxa.HXANodeType(read_u8(f))
    num_metas = read_u32(f)
    node["meta_data"] = [read_meta(f) for _ in range(num_metas)]

    content = {}

    if node["type"] == hxa.HXANodeType.GEOMETRY:
        content["vertex_count"] = read_u32(f)
        content["vertex_stack"] = read_layerstack(f, content["vertex_count"])

        content["edge_corner_count"] = read_u32(f)
        content["corner_stack"] = read_layerstack(f, content["edge_corner_count"])
        content["edge_stack"] = read_layerstack(f, content["edge_corner_count"])

        content["face_count"] = read_u32(f)
        content["face_stack"] = read_layerstack(f, content["face_count"])

    elif node["type"] == hxa.HXANodeType.IMAGE:
        content["type"] = hxa.HXAImageType(read_u8(f))
        width = read_u32(f)
        height = read_u32(f)
        depth = read_u32(f)
        content["resolution"] = [width, height, depth]
        content["image_stack"] = read_layerstack(f, width * height * depth)

    node["content"] = content

    return node


def read_hxa(f):
    magic = f.read(4)
    if magic != b"HxA\0":
        raise RuntimeError("not a HxA file (incorrect magic number)")

    hxa_dict = {}
    hxa_dict["version"] = read_u8(f)
    node_count = read_u32(f)
    hxa_dict["nodes"] = [read_node(f) for _ in range(node_count)]

    return hxa_dict


# *** Write functions


def ensure_array(v):
    """If v isn't already an array, makes it a one-element list."""
    try:
        _ = len(v)
        has_len = True
    except Exception:
        has_len = False

    is_array = not isinstance(v, dict) and has_len
    return v if is_array else [v]


def write_u8(f, v):
    f.write(struct.pack("<B", v))


def write_u32(f, v):
    f.write(struct.pack("<I", v))


def write_str(f, s):
    f.write(s.encode())


def write_name(f, name):
    assert len(name) <= 255
    write_u8(f, len(name))
    write_str(f, name)


def write_array(f, typecode, arr):
    if isinstance(arr, array.array) and arr.typecode == typecode:
        arr.tofile(f)
    else:
        fmt = f"<{len(arr)}{typecode}"
        f.write(struct.pack(fmt, *arr))


def write_meta(f, meta):
    mtype = meta["type"]
    data = ensure_array(meta["data"])

    write_name(f, meta["name"])
    write_u8(f, mtype)
    write_u32(f, len(data))

    if mtype == hxa.HXAMetaDataType.INT64:
        write_array(f, "Q", data)
    elif mtype == hxa.HXAMetaDataType.DOUBLE:
        write_array(f, "d", data)
    elif mtype == hxa.HXAMetaDataType.NODE:
        write_array(f, "I", data)
    elif mtype == hxa.HXAMetaDataType.TEXT:
        write_str(f, data)
    elif mtype == hxa.HXAMetaDataType.BINARY:
        f.write(data)
    elif mtype == hxa.HXAMetaDataType.META:
        for child_meta in data:
            write_meta(f, child_meta)
    else:
        assert False


def write_layer(f, layer):
    ltype = layer["type"]
    data = layer["data"]

    write_name(f, layer["name"])
    write_u8(f, layer["components"])
    write_u8(f, ltype)

    if ltype == hxa.HXALayerDataType.UINT8:
        write_array(f, "B", data)
    elif ltype == hxa.HXALayerDataType.INT32:
        write_array(f, "i", data)
    elif ltype == hxa.HXALayerDataType.FLOAT:
        write_array(f, "f", data)
    elif ltype == hxa.HXALayerDataType.DOUBLE:
        write_array(f, "d", data)
    else:
        assert False


def write_layerstack(f, stack):
    write_u32(f, len(stack["layers"]))

    for layer in stack["layers"]:
        write_layer(f, layer)


def write_node(f, node):
    write_u8(f, node["type"])
    write_u32(f, len(node["meta_data"]))

    for meta in node["meta_data"]:
        write_meta(f, meta)

    if node["type"] == hxa.HXANodeType.GEOMETRY:
        content = node["content"]

        write_u32(f, content["vertex_count"])
        write_layerstack(f, content["vertex_stack"])
        write_u32(f, content["edge_corner_count"])
        write_layerstack(f, content["corner_stack"])
        write_layerstack(f, content["edge_stack"])
        write_u32(f, content["face_count"])
        write_layerstack(f, content["face_stack"])

    elif node["type"] == hxa.HXANodeType.IMAGE:
        content = node["content"]

        write_u8(f, content["type"])
        write_u32(f, content["resolution"][0])
        write_u32(f, content["resolution"][1])
        write_u32(f, content["resolution"][2])
        write_layerstack(f, content["image_stack"])


def write_hxa(f, hxa_dict):
    f.write(b"HxA\0")
    write_u8(f, hxa_dict["version"])
    write_u32(f, len(hxa_dict["nodes"]))

    for node in hxa_dict["nodes"]:
        write_node(f, node)