diff --git a/src/raster/r.vif/r.vif.py b/src/raster/r.vif/r.vif.py index cda46d309e..105c0e9617 100755 --- a/src/raster/r.vif/r.vif.py +++ b/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 .""" + 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):