dimcheck.py

dimcheck.py

from typing import Callable

from numpy import array
from numpy import ndarray


def dimension(*expected: int):
    def wrap_up(func: Callable) -> Callable:
        def wrapped(*args, **kwargs) -> ndarray:
            result: ndarray = func(*args, **kwargs)
            assert tuple(result.shape) == tuple(expected), "Mismatch dimensions"
            return result
        return wrapped
    return wrap_up


@dimension(2, 1)
def f():
    return array([[1], [1]])