Fix/update geo layer export

src/gisaf/exporters.py

@@ -0,0 +1,108 @@
+from os import remove
+from pathlib import Path
+from datetime import datetime
+from io import BytesIO
+from zipfile import ZipFile
+import logging
+import tempfile
+
+from geopandas.io.file import infer_schema
+
+import fiona
+
+from gisaf.registry import registry
+from gisaf.redis_tools import store as redis_store, RedisError
+
+logger = logging.getLogger(__name__)
+
+async def export_with_fiona(store_names, driver, mimetype, extension, filter_columns=None, reproject=False):
+    """
+    Use fiona to export geo data.
+    registry: gisaf.registry.ModelRegistry
+    store_names: comma separated string of store (aka. layer) names
+    driver: fiona driver (one of fiona.supported_drivers)
+    extension: extension of the file name
+    filter_columns: list of column names to filter out
+    reproject: if true-ish, the geometries are reprojected to the srid specified in conf.srid_for_proj
+    """
+    layers_features = {}
+    for store_name in store_names.split(','):
+        try:
+            if store_name in registry.geom:
+                layers_features[store_name] = await registry.geom[store_name].get_geo_df(reproject=reproject)
+            else:
+                ## Live
+                ## TODO: make live check more explicit
+                layers_features[store_name] = await redis_store.get_gdf(store_name, reproject=reproject)
+        except RedisError as err:
+            logger.warn(f'Cannot get store {store_name}: {err}')
+        except Exception as err:
+            logger.warn(f'Cannot get store {store_name}, see below')
+            logger.exception(err)
+
+    ## FIXME: only 1 layer gets exported with BytesIO, so use a real file
+
+    #filename = '{}_{:%Y-%m-%d_%H:%M:%S}.{}'.format(layers, datetime.now(), extension)
+    filename = 'Gisaf export {:%Y-%m-%d_%H:%M:%S}_{}.{}'.format(
+        datetime.now(), next(tempfile._get_candidate_names()), extension)
+    ## XXX: fails in case of a lot of layers
+    data_file_name = Path(tempfile._get_default_tempdir()) / filename
+    #data_file_name = Path(tempfile._get_default_tempdir()) / next(tempfile._get_candidate_names())
+
+    ## XXX: MemoryFile doesn't support multiple layers (I opened https://github.com/Toblerity/Fiona/issues/830)
+    #with fiona.io.MemoryFile(filename='selected_layers.gpkg') as mem_file:
+    #    for layer_name, gdf in layers_features.items():
+    #        if filter_columns:
+    #            gdf = gdf.filter(filter_columns)
+    #        schema = infer_schema(gdf)
+    #        with mem_file.open(layer=layer_name, driver=driver, crs=gdf.crs, schema=schema) as mem_sink:
+    #            mem_sink.writerecords(gdf.iterfeatures())
+    #return mem_file, filename, mimetype
+
+    with fiona.Env():
+        for layer_name, gdf in layers_features.items():
+            ## XXX: geopandas doesn't accept BytesIO: using fiona directly
+            #gdf.to_file(data, driver=driver, mode='a')
+            _gdf = gdf.reset_index()
+            _gdf['fid'] = _gdf['id']
+            if filter_columns:
+                _gdf = _gdf.filter(filter_columns)
+            schema = infer_schema(_gdf)
+            with fiona.Env(OSR_WKT_FORMAT="WKT2_2018"), fiona.open(
+                data_file_name, 'w',
+                driver=driver,
+                crs=_gdf.crs.to_string(),
+                layer=layer_name,
+                schema=schema) as colxn:
+                colxn.writerecords(_gdf.iterfeatures())
+        #data.seek(0)
+        with open(data_file_name, 'rb') as data_file:
+            data = data_file.read()
+        remove(data_file_name)
+        return data, filename, mimetype
+
+
+async def export_with_pyshp(store_names, reproject=False):
+    """
+    Zip and return data using "old style", ie. with pyshp 1.2
+    """
+    ## TODO: migrate to fiona, see below
+    zip_file = BytesIO()
+    with ZipFile(zip_file, 'w') as zip:
+        for layer_name in store_names.split(','):
+            model = registry.geom[layer_name]
+            dbf, shp, shx, qml, proj_str = await model.get_shapefile_files()
+            zip.writestr('{}.dbf'.format(layer_name), dbf.getvalue())
+            zip.writestr('{}.shp'.format(layer_name), shp.getvalue())
+            zip.writestr('{}.shx'.format(layer_name), shx.getvalue())
+            if qml:
+                zip.writestr('{}.qml'.format(layer_name), qml)
+            if proj_str:
+                zip.writestr('{}.prj'.format(layer_name), proj_str)
+
+    zip_file.seek(0)
+    filename = '{}_{:%Y-%m-%d_%H:%M}.zip'.format(store_names, datetime.now(), )
+    content_type = 'application/zip'
+
+    return zip_file.read(), filename, content_type
+