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,36 +68,38 @@ 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')
-    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]
-      if velocities is not None: 
-        data[:,3:6] = velocities[:,:3]
-      else:
-        data[:,3:6] = position[:,3:]
-      data[:,6] = 1
+  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]
+    if velocities is not None: 
+      data[:,3:6] = velocities[:,:3]
+    else:
+      data[:,3:6] = position[:,3:]
+    data[:,6] = 1
 
 
-    ret =  \
-       subprocess.run([
-           os.path.join(install_prefix,'bin','simple3DFilter'),
-           h5_file,
-           str(radius),
-           str(boxsize),
-           str(resolution),
-           str(cx), str(cy), str(cz)
-       ], cwd=tmpdir)
+  ret =  \
+     subprocess.run([
+         os.path.join(install_prefix,'bin','simple3DFilter'),
+         h5_file,
+         str(radius),
+         str(boxsize),
+         str(resolution),
+         str(cx), str(cy), str(cz)
+     ], cwd=tmpdir.name)
 
-    
-    f0 = h5.File(os.path.join(tmpdir,'fields.h5'), mode="r")
-    def cleanup_f0():
-      f0.close()
+  
+  f0 = h5.File(os.path.join(tmpdir.name,'fields.h5'), mode="r")
+  def cleanup_f0():
+    f0.close()
+    tmpdir.cleanup()
 
-    class Dict(dict):
-      pass
+  class Dict(dict):
+    pass
 
-    t = Dict(rho=f0['density'], p=[f0['p0'], f0['p1'], f0['p2']])
-    weakref.finalize(t, cleanup_f0)
-    return t
+  t = Dict(rho=f0['density'], p=[f0['p0'], f0['p1'], f0['p2']])
+  t._tmpdir_=tmpdir
+  weakref.finalize(t, cleanup_f0)
+  return t