grass.script: Allow init to modify only specified environment (OSGeo#…

…3438)

The _grass.script.setup.init_ function modifies os.environ. While os.environ is fit for many, if not most, cases, usage in tests and parallel processing is limited and in all cases, os.environ stays partially modified even after session is finished (we don't destroy the runtime environment, i.e., variables such as GISBASE). With this change, _init_ takes _env_ parameter specifying the environment to modify and it modifies that environment. When no _env_ is provided, it still modifies os.environ, so the default behavior is not changed.

This required only few changes to the initialization code, but more changes were needed for the cleanup code. A lot of tests can now take advantage of this functionality and this PR updates some of them. I plan to update all where it is applicable, but will leave as is some others, namely those which use ctypes (situation is more complex there) and _grass.jupyter_ (_env_ parameter is not implemented for most of them).

There is plenty of tests covering it, but the cleanup part which needed most changes does not have any coverage. I don't think there is much concern in terms of API because _env_ parameters is what we now have in many functions. The only uncertainty with that is whether it should create its own copy or modify the provided environment. I go with modify because that makes the copy explicit in the call which is more clear (caller or reader is sure a copy is created) and it is consistent with how os.environ is modified.

python/grass/experimental/tests/conftest.py

@@ -1,11 +1,11 @@
 """Fixtures for grass.script"""
 
 import uuid
+import os
 
 import pytest
 
 import grass.script as gs
-import grass.script.setup as grass_setup
 import grass.experimental as experimental
 
 
@@ -14,7 +14,7 @@ def xy_session(tmp_path):
     """Active session in an XY location (scope: function)"""
     location = "xy_test"
     gs.core._create_location_xy(tmp_path, location)  # pylint: disable=protected-access
-    with grass_setup.init(tmp_path / location) as session:
+    with gs.setup.init(tmp_path / location, env=os.environ.copy()) as session:
         yield session
 
 
@@ -29,7 +29,7 @@ def xy_session_for_module(tmp_path_factory):
     tmp_path = tmp_path_factory.mktemp("xy_session_for_module")
     location = "xy_test"
     gs.core._create_location_xy(tmp_path, location)  # pylint: disable=protected-access
-    with grass_setup.init(tmp_path / location) as session:
+    with gs.setup.init(tmp_path / location, env=os.environ.copy()) as session:
         yield session
 
 

python/grass/experimental/tests/grass_script_mapset_session_test.py

@@ -59,7 +59,7 @@ def test_create_overwrite(xy_session):
     """Session creates and creates again with overwrite"""
     name = "test_mapset_1"
     session_file = xy_session.env["GISRC"]
-    with experimental.MapsetSession(name, create=True) as session:
+    with experimental.MapsetSession(name, create=True, env=xy_session.env) as session:
         session_mapset = gs.read_command("g.mapset", flags="p", env=session.env).strip()
         assert name == session_mapset
         gs.run_command("r.mapcalc", expression="a = 1", env=session.env)
@@ -69,7 +69,9 @@ def test_create_overwrite(xy_session):
             .split()
         )
         assert len(rasters) == 1 and rasters[0] == "a"
-    with experimental.MapsetSession(name, create=True, overwrite=True) as session:
+    with experimental.MapsetSession(
+        name, create=True, overwrite=True, env=xy_session.env
+    ) as session:
         session_mapset = gs.read_command("g.mapset", flags="p", env=session.env).strip()
         assert name == session_mapset
         rasters = (
@@ -92,7 +94,7 @@ def test_ensure(xy_session):
     """Session ensures and does not delete"""
     name = "test_mapset_1"
     session_file = xy_session.env["GISRC"]
-    with experimental.MapsetSession(name, ensure=True) as session:
+    with experimental.MapsetSession(name, ensure=True, env=xy_session.env) as session:
         session_mapset = gs.read_command("g.mapset", flags="p", env=session.env).strip()
         assert name == session_mapset
         gs.run_command("r.mapcalc", expression="a = 1", env=session.env)
@@ -102,7 +104,7 @@ def test_ensure(xy_session):
             .split()
         )
         assert len(rasters) == 1 and rasters[0] == "a"
-    with experimental.MapsetSession(name, ensure=True) as session:
+    with experimental.MapsetSession(name, ensure=True, env=xy_session.env) as session:
         session_mapset = gs.read_command("g.mapset", flags="p", env=session.env).strip()
         assert name == session_mapset
         rasters = (

python/grass/script/setup.py

@@ -268,7 +268,7 @@ def setup_runtime_env(gisbase=None, *, env=None):
     env["PYTHONPATH"] = path
 
 
-def init(path, location=None, mapset=None, *, grass_path=None):
+def init(path, location=None, mapset=None, *, grass_path=None, env=None):
     """Initialize system variables to run GRASS modules
 
     This function is for running GRASS GIS without starting it with the
@@ -320,6 +320,7 @@ def init(path, location=None, mapset=None, *, grass_path=None):
 
     :returns: reference to a session handle object which is a context manager
     """
+    # The path heuristic always uses the global environment.
     grass_path = get_install_path(grass_path)
     if not grass_path:
         raise ValueError(
@@ -351,15 +352,18 @@ def init(path, location=None, mapset=None, *, grass_path=None):
             )
         )
 
-    setup_runtime_env(grass_path)
+    # If environment is not provided, use the global one.
+    if not env:
+        env = os.environ
+    setup_runtime_env(grass_path, env=env)
 
     # TODO: lock the mapset?
-    os.environ["GIS_LOCK"] = str(os.getpid())
+    env["GIS_LOCK"] = str(os.getpid())
 
-    os.environ["GISRC"] = write_gisrc(
+    env["GISRC"] = write_gisrc(
         mapset_path.directory, mapset_path.location, mapset_path.mapset
     )
-    return SessionHandle()
+    return SessionHandle(env=env)
 
 
 class SessionHandle:
@@ -390,9 +394,27 @@ class SessionHandle:
         with gs.setup.init("~/grassdata/nc_spm_08/user1"):
             # ... use GRASS modules here
         # session ends automatically here
+
+    The example above is modifying the global, process environment (`os.environ`).
+    If you don't want to modify the global environment, use the _env_ parameter
+    for the _init_ function to modify the provided environment instead.
+    This environment is then available as an attribute of the session object.
+    The attribute then needs to be passed to all calls of GRASS
+    tools and functions that wrap them.
+    Context manager usage with custom environment::
+
+        # ... setup sys.path before import as needed
+
+        import grass.script as gs
+
+        with gs.setup.init("~/grassdata/nc_spm_08/user1", env=os.environ.copy()):
+            # ... use GRASS modules here with env parameter
+            gs.run_command("g.region", flags="p", env=session.env)
+        # session ends automatically here, global environment was never modifed
     """
 
-    def __init__(self, active=True):
+    def __init__(self, *, env, active=True):
+        self._env = env
         self._active = active
         self._start_time = datetime.datetime.now(datetime.timezone.utc)
 
@@ -403,7 +425,7 @@ def active(self):
 
     @property
     def env(self):
-        return os.environ
+        return self._env
 
     def __enter__(self):
         """Enter the context manager context.
@@ -434,14 +456,14 @@ def finish(self):
         if not self.active:
             raise ValueError("Attempt to finish an already finished session")
         self._active = False
-        finish(start_time=self._start_time)
+        finish(env=self._env, start_time=self._start_time)
 
 
 # clean-up functions when terminating a GRASS session
 # these fns can only be called within a valid GRASS session
 
 
-def clean_default_db(*, modified_after=None):
+def clean_default_db(*, modified_after=None, env=None):
     """Clean (vacuum) the default db if it is SQLite
 
     When *modified_after* is set, database is cleaned only when it was modified
@@ -451,11 +473,11 @@ def clean_default_db(*, modified_after=None):
     # pylint: disable=import-outside-toplevel
     import grass.script as gs
 
-    conn = gs.db_connection()
+    conn = gs.db_connection(env=env)
     if not conn or conn["driver"] != "sqlite":
         return
     # check if db exists
-    gis_env = gs.gisenv()
+    gis_env = gs.gisenv(env=env)
     database = conn["database"]
     database = database.replace("$GISDBASE", gis_env["GISDBASE"])
     database = database.replace("$LOCATION_NAME", gis_env["LOCATION_NAME"])
@@ -477,8 +499,8 @@ def clean_default_db(*, modified_after=None):
     # Start the vacuum process, then show the message in parallel while
     # the vacuum is running. Finally, wait for the vacuum process to finish.
     # Error handling is the same as errors="ignore".
-    process = gs.start_command("db.execute", sql="VACUUM")
-    gs.verbose(_("Cleaning up default SQLite database..."))
+    process = gs.start_command("db.execute", sql="VACUUM", env=env)
+    gs.verbose(_("Cleaning up SQLite attribute database..."), env=env)
     process.wait()
 
 
@@ -489,18 +511,23 @@ def call(cmd, **kwargs):
     return subprocess.call(cmd, **kwargs)
 
 
-def clean_temp():
+def clean_temp(env=None):
     """Clean mapset temporary directory"""
     # Lazy-importing to reduce dependencies (this can be eventually removed).
     # pylint: disable=import-outside-toplevel
     import grass.script as gs
 
-    gs.verbose(_("Cleaning up temporary files..."))
-    gisbase = os.environ["GISBASE"]
-    call([os.path.join(gisbase, "etc", "clean_temp")], stdout=subprocess.DEVNULL)
+    if not env:
+        env = os.environ
+
+    gs.verbose(_("Cleaning up temporary files..."), env=env)
+    gisbase = env["GISBASE"]
+    call(
+        [os.path.join(gisbase, "etc", "clean_temp")], stdout=subprocess.DEVNULL, env=env
+    )
 
 
-def finish(*, start_time=None):
+def finish(*, env=None, start_time=None):
     """Terminate the GRASS session and clean up
 
     GRASS commands can no longer be used after this function has been
@@ -518,13 +545,16 @@ def finish(*, start_time=None):
     Currently, it is used to do SQLite database vacuum only when database was modified
     since the session started.
     """
-    clean_default_db(modified_after=start_time)
-    clean_temp()
+    if not env:
+        env = os.environ
+
+    clean_default_db(modified_after=start_time, env=env)
+    clean_temp(env=env)
     # TODO: unlock the mapset?
     # unset the GISRC and delete the file
     from grass.script import utils as gutils
 
-    gutils.try_remove(os.environ["GISRC"])
-    os.environ.pop("GISRC")
+    gutils.try_remove(env["GISRC"])
+    del env["GISRC"]
     # remove gislock env var (not the gislock itself
-    os.environ.pop("GIS_LOCK")
+    del env["GIS_LOCK"]

python/grass/script/tests/grass_script_core_location_test.py

@@ -1,14 +1,15 @@
 """Test functions in grass.script.setup"""
 
 import multiprocessing
+import os
+
+import pytest
 
 import grass.script as gs
 
 
-# All init tests change the global environment, but when it really matters,
-# we use a separate process.
-# Ideally, the functions would support env parameter and the test
-# would mostly use that.
+# This is useful when we want to ensure that function like init does
+# not change the global environment.
 def run_in_subprocess(function):
     """Run function in a separate process
 
@@ -30,15 +31,17 @@ def create_and_get_srid(tmp_path):
     gs.core._create_location_xy(
         tmp_path, bootstrap_location
     )  # pylint: disable=protected-access
-    with gs.setup.init(tmp_path / bootstrap_location):
+    with gs.setup.init(tmp_path / bootstrap_location, env=os.environ.copy()) as session:
         gs.create_location(tmp_path, desired_location, epsg="3358")
         assert (tmp_path / desired_location).exists()
         wkt_file = tmp_path / desired_location / "PERMANENT" / "PROJ_WKT"
         assert wkt_file.exists()
-        gs.run_command("g.gisenv", set=f"GISDBASE={tmp_path}")
-        gs.run_command("g.gisenv", set=f"LOCATION_NAME={desired_location}")
-        gs.run_command("g.gisenv", set="MAPSET=PERMANENT")
-        return gs.parse_command("g.proj", flags="g")["srid"]
+        gs.run_command("g.gisenv", set=f"GISDBASE={tmp_path}", env=session.env)
+        gs.run_command(
+            "g.gisenv", set=f"LOCATION_NAME={desired_location}", env=session.env
+        )
+        gs.run_command("g.gisenv", set="MAPSET=PERMANENT", env=session.env)
+        return gs.parse_command("g.proj", flags="g", env=session.env)["srid"]
 
 
 def test_with_same_path(tmp_path):
@@ -58,19 +61,28 @@ def workload(queue):
     assert epsg == "EPSG:3358"
 
 
+@pytest.mark.usefixtures("mock_no_session")
 def test_without_session(tmp_path):
     """Check that creation works outside of session.
 
     Assumes that there is no session for the test. This can be ensured by running only
-    this test with pylint outside a session.
+    this test with pytest outside of a session.
+
+    Also checks that the global environment is intact after calling the function.
     """
+
     name = "desired"
     gs.create_location(tmp_path, name, epsg="3358")
+
+    # Check that the global environment is still intact.
+    assert not os.environ.get("GISRC"), "Session exists after the call"
+    assert not os.environ.get("GISBASE"), "Runtime exists after the call"
+
     assert (tmp_path / name).exists()
     wkt_file = tmp_path / name / "PERMANENT" / "PROJ_WKT"
     assert wkt_file.exists()
-    with gs.setup.init(tmp_path / name):
-        epsg = gs.parse_command("g.proj", flags="g")["srid"]
+    with gs.setup.init(tmp_path / name, env=os.environ.copy()) as session:
+        epsg = gs.parse_command("g.proj", flags="g", env=session.env)["srid"]
         assert epsg == "EPSG:3358"
 
 
@@ -84,15 +96,19 @@ def test_with_different_path(tmp_path):
     gs.core._create_location_xy(
         tmp_path_a, bootstrap_location
     )  # pylint: disable=protected-access
-    with gs.setup.init(tmp_path_a / bootstrap_location):
+    with gs.setup.init(
+        tmp_path_a / bootstrap_location, env=os.environ.copy()
+    ) as session:
         gs.create_location(tmp_path_b, desired_location, epsg="3358")
         assert (tmp_path_b / desired_location).exists()
         wkt_file = tmp_path_b / desired_location / "PERMANENT" / "PROJ_WKT"
         assert wkt_file.exists()
-        gs.run_command("g.gisenv", set=f"GISDBASE={tmp_path_b}")
-        gs.run_command("g.gisenv", set=f"LOCATION_NAME={desired_location}")
-        gs.run_command("g.gisenv", set="MAPSET=PERMANENT")
-        epsg = gs.parse_command("g.proj", flags="g")["srid"]
+        gs.run_command("g.gisenv", set=f"GISDBASE={tmp_path_b}", env=session.env)
+        gs.run_command(
+            "g.gisenv", set=f"LOCATION_NAME={desired_location}", env=session.env
+        )
+        gs.run_command("g.gisenv", set="MAPSET=PERMANENT", env=session.env)
+        epsg = gs.parse_command("g.proj", flags="g", env=session.env)["srid"]
         assert epsg == "EPSG:3358"
 
 
@@ -105,8 +121,8 @@ def test_path_only(tmp_path):
     assert full_path.exists()
     assert mapset_path.exists()
     assert wkt_file.exists()
-    with gs.setup.init(full_path):
-        epsg = gs.parse_command("g.proj", flags="g")["srid"]
+    with gs.setup.init(full_path, env=os.environ.copy()) as session:
+        epsg = gs.parse_command("g.proj", flags="g", env=session.env)["srid"]
         assert epsg == "EPSG:3358"
 
 
@@ -116,8 +132,8 @@ def test_create_project(tmp_path):
     assert (tmp_path / name).exists()
     wkt_file = tmp_path / name / "PERMANENT" / "PROJ_WKT"
     assert wkt_file.exists()
-    with gs.setup.init(tmp_path / name):
-        epsg = gs.parse_command("g.proj", flags="g")["srid"]
+    with gs.setup.init(tmp_path / name, env=os.environ.copy()) as session:
+        epsg = gs.parse_command("g.proj", flags="g", env=session.env)["srid"]
         assert epsg == "EPSG:3358"
 
 
@@ -128,7 +144,7 @@ def test_files(tmp_path):
     gs.core._create_location_xy(
         tmp_path, bootstrap_location
     )  # pylint: disable=protected-access
-    with gs.setup.init(tmp_path / bootstrap_location):
+    with gs.setup.init(tmp_path / bootstrap_location, env=os.environ.copy()):
         description = "This is a test (not Gauss-Krüger or Křovák)"
         gs.create_location(tmp_path, desired_location, epsg="3358", desc=description)
         assert (tmp_path / desired_location).exists()