Deprecation fix the 2nd: replace mktemp with TemporaryDirectory

CHANGES.rst

@@ -1,6 +1,8 @@
 0.6.1.dev (unreleased)
 ----------------------
 - Fix memory issue when calling statistics() on FITS cubes. #752
+- Replace deprecated and insecure `mktemp` with `TemporaryDirectory` for
+ ``save_to_tmp_dir`` and add custom deletor to ensure cleanup. #796
 
 0.6 (2021-09-30)
 ----------------

spectral_cube/dask_spectral_cube.py

@@ -42,8 +42,8 @@
  SCIPY_INSTALLED = False
 
 try:
- import zarr
- import fsspec
+ import zarr # noqa: F401
+ import fsspec # noqa: F401
 except ImportError:
  ZARR_INSTALLED = False
 else:
@@ -80,10 +80,10 @@ def wrapper(self, *args, **kwargs):
  raise ImportError("saving the cube to a temporary directory "
  "requires the zarr and fsspec packages to "
  "be installed.")
- filename = tempfile.mktemp()
+ cube._tmpdir = tempfile.TemporaryDirectory()
  with dask.config.set(**cube._scheduler_kwargs):
- cube._data.to_zarr(filename)
- cube._data = da.from_zarr(filename)
+ cube._data.to_zarr(cube._tmpdir.name)
+ cube._data = da.from_zarr(cube._tmpdir.name)
  return cube
 
  return wrapper
@@ -238,6 +238,11 @@ def _new_cube_with(self, *args, **kwargs):
  new_cube._scheduler_kwargs = self._scheduler_kwargs
  return new_cube
 
+ def __del__(self):
+ # Clean up any `tempfile.TemporaryDirectory` created by `save_to_tmp_dir`.
+ if hasattr(self, '_tmpdir'):
+ self._tmpdir.cleanup()
+
  @property
  def _data(self):
  return self.__data

spectral_cube/tests/test_dask.py

@@ -80,6 +80,19 @@ def test_save_to_tmp_dir(data_adv):
  assert cube_new._data.name.startswith('from-zarr')
 
 
+def test_cleanup_tmp_dir(data_adv):
+ '''Test that temporary directory is removed after deleting cube.'''
+ pytest.importorskip('zarr')
+ cube = DaskSpectralCube.read(data_adv)
+ cube_new = cube.sigma_clip_spectrally(3, save_to_tmp_dir=True)
+ tmp_dir = cube_new._tmpdir.name
+ assert os.path.exists(tmp_dir)
+ assert '.zarray' in os.listdir(tmp_dir)
+
+ del cube_new
+ assert not os.path.exists(tmp_dir)
+
+
 def test_rechunk(data_adv):
  cube = DaskSpectralCube.read(data_adv)
  assert cube._data.chunksize == (4, 3, 2)