From f31661fb8096983a9b512648178735b3c2590bd2 Mon Sep 17 00:00:00 2001
From: bennahugo <bennahugo@gmail.com>
Date: Tue, 26 Mar 2024 18:10:17 +0200
Subject: [PATCH] Add Python 3.10 test and replace generator_jit with overload

---
 .github/workflows/ci.yml |  2 +-
 Jenkinsfile.sh           | 11 +++++++++++
 docker/python310.docker  | 22 ++++++++++++++++++++++
 docker/python36.docker   |  4 +---
 docker/python38.docker   |  4 +---
 setup.py                 |  4 ++--
 tricolour/flagging.py    | 16 ++++++++++++++--
 7 files changed, 52 insertions(+), 11 deletions(-)
 create mode 100644 docker/python310.docker

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index af2eab0..cced2e8 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -9,7 +9,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8"]
+        python-version: ["3.8", "3.9", "3.10"]
 
     steps:
       - name: Set up Python ${{ matrix.python-version }}
diff --git a/Jenkinsfile.sh b/Jenkinsfile.sh
index c72fb01..710acdc 100644
--- a/Jenkinsfile.sh
+++ b/Jenkinsfile.sh
@@ -12,6 +12,7 @@ TEST_DATA_DIR="$WORKSPACE/../../../test-data"
 # build
 docker build -f ${WORKSPACE_ROOT}/projects/tricolour/docker/python36.docker -t tricolour.1804.py36:${BUILD_NUMBER} ${WORKSPACE_ROOT}/projects/tricolour/
 docker build -f ${WORKSPACE_ROOT}/projects/tricolour/docker/python38.docker -t tricolour.2004.py38:${BUILD_NUMBER} ${WORKSPACE_ROOT}/projects/tricolour/
+docker build -f ${WORKSPACE_ROOT}/projects/tricolour/docker/python310.docker -t tricolour.2204.py310:${BUILD_NUMBER} ${WORKSPACE_ROOT}/projects/tricolour/
 
 #run tests
 tar xvf $TEST_DATA_DIR/acceptance_test_data.tar.gz -C $TEST_OUTPUT_DIR
@@ -38,4 +39,14 @@ docker run \
     --workdir /code \
     --entrypoint /bin/sh \
     tricolour.2004.py38:${BUILD_NUMBER} \
+    -c "python3 -m pytest --flake8 -s -vvv ."
+
+# test 3.10
+docker run \
+    --rm \
+    -v $TEST_OUTPUT_DIR:/testdata \
+    --env TRICOLOUR_TEST_MS=/testdata/$TEST_MS_REL \
+    --workdir /code \
+    --entrypoint /bin/sh \
+    tricolour.2204.py310:${BUILD_NUMBER} \
     -c "python3 -m pytest --flake8 -s -vvv ."
\ No newline at end of file
diff --git a/docker/python310.docker b/docker/python310.docker
new file mode 100644
index 0000000..00d6af6
--- /dev/null
+++ b/docker/python310.docker
@@ -0,0 +1,22 @@
+FROM ubuntu:22.04
+
+ENV DEBIAN_FRONTEND noninteractive
+ENV DEBIAN_PRIORITY critical
+
+# Install base requirements
+RUN apt update && apt install -y python3-pip \
+                              build-essential \
+                              cmake \
+                              casacore-dev \
+                              python3-numpy \
+                              python3-setuptools \
+                              libboost-python-dev \
+                              libcfitsio-dev \
+                              wcslib-dev
+
+RUN pip3 install --upgrade pip setuptools
+ADD . /code
+WORKDIR /code
+
+# Install base + testing packages
+RUN pip3 install .[testing]
\ No newline at end of file
diff --git a/docker/python36.docker b/docker/python36.docker
index 750d5db..899b5d5 100644
--- a/docker/python36.docker
+++ b/docker/python36.docker
@@ -16,6 +16,4 @@ ADD . /code
 WORKDIR /code
 
 # Install base + testing packages
-RUN pip3 install .[testing]
-
-RUN python3 -m pytest --flake8 -s -vvv .
\ No newline at end of file
+RUN pip3 install .[testing]
\ No newline at end of file
diff --git a/docker/python38.docker b/docker/python38.docker
index 7e104a9..d413cec 100644
--- a/docker/python38.docker
+++ b/docker/python38.docker
@@ -16,6 +16,4 @@ ADD . /code
 WORKDIR /code
 
 # Install base + testing packages
-RUN pip3 install .[testing]
-
-RUN python3 -m pytest --flake8 -s -vvv .
\ No newline at end of file
+RUN pip3 install .[testing]
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 5312bff..0bf88ee 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
 requirements = [
     'dask[array] >= 2021.2.0, <=2024.3.1',
     'donfig >= 0.4.0, <0.8.2',
-    'numpy >= 1.14.0, <= 1.22.0',  # breakage in newer numpy + numerical errors
+    'numpy >= 1.14.0, <= 1.22.0',
     'numba >= 0.43.0, <= 0.59.1',
     'scipy >= 1.2.0, <=1.12.0',
     'threadpoolctl >= 1.0.0, <=3.4.0',
@@ -21,7 +21,7 @@
 
 extras_require = {'testing': ['pytest',
                               'pytest-flake8',
-                              'flake8 >= 4.0.0, <5.0.0',
+                              "flake8 >= 4.0.0, <5.0.0",
                               'requests', 'gdown']}
 
 setup(
diff --git a/tricolour/flagging.py b/tricolour/flagging.py
index 04fb045..710ba72 100644
--- a/tricolour/flagging.py
+++ b/tricolour/flagging.py
@@ -26,7 +26,7 @@
 """  # noqa
 
 
-@numba.njit(nogil=True, cache=True)
+@numba.jit(nopython=True, nogil=True, cache=True)
 def flag_nans_and_zeros(vis_windows, flag_windows):
     """
     Flag nan and zero visibilities.
@@ -190,8 +190,20 @@ def _as_min_dtype(value):
     return np.array(value, dtype)
 
 
-@numba.generated_jit(nopython=True, nogil=True, cache=True)
+@numba.jit(nopython=True, nogil=True, cache=True)
 def _asbool(data):
+    return _asbool_impl(data)
+
+
+def _asbool_impl(data):
+    if data.dtype.itemsize == 1:
+        return data.view(np.bool_)
+    else:
+        return data.astype(np.bool_)
+
+
+@numba.extending.overload(_asbool_impl, nopython=True, nogil=True, cache=True)
+def _asbool_impl_jit(data):
     """Create a boolean array with the same values as `data`.
 
     The `data` contain only 0's and 1's. If possible, a view is returned,