grass.pygrass: GridModule clean up temporary mapsets when exception occurs

python/grass/pygrass/modules/grid/grid.py

@@ -7,6 +7,7 @@
     print_function,
     unicode_literals,
 )
+import contextlib
 import os
 import sys
 import multiprocessing as mltp
@@ -645,8 +646,19 @@ def run(self, patch=True, clean=True):
                       created by GridModule
         :type clean: bool
         """
+        with contextlib.ExitStack() as stack:
+            if clean:
+                stack.callback(self._clean)
+            self._actual_run(patch=patch)
+
+    def _actual_run(self, patch):
+        """Run the GRASS command
+
+        :param patch: set False if you does not want to patch the results
+        """
         self.module.flags.overwrite = True
         self.define_mapset_inputs()
+
         if self.debug:
             for wrk in self.get_works():
                 cmd_exe(wrk)
@@ -689,17 +701,18 @@ def run(self, patch=True, clean=True):
                     fil = open(os.path.join(dirpath, self.out_prefix + par.value), "w+")
                     fil.close()
 
-        if clean:
-            self.clean_location()
-            self.rm_tiles()
-            if self.n_mset:
-                gisdbase, location = os.path.split(self.move)
-                self.clean_location(Location(location, gisdbase))
-                # rm temporary gis_rc
-                os.remove(self.gisrc_dst)
-                self.gisrc_dst = None
-                sht.rmtree(os.path.join(self.move, "PERMANENT"))
-                sht.rmtree(os.path.join(self.move, self.mset.name))
+    def _clean(self):
+        """Cleanup temporary data"""
+        self.clean_location()
+        self.rm_tiles()
+        if self.n_mset:
+            gisdbase, location = os.path.split(self.move)
+            self.clean_location(Location(location, gisdbase))
+            # rm temporary gis_rc
+            os.remove(self.gisrc_dst)
+            self.gisrc_dst = None
+            sht.rmtree(os.path.join(self.move, "PERMANENT"))
+            sht.rmtree(os.path.join(self.move, self.mset.name))
 
     def patch(self):
         """Patch the final results."""

python/grass/pygrass/modules/tests/grass_pygrass_grid_test.py

@@ -128,15 +128,16 @@ def run_grid_module():
 
 
 @pytest.mark.parametrize("clean", [True, False])
-def test_cleans(tmp_path, clean):
+@pytest.mark.parametrize("surface", ["surface", "non_exist_surface"])
+def test_cleans(tmp_path, clean, surface):
     """Check that temporary mapsets are cleaned when appropriate"""
     location = "test"
     mapset_prefix = "abc"
     gs.core._create_location_xy(tmp_path, location)  # pylint: disable=protected-access
     with gs.setup.init(tmp_path / location):
         gs.run_command("g.region", s=0, n=50, w=0, e=50, res=1)
-        surface = "surface"
-        gs.run_command("r.surf.fractal", output=surface)
+        if surface == "surface":
+            gs.run_command("r.surf.fractal", output=surface)
 
         def run_grid_module():
             # modules/shortcuts calls get_commands which requires GISBASE.