Implement local caching for WMTS requests

lib/cartopy/io/ogc_clients.py

@@ -22,11 +22,15 @@
 import warnings
 import weakref
 from xml.etree import ElementTree
+import os
 
 import numpy as np
 from PIL import Image
+from pathlib import Path
 import shapely.geometry as sgeom
 
+import cartopy
+
 
 try:
     import owslib.util
@@ -357,7 +361,7 @@ class WMTSRasterSource(RasterSource):
 
     """
 
-    def __init__(self, wmts, layer_name, gettile_extra_kwargs=None):
+    def __init__(self, wmts, layer_name, gettile_extra_kwargs=None, cache=False):
         """
         Parameters
         ----------
@@ -368,6 +372,9 @@ def __init__(self, wmts, layer_name, gettile_extra_kwargs=None):
         gettile_extra_kwargs: dict, optional
             Extra keywords (e.g. time) to pass through to the
             service's gettile method.
+        cache : bool or str, optional
+            If True, the default cache directory is used. If False, no cache is
+            used. If a string, the string is used as the path to the cache.
 
         """
         if WebMapService is None:
@@ -397,6 +404,18 @@ def __init__(self, wmts, layer_name, gettile_extra_kwargs=None):
 
         self._matrix_set_name_map = {}
 
+        # Enable a cache mechanism when cache is equal to True or to a path.
+        self._default_cache = False
+        if cache is True:
+            self._default_cache = True
+            self.cache_path = Path(cartopy.config["cache_dir"])
+        elif cache is False:
+            self.cache_path = None
+        else:
+            self.cache_path = Path(cache)
+        self.cache = set({})
+        self._load_cache()
+
     def _matrix_set_name(self, target_projection):
         key = id(target_projection)
         matrix_set_name = self._matrix_set_name_map.get(key)
@@ -510,6 +529,23 @@ def fetch_raster(self, projection, extent, target_resolution):
 
         return located_images
 
+    @property
+    def _cache_dir(self):
+        """Return the name of the cache directory"""
+        return self.cache_path / self.__class__.__name__
+
+    def _load_cache(self):
+        """Load the cache"""
+        if self.cache_path is not None:
+            cache_dir = self._cache_dir
+            if not cache_dir.exists():
+                os.makedirs(cache_dir)
+                if self._default_cache:
+                    warnings.warn(
+                        'Cartopy created the following directory to cache '
+                        'WMTSRasterSource tiles: {}'.format(cache_dir))
+            self.cache = self.cache.union(set(os.listdir(cache_dir)))
+
     def _choose_matrix(self, tile_matrices, meters_per_unit, max_pixel_span):
         # Get the tile matrices in order of increasing resolution.
         tile_matrices = sorted(tile_matrices,
@@ -642,21 +678,41 @@ def _wmts_images(self, wmts, layer, matrix_set_name, extent,
                 # Get the tile's Image from the cache if possible.
                 img_key = (row, col)
                 img = image_cache.get(img_key)
+
                 if img is None:
-                    try:
-                        tile = wmts.gettile(
-                            layer=layer.id,
-                            tilematrixset=matrix_set_name,
-                            tilematrix=str(tile_matrix_id),
-                            row=str(row), column=str(col),
-                            **self.gettile_extra_kwargs)
-                    except owslib.util.ServiceException as exception:
-                        if ('TileOutOfRange' in exception.message and
-                                ignore_out_of_range):
-                            continue
-                        raise exception
-                    img = Image.open(io.BytesIO(tile.read()))
-                    image_cache[img_key] = img
+                    # Try it from disk cache
+                    if self.cache_path is not None:
+                        filename = f"{img_key[0]}_{img_key[1]}.npy"
+                        cached_file = os.path.join(
+                            self._cache_dir,
+                            filename
+                        )
+                    else:
+                        filename = None
+                        cached_file = None
+
+                    if filename in self.cache:
+                        img = Image.fromarray(np.load(cached_file, allow_pickle=False))
+                    else:
+                        try:
+                            tile = wmts.gettile(
+                                layer=layer.id,
+                                tilematrixset=matrix_set_name,
+                                tilematrix=str(tile_matrix_id),
+                                row=str(row), column=str(col),
+                                **self.gettile_extra_kwargs)
+                        except owslib.util.ServiceException as exception:
+                            if ('TileOutOfRange' in exception.message and
+                                    ignore_out_of_range):
+                                continue
+                            raise exception
+                        img = Image.open(io.BytesIO(tile.read()))
+                        image_cache[img_key] = img
+                        # save image to local cache
+                        if self.cache_path is not None:
+                            np.save(cached_file, img, allow_pickle=False)
+                            self.cache.add(filename)
+
                 if big_img is None:
                     size = (img.size[0] * n_cols, img.size[1] * n_rows)
                     big_img = Image.new('RGBA', size, (255, 255, 255, 255))

lib/cartopy/mpl/geoaxes.py

@@ -2224,7 +2224,7 @@ def streamplot(self, x, y, u, v, **kwargs):
             sp = super().streamplot(x, y, u, v, **kwargs)
         return sp
 
-    def add_wmts(self, wmts, layer_name, wmts_kwargs=None, **kwargs):
+    def add_wmts(self, wmts, layer_name, wmts_kwargs=None, cache=False, **kwargs):
         """
         Add the specified WMTS layer to the axes.
 
@@ -2249,7 +2249,7 @@ def add_wmts(self, wmts, layer_name, wmts_kwargs=None, **kwargs):
         """
         from cartopy.io.ogc_clients import WMTSRasterSource
         wmts = WMTSRasterSource(wmts, layer_name,
-                                gettile_extra_kwargs=wmts_kwargs)
+                                gettile_extra_kwargs=wmts_kwargs, cache=cache)
         return self.add_raster(wmts, **kwargs)
 
     def add_wms(self, wms, layers, wms_kwargs=None, **kwargs):