22import datetime
33import faulthandler
44import os
5+ import shutil
56import tempfile
67import unittest
78
2021]
2122[tmpfile1 , tmpfile2 ] = tmpfiles
2223
24+ # Set up temporary directories
25+ tmpdirs = [
26+ tempfile .mkdtemp ("_test_quantization.zarr" , dir = os .getcwd ())
27+ for i in range (1 )
28+ ]
29+ [tmpdir ] = tmpdirs
30+
2331
2432def _remove_tmpfiles ():
2533 """Remove temporary files created during tests."""
@@ -29,6 +37,13 @@ def _remove_tmpfiles():
2937 except OSError :
3038 pass
3139
40+ for d in tmpdirs :
41+ try :
42+ shutil .rmtree (d )
43+ os .rmdir (d )
44+ except OSError :
45+ pass
46+
3247
3348atexit .register (_remove_tmpfiles )
3449
@@ -222,6 +237,24 @@ def test_quantization_copy(self):
222237 g = f .copy ()
223238 self .assertTrue (g .get_quantization ().equals (q ))
224239
240+ def test_quantization_backends (self ):
241+ """Test that quantization-on-write with different backends."""
242+ f = self .f1 .copy ()
243+ f .set_quantize_on_write (
244+ algorithm = "granular_bitround" , quantization_nsd = 8
245+ )
246+
247+ # Backends that allow quantisation-on-write
248+ for backend in ("netCDF4" ,):
249+ cf .write (f , tmpfile1 , netcdf_backend = backend )
250+
251+ # Backends that do not allow quantisation-on-write
252+ with self .assertRaises (NotImplementedError ):
253+ cf .write (f , tmpdir , fmt = "ZARR3" , netcdf_backend = "zarr" )
254+
255+ with self .assertRaises (NotImplementedError ):
256+ cf .write (f , tmpfile1 , netcdf_backend = "h5netcdf-h5py" )
257+
225258
226259if __name__ == "__main__" :
227260 print ("Run date:" , datetime .datetime .now ())
0 commit comments