Merge pull request #447 from emlys/task/446

Move `GDALUseExceptions` and `gdal_use_exceptions` into new utils module

Author: claire-simpson

src/pygeoprocessing/__init__.py

@@ -50,6 +50,8 @@
 from .geoprocessing_core import calculate_slope
 from .geoprocessing_core import raster_band_percentile
 from .slurm_utils import log_warning_if_gdal_will_exhaust_slurm_memory
+from .utils import GDALUseExceptions
+from .utils import gdal_use_exceptions
 
 try:
     __version__ = version('pygeoprocessing')

src/pygeoprocessing/geoprocessing.py

@@ -36,9 +36,8 @@
 from .geoprocessing_core import DEFAULT_CREATION_OPTIONS
 from .geoprocessing_core import DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS
 from .geoprocessing_core import DEFAULT_OSR_AXIS_MAPPING_STRATEGY
-from .geoprocessing_core import gdal_use_exceptions
-from .geoprocessing_core import GDALUseExceptions
 from .geoprocessing_core import INT8_CREATION_OPTIONS
+from .utils import GDALUseExceptions, gdal_use_exceptions
 
 # This is used to efficiently pass data to the raster stats worker if available
 if sys.version_info >= (3, 8):

src/pygeoprocessing/geoprocessing_core.pyx

@@ -28,6 +28,7 @@ from osgeo import osr
 import numpy
 import pygeoprocessing
 from .extensions cimport FastFileIterator, FastFileIteratorCompare, is_close
+from .utils import gdal_use_exceptions
 
 
 BLOCK_BITS = 8
@@ -50,38 +51,6 @@ DEFAULT_OSR_AXIS_MAPPING_STRATEGY = osr.OAMS_TRADITIONAL_GIS_ORDER
 
 LOGGER = logging.getLogger('pygeoprocessing.geoprocessing_core')
 
-
-class GDALUseExceptions:
-    """Context manager that enables GDAL exceptions and restores state after."""
-
-    def __init__(self):
-        pass
-
-    def __enter__(self):
-        self.currentUseExceptions = gdal.GetUseExceptions()
-        gdal.UseExceptions()
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.currentUseExceptions == 0:
-            gdal.DontUseExceptions()
-
-
-def gdal_use_exceptions(func):
-    """Decorator that enables GDAL exceptions and restores state after.
-
-    Args:
-        func (callable): function to call with GDAL exceptions enabled
-
-    Returns:
-        Wrapper function that calls ``func`` with GDAL exceptions enabled
-    """
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        with GDALUseExceptions():
-            return func(*args, **kwargs)
-    return wrapper
-
-
 cdef float _NODATA = -1.0
 
 # This resolves an issue on Mac OS X Catalina where cimporting ``push_heap``

src/pygeoprocessing/kernels.py

@@ -41,7 +41,7 @@
 from numpy.typing import ArrayLike
 from osgeo import gdal
 
-from .geoprocessing_core import gdal_use_exceptions
+from .utils import gdal_use_exceptions
 
 FLOAT32_NODATA = float(numpy.finfo(numpy.float32).min)
 LOGGER = logging.getLogger(__name__)

src/pygeoprocessing/multiprocessing/raster_calculator.py

@@ -22,7 +22,7 @@
 from ..geoprocessing import iterblocks
 from ..geoprocessing import LOGGER
 from ..geoprocessing_core import DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS
-from ..geoprocessing_core import GDALUseExceptions
+from ..utils import GDALUseExceptions
 
 if sys.version_info >= (3, 8):
     import multiprocessing.shared_memory

src/pygeoprocessing/routing/helper_functions.py

@@ -4,7 +4,7 @@
 from ..geoprocessing import get_raster_info
 from ..geoprocessing import raster_calculator
 from ..geoprocessing_core import DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS
-from ..geoprocessing_core import gdal_use_exceptions
+from ..utils import gdal_use_exceptions
 
 
 @gdal_use_exceptions

src/pygeoprocessing/routing/routing.pyx

@@ -61,11 +61,11 @@ import shapely.ops
 import scipy.stats
 
 from ..geoprocessing_core import DEFAULT_OSR_AXIS_MAPPING_STRATEGY
-from ..geoprocessing_core import gdal_use_exceptions, GDALUseExceptions
 from ..geoprocessing_core import DEFAULT_GTIFF_CREATION_TUPLE_OPTIONS
 from ..geoprocessing_core import SPARSE_CREATION_OPTIONS, BLOCK_BITS
 from ..extensions cimport ManagedRaster, is_close
 from ..extensions cimport COL_OFFSETS, ROW_OFFSETS, INFLOW_OFFSETS
+from ..utils import gdal_use_exceptions, GDALUseExceptions
 import pygeoprocessing
 
 LOGGER = logging.getLogger(__name__)

src/pygeoprocessing/routing/watershed.pyx

@@ -26,9 +26,9 @@ import shapely.prepared
 import shapely.wkb
 
 import pygeoprocessing
-from ..geoprocessing_core import gdal_use_exceptions, GDALUseExceptions
 from ..geoprocessing_core import SPARSE_CREATION_OPTIONS
 from ..extensions cimport ManagedRaster
+from ..utils import gdal_use_exceptions, GDALUseExceptions
 
 LOGGER = logging.getLogger(__name__)
 

src/pygeoprocessing/slurm_utils.py

@@ -4,7 +4,7 @@
 
 from osgeo import gdal
 
-from .geoprocessing_core import gdal_use_exceptions
+from .utils import gdal_use_exceptions
 
 LOGGER = logging.getLogger(__name__)
 

src/pygeoprocessing/utils.py

@@ -0,0 +1,33 @@
+import functools
+
+from osgeo import gdal
+
+class GDALUseExceptions:
+    """Context manager that enables GDAL exceptions and restores state after."""
+
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        self.currentUseExceptions = gdal.GetUseExceptions()
+        gdal.UseExceptions()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.currentUseExceptions == 0:
+            gdal.DontUseExceptions()
+
+
+def gdal_use_exceptions(func):
+    """Decorator that enables GDAL exceptions and restores state after.
+
+    Args:
+        func (callable): function to call with GDAL exceptions enabled
+
+    Returns:
+        Wrapper function that calls ``func`` with GDAL exceptions enabled
+    """
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        with GDALUseExceptions():
+            return func(*args, **kwargs)
+    return wrapper