grass.script.setup: Add return context manager from init (#1912)

The grass.script.setup.init function now returns a SessionHandle handle object which is a context manager (a class with __enter__ and __exit__ methods).

Works as the Python open function which can be used without a context manager. This also makes it compatible with the current usage of init.

Supports cases when session object is used to do finish and when reference to the session object is not assigned to variable. Use of global finish is still supported. Multiple calls of finish on session raises exception. Multiple sessions in parallel are not supported (the underlying global finish function does not currently support that). Cleaning is not enforced. However, the context manager makes it easier to do that right.

A simple test using pytest of context manager API uses private (protected) function from grass.script.core (for now).

Co-authored-by: Vaclav Petras <wenzeslaus@gmail.com>

python/grass/script/setup.py

@@ -286,12 +286,20 @@ def init(path, location=None, mapset=None, grass_path=None):
         # end the session
         gs.setup.finish()
 
+    The returned object is a context manager, so the ``with`` statement can be used to
+    ensure that the session is finished (closed) at the end::
+
+        # ... setup sys.path before import
+        import grass.script as gs
+        with gs.setup.init("~/grassdata/nc_spm_08/user1")
+            # ... use GRASS modules here
+
     :param path: path to GRASS database
     :param location: location name
     :param mapset: mapset within given location (default: 'PERMANENT')
     :param grass_path: path to GRASS installation or executable
 
-    :returns: path to ``gisrc`` file (may change in future versions)
+    :returns: reference to a session handle object which is a context manager
     """
     grass_path = get_install_path(grass_path)
     if not grass_path:
@@ -332,7 +340,79 @@ def init(path, location=None, mapset=None, grass_path=None):
     os.environ["GISRC"] = write_gisrc(
         mapset_path.directory, mapset_path.location, mapset_path.mapset
     )
-    return os.environ["GISRC"]
+    return SessionHandle()
+
+
+class SessionHandle:
+    """Object used to manage GRASS sessions.
+
+    Do not create objects of this class directly. Use the *init* function
+    to get a session object.
+
+    Basic usage::
+
+        # ... setup sys.path before import as needed
+
+        import grass.script as gs
+        import grass.script.setup
+
+        session = gs.setup.init("~/grassdata/nc_spm_08/user1")
+
+        # ... use GRASS modules here
+
+        # end the session
+        session.finish()
+
+    Context manager usage::
+
+        # ... setup sys.path before import as needed
+
+        import grass.script as gs
+        import grass.script.setup
+
+        with gs.setup.init("~/grassdata/nc_spm_08/user1"):
+            # ... use GRASS modules here
+        # session ends automatically here
+    """
+
+    def __init__(self, active=True):
+        self._active = active
+
+    @property
+    def active(self):
+        """True if session is active (not finished)"""
+        return self._active
+
+    def __enter__(self):
+        """Enter the context manager context.
+
+        Notably, the session is activated using the *init* function.
+
+        :returns: reference to the object (self)
+        """
+        if not self.active:
+            raise ValueError(
+                "Attempt to use inactive (finished) session as a context manager"
+            )
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Exit the context manager context.
+
+        Finishes the existing session.
+        """
+        self.finish()
+
+    def finish(self):
+        """Finish the session.
+
+        If not used as a context manager, call explicitly to clean and close the mapset
+        and finish the session. No GRASS modules can be called afterwards.
+        """
+        if not self.active:
+            raise ValueError("Attempt to finish an already finished session")
+        self._active = False
+        finish()
 
 
 # clean-up functions when terminating a GRASS session

python/grass/script/tests/grass_script_setup_test.py

@@ -0,0 +1,12 @@
+"""Test functions in grass.script.setup"""
+
+import grass.script as gs
+import grass.script.setup as grass_setup
+
+
+def test_init_as_context_manager(tmp_path):
+    """Check that init function return value works as a context manager"""
+    location = "test"
+    gs.core._create_location_xy(tmp_path, location)  # pylint: disable=protected-access
+    with grass_setup.init(tmp_path / location):
+        gs.run_command("g.region", flags="p")