From 8c76b8b27ea09fb8c64e430d8b33aef57971e58a Mon Sep 17 00:00:00 2001
From: Chris Byrohl <9221545+cbyrohl@users.noreply.github.com>
Date: Wed, 21 Aug 2024 19:35:13 +0200
Subject: [PATCH] allow numpy 2.0 and fix python 3.12 support

* catch unit parse exception raise since py3.12
* numpy 2.0 deprec of np.prod
* unpin astropy, numpy; bump pint
---
 pyproject.toml                       | 14 ++++----------
 src/scida/customs/arepo/dataset.py   |  2 +-
 src/scida/interfaces/mixins/units.py |  8 ++++++++
 tests/test_simulation_units.py       |  2 +-
 4 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index b5ad0b78..377369d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,10 +7,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.9"
-numpy = [
-    {version = "^1.21", python = "<3.12"},
-    {version = "^1.26", python = ">=3.12"}
-]
+numpy = ">=1.21"
 # need dask>2024.4.1 for py>=3.11.9 (https://github.com/dask/dask/pull/11035)
 dask = [{extras = ["array", "dataframe", "distributed"], version = ">=2023", python = "<3.11.9"},
         {extras = ["array", "dataframe", "distributed"], version = ">=2024.4.1", python = ">=3.11.9"}
@@ -18,15 +15,12 @@ dask = [{extras = ["array", "dataframe", "distributed"], version = ">=2023", pyt
 distributed = ">=2023"
 h5py = "^3.7.0"
 zarr = "^v2.10.0"
-astropy = "^5.0"
-numba = [
-    {version = "^0.57", python = "<3.12"},
-    {version = "^0.59", python = ">=3.12"}
-]
+astropy = ">=5.0"
+numba = ">=0.57"
 pyyaml = ">=5.3.1"
 jupyter = "^1.0.0"
 tqdm = "^4.64.1"
-pint = "^0.22"
+pint = "^0.24"
 requests = "^2.31.0"
 
 
diff --git a/src/scida/customs/arepo/dataset.py b/src/scida/customs/arepo/dataset.py
index 80480fee..86f7850a 100644
--- a/src/scida/customs/arepo/dataset.py
+++ b/src/scida/customs/arepo/dataset.py
@@ -1531,7 +1531,7 @@ def map_group_operation(
     # TODO: very messy and inefficient routine. improve some time.
     # TODO: Set entry_bytes_out
     nbytes_dtype_out = 4  # TODO: hardcode 4 byte output dtype as estimate for now
-    entry_nbytes_out = nbytes_dtype_out * np.product(shape)
+    entry_nbytes_out = nbytes_dtype_out * np.prod(shape)
 
     # list_chunkedges refers to bounds of index intervals to be processed together
     # if idxlist is specified, then these indices do not have to refer to group indices.
diff --git a/src/scida/interfaces/mixins/units.py b/src/scida/interfaces/mixins/units.py
index 4cefdb92..136c78aa 100644
--- a/src/scida/interfaces/mixins/units.py
+++ b/src/scida/interfaces/mixins/units.py
@@ -4,6 +4,7 @@
 
 import contextlib
 import logging
+import tokenize
 from enum import Enum
 from typing import Optional, Union
 
@@ -55,6 +56,9 @@ def str_to_unit(
     unit = None
     try:
         unit = ureg(unitstr)
+    except tokenize.TokenError as e:  # cannot parse; raises exception since python 3.12
+        log.debug("Cannot parse unit string '%s'. Skipping." % unitstr)
+        raise e
     except pint.errors.UndefinedUnitError as e:
         log.debug(
             "Cannot parse unit string '%s' from metadata description. Skipping."
@@ -470,6 +474,10 @@ def add_units(container: FieldContainer, basepath: str):
                         )
                     except pint.errors.UndefinedUnitError:
                         self._unitstates[path] = UnitState.parse_error
+                    except (
+                        tokenize.TokenError
+                    ):  # cannot parse; raises exception since python 3.12
+                        self._unitstates[path] = UnitState.parse_error
                     except ValueError as e:
                         if str(e) != "Could not find units.":
                             raise e
diff --git a/tests/test_simulation_units.py b/tests/test_simulation_units.py
index 96d04c7e..51e352e5 100644
--- a/tests/test_simulation_units.py
+++ b/tests/test_simulation_units.py
@@ -16,7 +16,7 @@ def test_simulation_tng50(testdatapath):
     # Compare dark matter mass
     cosmology: FlatLambdaCDM = snp.cosmology
     ureg = snp.unitregistry
-    vol = np.product(snp.boxsize) / cosmology.h**3 * ureg.kpc**3
+    vol = np.prod(snp.boxsize) / cosmology.h**3 * ureg.kpc**3
     dens = (cosmology.Om0 - cosmology.Ob0) * cosmology.critical_density0
     dens = dens.value * ureg(dens.unit.to_string("ogip"))
     boxmass = vol * dens