r.vif: preserve active MASK when read_data() exits with an error

src/raster/r.vif/r.vif.py

@@ -121,6 +121,13 @@
     from cStringIO import StringIO
 import grass.script as gs
 
+try:
+    from grass.script import MaskManager
+
+    HAS_MASK_MANAGER = True
+except ImportError:
+    HAS_MASK_MANAGER = False
+
 
 # Functions
 def prRed(skk):
@@ -133,10 +140,47 @@ def prGreen(skk):
 
 CLEAN_RAST = []
 
+mask_backup_name = None
+
+
+def restore_mask_backup():
+    """Best-effort restore of MASK from the tracked backup.
+
+    If a MASK already exists, it is one r.vif created internally (the script
+    held control throughout, so the user could not have set one). Remove it
+    before renaming the backup back. On success, clears mask_backup_name.
+    On failure, leaves it set and emits a warning with the backup name so the
+    user can restore manually. Never raises.
+    """
+    global mask_backup_name
+    if mask_backup_name is None:
+        return
+    mapset = gs.gisenv()["MAPSET"]
+    found_backup = gs.find_file(name=mask_backup_name, element="cell", mapset=mapset)
+    if not found_backup["fullname"]:
+        mask_backup_name = None
+        return
+    found_mask = gs.find_file(name="MASK", element="cell", mapset=mapset)
+    if found_mask["fullname"]:
+        try:
+            gs.run_command("r.mask", flags="r", quiet=True)
+        except Exception:
+            pass  # best effort; the rename below will fail with a clear warning
+    try:
+        gs.run_command("g.rename", raster=[mask_backup_name, "MASK"], quiet=True)
+        mask_backup_name = None
+    except Exception as exc:
+        gs.warning(
+            _(
+                "Failed to restore MASK from backup <{name}>: {err}. "
+                "Restore manually with: g.rename raster={name},MASK"
+            ).format(name=mask_backup_name, err=exc)
+        )
+
 
 def cleanup():
-    """Remove temporary maps specified in the global list. In addition,
-    remove temporary files"""
+    """Restore the original MASK if needed, then remove temporary maps."""
+    restore_mask_backup()
     cleanrast = list(reversed(CLEAN_RAST))
     for rast in cleanrast:
         gs.run_command("g.remove", flags="f", type="all", name=rast, quiet=True)
@@ -165,49 +209,68 @@ def check_layer(envlay):
 
 def read_data(raster, n, flag_s, seed):
     """Read in the raster layers as a numpy array."""
+    global mask_backup_name
     gs.message("Reading in the data ...")
-    if n:
-        # Create mask random locations
-        new_mask = tmpname("rvif")
+
+    def read_sampled():
+        """Read the raster values at the currently unmasked cells."""
+        tmpcov = StringIO(
+            gs.read_command(
+                "r.stats", flags="1n", input=raster, quiet=True, separator="comma"
+            ).rstrip("\n")
+        )
+        return np.loadtxt(tmpcov, skiprows=0, delimiter=",")
+
+    def random_sample(out):
+        """Write a raster of n random sample points from raster[0] to <out>."""
+        kwargs = {"input": raster[0], "npoints": n, "raster": out, "quiet": True}
         if flag_s:
-            gs.run_command(
-                "r.random",
-                input=raster[0],
-                flags="s",
-                npoints=n,
-                raster=new_mask,
-                quiet=True,
-            )
+            kwargs["flags"] = "s"
         else:
-            gs.run_command(
-                "r.random",
-                input=raster[0],
-                seed=seed,
-                npoints=n,
-                raster=new_mask,
+            kwargs["seed"] = seed
+        gs.run_command("r.random", **kwargs)
+
+    if not n:
+        return read_sampled()
+
+    if HAS_MASK_MANAGER:
+        # GRASS 8.5+: draw the sample first so it respects any active user MASK
+        # (as on 8.4), then let MaskManager apply the sample as the mask for
+        # reading and restore the user's MASK automatically on exit, including
+        # on error. No manual backup or rename is needed.
+        sample = tmpname("rvif")
+        random_sample(sample)
+        with MaskManager() as manager:
+            gs.mapcalc(
+                "{m} = if(isnull({s}), null(), 1)".format(
+                    m=manager.mask_name, s=sample
+                ),
                 quiet=True,
             )
-        exist_mask = gs.find_file(
-            name="MASK", element="cell", mapset=gs.gisenv()["MAPSET"]
-        )
-        if exist_mask["fullname"]:
-            mask_backup = tmpname("rvifoldmask")
-            gs.run_command("g.rename", raster=["MASK", mask_backup], quiet=True)
-        gs.run_command("r.mask", raster=new_mask, quiet=True)
-
-    # Get the raster values at sample points
-    tmpcov = StringIO(
-        gs.read_command(
-            "r.stats", flags="1n", input=raster, quiet=True, separator="comma"
-        ).rstrip("\n")
-    )
-    p = np.loadtxt(tmpcov, skiprows=0, delimiter=",")
-
-    if n:
-        gs.run_command("r.mask", flags="r", quiet=True)
+            return read_sampled()
+
+    # GRASS 8.4: no MaskManager. Back up the user's MASK, apply the sample
+    # mask, and restore it in the finally block (with the atexit cleanup() as
+    # a backstop if the process is killed).
+    sample = tmpname("rvif")
+    random_sample(sample)
+    exist_mask = gs.find_file(name="MASK", element="cell", mapset=gs.gisenv()["MAPSET"])
+    if exist_mask["fullname"]:
+        # Track the backup in mask_backup_name (NOT CLEAN_RAST) so cleanup()
+        # can restore it. Set only after a successful rename.
+        backup = create_unique_name("rvifoldmask")
+        gs.run_command("g.rename", raster=["MASK", backup], quiet=True)
+        mask_backup_name = backup
+    try:
+        gs.run_command("r.mask", raster=sample, quiet=True)
+        return read_sampled()
+    finally:
+        try:
+            gs.run_command("r.mask", flags="r", quiet=True)
+        except Exception:
+            pass
         if exist_mask["fullname"]:
-            gs.run_command("g.rename", raster=[mask_backup, "MASK"], quiet=True)
-    return p
+            restore_mask_backup()
 
 
 def compute_vif(mapx, mapy):