Manage zarr and netcdf output in the same way

geopython · Aug 27, 2024 · 0b09281 · 0b09281
1 parent 1a604f8
commit 0b09281
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 29 deletions.
diff --git a/pygeoapi/api/__init__.py b/pygeoapi/api/__init__.py
@@ -88,6 +88,7 @@
 F_JPEG = 'jpeg'
 F_MVT = 'mvt'
 F_NETCDF = 'NetCDF'
+F_ZARR = 'zarr'
 
 #: Formats allowed for ?f= requests (order matters for complex MIME types)
 FORMAT_TYPES = OrderedDict((
@@ -98,6 +99,7 @@
     (F_JPEG, 'image/jpeg'),
     (F_MVT, 'application/vnd.mapbox-vector-tile'),
     (F_NETCDF, 'application/x-netcdf'),
+    (F_ZARR, 'application/zip+zarr'),
 ))
 
 #: Locale used for system responses (e.g. exceptions)
@@ -1113,27 +1115,28 @@ def describe_collections(self, request: Union[APIRequest, Any],
                     'title': l10n.translate('Coverage data', request.locale),
                     'href': f'{self.get_collections_url()}/{k}/coverage?f={F_JSON}'  # noqa
                 })
-                if collection_data_format is not None:
-                    title_ = l10n.translate('Coverage data as', request.locale)  # noqa
-                    title_ = f"{title_} {collection_data_format['name']}"
-                    collection['links'].append({
-                        'type': collection_data_format['mimetype'],
-                        'rel': f'{OGC_RELTYPES_BASE}/coverage',
-                        'title': title_,
-                        'href': f"{self.get_collections_url()}/{k}/coverage?f={collection_data_format['name']}"  # noqa
-                    })
-
                 # Hardcode netcdf format for xarray provider
-                if (collection_data['name'] == 'xarray' and
-                   collection_data_format['name'] == 'zarr'):
-                    title_ = l10n.translate('Coverage data as', request.locale)
-                    title_ = f"{title_} {F_NETCDF}"
-                    collection['links'].append({
-                        'type': FORMAT_TYPES[F_NETCDF],
-                        'rel': f'{OGC_RELTYPES_BASE}/coverage',
-                        'title': title_,
-                        'href': f"{self.get_collections_url()}/{k}/coverage?f={F_NETCDF}"  # noqa
-                    })
+                if collection_data['name'] == 'xarray':
+                    data_formats = [
+                        {'name': F_NETCDF, 'mimetype': FORMAT_TYPES[F_NETCDF]},
+                        {'name': F_ZARR, 'mimetype': FORMAT_TYPES[F_ZARR]}
+                    ]
+                elif collection_data_format is not None:
+                    data_formats = [collection_data_format]
+                else:
+                    data_formats = []
+
+                for data_format in data_formats:
+                    title_ = l10n.translate('Coverage data as', request.locale)  # noqa
+                    title_ = f"{title_} {data_format['name']}"
+                    collection['links'].append(
+                        {
+                            'type': data_format['mimetype'],
+                            'rel': f'{OGC_RELTYPES_BASE}/coverage',
+                            'title': title_,
+                            'href': f"{self.get_collections_url()}/{k}/coverage?f={data_format['name']}",  # noqa
+                        }
+                    )
 
                 if dataset is not None:
                     LOGGER.debug('Creating extended coverage metadata')

diff --git a/pygeoapi/provider/xarray_.py b/pygeoapi/provider/xarray_.py
@@ -59,10 +59,13 @@ def __init__(self, provider_def):
 
         super().__init__(provider_def)
 
+        self.native_format = None
         try:
             if provider_def['data'].endswith('.zarr'):
+                self.native_format = 'zarr'
                 open_func = xarray.open_zarr
             else:
+                self.native_format = 'netcdf'
                 if '*' in self.data:
                     LOGGER.debug('Detected multi file dataset')
                     open_func = xarray.open_mfdataset
@@ -250,19 +253,15 @@ def query(self, properties=[], subsets={}, bbox=[], bbox_crs=4326,
             LOGGER.debug('Creating output in CoverageJSON')
             return self.gen_covjson(out_meta, data, properties)
         elif format_ == 'zarr':
-            LOGGER.debug('Returning data in native zarr format')
+            LOGGER.debug('Returning data in zarr format')
             return _get_zarr_data(data)
         elif format_ == 'netcdf':
             LOGGER.debug('Returning data in netcdf format')
             return _get_netcdf_data(data)
-        else:  # return data in native format
-            with tempfile.NamedTemporaryFile() as fp:
-                LOGGER.debug('Returning data in native NetCDF format')
-                data.to_netcdf(
-                    fp.name
-                )  # we need to pass a string to be able to use the "netcdf4" engine  # noqa
-                fp.seek(0)
-                return fp.read()
+        else:
+            msg = f'Unsupported format: {format_}'
+            LOGGER.error(msg)
+            raise ProviderQueryError(msg)
 
     def gen_covjson(self, metadata, data, fields):
         """