UniformBlock.py

import functools
from typing import Type, Union, Sequence, Any, List, Dict

import moderngl as mgl
import numpy as np


class Uniform:
    def __init__(self, shape: Union[int, Sequence[int]] = 1,
                 dtype: Union[np.dtype, Type, str] = "f4"):
        self._array = np.empty(shape=shape, dtype=dtype)
        self._offset = 0

    @property
    def nbytes(self) -> int:
        return self._array.nbytes

    def tobytes(self) -> bytes:
        return self._array.tobytes()

    def __get__(self, instance, owner=None) -> Any:
        if instance is None:
            return self

        if len(self._array > 1):
            return self._array.tolist()
        elif len(self._array == 1):
            return self._array[0]
        else:
            raise RuntimeError("array has 0 length")

    def __set__(self, instance: "UniformBlock", value: Any):
        try:
            val_len = len(value)
        except TypeError:
            val_len = 1

        if val_len == len(self._array):
            self._array[:] = value
        else:
            raise ValueError("provided value does not have matching length")

        instance.buffer.write(self._array.tobytes(), offset=self._offset)


class UniformBlock:
    __buffer_size__ = 0

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

        offset = 0
        for name in dir(cls):
            attrib, = getattr(cls, name),
            if isinstance(attrib, Uniform):
                # move uniform to private member
                attrib._offset = offset
                offset += attrib.nbytes

        cls.__buffer_size__ = offset

    def __init__(self, ctx: mgl.Context):
        self.ctx = ctx
        self.buffer = ctx.buffer(reserve=self.__class__.__buffer_size__)

    def use(self):
        self.buffer.bind_to_uniform_block(0)




# =======================================================================================
# EXAMPLE


class MyUniformBlock(UniformBlock):
    color = Uniform(4, "i4")
    size = Uniform(2)  # defualts to f4
    width = Uniform()  # defaults to one f4


ctx = mgl.create_standalone_context()
ub = MyUniformBlock(ctx)
ub.width = 1
ub.size = 7, 8
ub.color = 1, 2, 3, 4

print(ub.width)
print(ub.size)
print(ub.color)
print(ub.buffer.read())

ub.use()
# vbo.render(...)