asdf: use asdf._compression to support asdf 3.1.0

abacusnbody/data/asdf.py

@@ -17,11 +17,6 @@
 from asdf.extension import Compressor, Extension
 
 
-def _monkey_patch(*args,**kwargs):
-    raise Exception("Please use abacusnbody.data.asdf.set_nthreads(nthreads)")
-
-asdf.compression.set_decompression_options = _monkey_patch
-
 def set_nthreads(nthreads):
     blosc.set_nthreads(nthreads)
 

abacusnbody/data/compaso_halo_catalog.py

@@ -20,17 +20,21 @@
 from pathlib import PurePath
 
 import asdf
-import asdf.compression
 import astropy.table
 import numba as nb
 import numpy as np
 from astropy.table import Table
 
+try:
+    import asdf._compression as asdf_compression
+except ImportError:
+    import asdf.compression as asdf_compression
+
 from . import asdf as _asdf
 from . import bitpacked
 
 try:
-    asdf.compression.validate('blsc')
+    asdf_compression.validate('blsc')
 except Exception as e:
     raise Exception("Abacus ASDF extension not properly loaded! Try reinstalling abacusutils, or updating ASDF: `pip install asdf>=2.8`") from e
 

abacusnbody/data/pipe_asdf.py

@@ -86,11 +86,15 @@
 from timeit import default_timer as timer
 
 import asdf
-import asdf.compression
 import numpy as np
 
 try:
-    asdf.compression.validate('blsc')
+    import asdf._compression as asdf_compression
+except ImportError:
+    import asdf.compression as asdf_compression
+
+try:
+    asdf_compression.validate('blsc')
 except Exception as e:
     raise Exception("Abacus ASDF extension not properly loaded! Try reinstalling abacusutils, or updating ASDF: `pip install asdf>=2.8`") from e
 

abacusnbody/data/read_abacus.py

@@ -71,9 +71,14 @@ def read_asdf(fn, load=None, colname=None, dtype=np.float32, verbose=True, **kwa
     '''
 
     import asdf
-    import asdf.compression
+
+    try:
+        import asdf._compression as asdf_compression
+    except ImportError:
+        import asdf.compression as asdf_compression
+
     try:
-        asdf.compression.validate('blsc')
+        asdf_compression.validate('blsc')
     except Exception as e:
         raise Exception("Abacus ASDF extension not properly loaded! \
                         Try reinstalling abacusutils: `pip install 'abacusutils>=1'`, \