Rearrange cleanup for NFS stale handles

python/cosmotool/smooth.py

@@ -11,7 +11,7 @@ def smooth_particle_density(
            velocities=None,
            radius=1e6,
            boxsize=None, 
-           resolution=128,i
+           resolution=128,
            center=None, tmpprefix=None ):
   """Use adaptive smoothing to produce density and momentum fields.
   The algorithm is originally described in [1].
@@ -68,8 +68,8 @@ def smooth_particle_density(
     raise ValueError("Need a boxsize")
 
   cx,cy,cz=center
-  with TemporaryDirectory(prefix=tmpprefix) as tmpdir:
-    h5_file = os.path.join(tmpdir, 'particles.h5')
+  tmpdir = TemporaryDirectory(prefix=tmpprefix) 
+  h5_file = os.path.join(tmpdir.name, 'particles.h5')
   with h5.File(h5_file, mode="w") as f:
     data = f.create_dataset('particles', shape=(position.shape[0],7), dtype=np.float32)
     data[:,:3] = position[:,:3]
@@ -88,16 +88,18 @@ def smooth_particle_density(
          str(boxsize),
          str(resolution),
          str(cx), str(cy), str(cz)
-       ], cwd=tmpdir)
+     ], cwd=tmpdir.name)
 
   
-    f0 = h5.File(os.path.join(tmpdir,'fields.h5'), mode="r")
+  f0 = h5.File(os.path.join(tmpdir.name,'fields.h5'), mode="r")
   def cleanup_f0():
     f0.close()
+    tmpdir.cleanup()
 
   class Dict(dict):
     pass
 
   t = Dict(rho=f0['density'], p=[f0['p0'], f0['p1'], f0['p2']])
+  t._tmpdir_=tmpdir
   weakref.finalize(t, cleanup_f0)
   return t