# Maintainer: Sven-Hendrik Haase <svenstaro@archlinux.org>
# Contributor: Daniel Bershatsky <bepshatsky@yandex.ru>

pkgname='python-ml-dtypes'
pkgver=0.3.2
pkgrel=1
pkgdesc='A stand-alone implementation of several NumPy dtype extensions used in ML'
arch=('x86_64' 'aarch64')
url='https://github.com/jax-ml/ml_dtypes'
license=('Apache')
depends=('python-numpy')
makedepends=('pybind11' 'python-absl' 'python-build' 'python-installer' 'python-pytest' 'python-setuptools' 'python-wheel')
source=("ml-dtypes-${pkgver}.tar.gz::https://github.com/jax-ml/ml_dtypes/archive/refs/tags/v${pkgver}.tar.gz"
        'ml-dtypes.diff'
        "eigen-7bf2968.tar.gz::https://gitlab.com/libeigen/eigen/-/archive/7bf2968fed5f246c0589e1111004cb420fcd7c71/eigen-7bf2968.tar.gz")
sha256sums=('029a866a745726954cd825f8e013363fa103d45ea7c4fb1c23dd0d6b8361ad5b'
            'a7b8074c063d17d178430f291f92c18bece68855af957018a6baefc2f54df296'
            'beca76e4663a75ef3ff00b0704009a7a60cec3783ce3cada0b828982da87130b')

prepare() {
    ln -fs \
        $srcdir/eigen-7bf2968fed5f246c0589e1111004cb420fcd7c71/Eigen \
        $srcdir/ml_dtypes-$pkgver/third_party/eigen/Eigen

    cd ml_dtypes-$pkgver
    patch -Np1 -i ../ml-dtypes.diff
}

build() {
    cd ml_dtypes-$pkgver
    python -m build --wheel --no-isolation
    python setup.py build_ext -i
}

check() {
    cd ml_dtypes-$pkgver
    PYTHONPATH=$(pwd) pytest
}

package() {
    python -m installer \
        --compile-bytecode 1 \
        --destdir $pkgdir \
        ml_dtypes-$pkgver/dist/ml_dtypes-$pkgver-*-*.whl

    # Remove artefacts of broken distribution.
    rm -rfv $pkgdir/usr/lib/python3.11/site-packages/third_party
}
