Added option to specify tmpdir in define_waterregion function

Author: doc78

.github/workflows/ci_env.yml

@@ -47,4 +47,4 @@ jobs:
       shell: bash -el {0}
       run: |
         pip install pytest pytest-cov
-        pytest tests/test_waterregions.py
+        pytest

src/lisfloodutilities/waterregions/define_waterregions.py

@@ -120,11 +120,12 @@ def parse_metadata(metadata_file):
             metadata = json.load(f)
     return metadata                            
 
-def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterregions_initial=None, output_wr=None, metadata_parsed=None):
-    
+def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterregions_initial=None, output_wr=None, metadata_parsed=None, tmpdir=None):
+    if tmpdir is None:
+        tmpdir = tempfile.gettempdir()
     #0. Check whether the input maps are in pcraster format, use nc2pcr if the condition is not satisfied
     if ldd[-3:]=='.nc':
-       ldd_pcr=tempfile.gettempdir() + '/ldd_pcr.map' 
+       ldd_pcr= tmpdir + '/ldd_pcr.map' 
        try:
           os.remove(ldd_pcr)
        except:
@@ -133,7 +134,7 @@ def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterreg
     else:
        ldd_pcr=ldd   
     if countries_id[-3:]=='.nc':
-       countries_id_pcr=tempfile.gettempdir() + '/countries_id_pcr.map' 
+       countries_id_pcr=tmpdir + '/countries_id_pcr.map' 
        try:
           os.remove(countries_id_pcr)
        except:
@@ -142,7 +143,7 @@ def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterreg
     else:
        countries_id_pcr=countries_id       
     if waterregions_initial[-3:]=='.nc':
-       waterregions_initial_pcr=tempfile.gettempdir() + '/waterregions_initial_pcr.map' 
+       waterregions_initial_pcr=tmpdir + '/waterregions_initial_pcr.map' 
        try:
           os.remove(waterregions_initial_pcr)
        except:
@@ -153,7 +154,7 @@ def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterreg
 
 
     #1. The calibration points are converted into a map
-    pointmap_file=tempfile.gettempdir() + '/points.map'
+    pointmap_file=tmpdir + '/points.map'
     try:
        os.remove(pointmap_file)
     except:
@@ -165,7 +166,10 @@ def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterreg
     ldd1 = pcr.readmap(ldd_pcr)
     points = pcr.readmap(pointmap_file)
     subcat1 = pcr.subcatchment(ldd1,points)
-   
+    try: 
+      os.remove(pointmap_file)
+    except:
+       pass   
     #3. Making map with all valid ldd cells
     land = pcr.scalar(subcat1)
     land = pcr.nominal(land >= 0.99999)
@@ -192,9 +196,9 @@ def define_waterregions(calib_points=None, countries_id=None, ldd=None, waterreg
     #9. Save the water region map
     if output_wr[-3:]==".nc":
         waterregion_nc=output_wr
-        output_wr=tempfile.gettempdir() + '/wr_pcr.map'
+        output_wr=tmpdir + '/wr_pcr.map'
     else:
-        waterregion_nc=tempfile.gettempdir() + '/wr_nc.nc'
+        waterregion_nc=tmpdir + '/wr_nc.nc'
         metadata_parsed=[]
 
     try:

tests/test_waterregions.py

@@ -17,7 +17,7 @@ def test_define_waterregions(self):
         output_wr = 'tests/data/waterregions/waterregions_output_test.map'
         reference_output = 'tests/data/waterregions/waterregions_reference_output.map'
 
-        define_waterregions(calib_points, countries_id, ldd, waterregion_init, output_wr)
+        define_waterregions(calib_points, countries_id, ldd, waterregion_init, output_wr, tmpdir='tests/data/waterregions/')
 
         out = PCRasterMap(output_wr)
         reference = PCRasterMap(reference_output)