Merge pull request #578 from mraspaud/add-attrs-to-swath-def

Add attrs to future swath definition

Author: djhoese

pyresample/future/geometry/swath.py

@@ -19,4 +19,64 @@
 
 from __future__ import annotations
 
-from pyresample.geometry import CoordinateDefinition, SwathDefinition  # noqa
+import numpy as np
+
+from pyresample.geometry import _get_highest_level_class  # noqa
+from pyresample.geometry import CoordinateDefinition, DimensionError  # noqa
+from pyresample.geometry import SwathDefinition as LegacySwathDefinition
+
+
+class SwathDefinition(LegacySwathDefinition):
+    """Swath defined by lons and lats.
+
+    Parameters
+    ----------
+    lons : numpy array
+    lats : numpy array
+    crs: pyproj.CRS,
+       The CRS to use. longlat on WGS84 by default.
+    attrs: dict,
+       A dictionary made to store metadata.
+
+    Attributes
+    ----------
+    shape : tuple
+        Swath shape
+    size : int
+        Number of elements in swath
+    ndims : int
+        Swath dimensions
+    lons : object
+        Swath lons
+    lats : object
+        Swath lats
+    cartesian_coords : object
+        Swath cartesian coordinates
+    """
+
+    def __init__(self, lons, lats, crs=None, attrs=None):
+        super().__init__(lons, lats, crs=crs)
+        self.attrs = attrs or {}
+
+    def __getitem__(self, key):
+        """Slice a 2D geographic definition."""
+        y_slice, x_slice = key
+        return self.__class__(
+            lons=self.lons[y_slice, x_slice],
+            lats=self.lats[y_slice, x_slice],
+            attrs=self.attrs
+        )
+
+    def concatenate(self, other):
+        """Concatenate coordinate definitions."""
+        if self.ndim != other.ndim:
+            raise DimensionError(('Unable to concatenate %sD and %sD '
+                                  'geometries') % (self.ndim, other.ndim))
+        if self.crs != other.crs:
+            raise ValueError("Incompatible CRSs.")
+        klass = _get_highest_level_class(self, other)
+        lons = np.concatenate((self.lons, other.lons))
+        lats = np.concatenate((self.lats, other.lats))
+        attrs = self.attrs.copy()
+        attrs.update(other.attrs)
+        return klass(lons, lats, attrs=attrs)

pyresample/test/conftest.py

@@ -76,6 +76,8 @@ def create_test_swath(swath_class):
 
     """
     def _create_test_swath(lons, lats, **kwargs):
+        if swath_class is SwathDefinition:
+            kwargs.pop("nproc", None)
         return swath_class(lons, lats, **kwargs)
     return _create_test_swath
 

pyresample/test/test_geometry/test_swath.py

@@ -462,3 +462,43 @@ def assert_np_dict_allclose(dict1, dict2):
             np.testing.assert_allclose(val, dict2[key])
         except TypeError:
             assert val == dict2[key]
+
+
+def test_future_swath_has_attrs():
+    """Test that future SwathDefinition has attrs."""
+    from pyresample.future.geometry import SwathDefinition
+    lons, lats = _gen_swath_lons_lats()
+    attrs = dict(meta="data")
+    swath = SwathDefinition(lons, lats, attrs=attrs)
+    assert swath.attrs == attrs
+
+
+def test_future_swath_slice_has_attrs():
+    """Test that future sliced SwathDefinition has attrs."""
+    from pyresample.future.geometry import SwathDefinition
+    lons, lats = _gen_swath_lons_lats()
+    attrs = dict(meta="data")
+    swath = SwathDefinition(lons, lats, attrs=attrs)[0:1, 0:1]
+    assert swath.attrs == attrs
+
+
+def test_future_swath_concat_has_attrs():
+    """Test that future concatenated SwathDefinition has attrs."""
+    from pyresample.future.geometry import SwathDefinition
+    lons, lats = _gen_swath_lons_lats()
+    attrs1 = dict(meta1="data")
+    swath1 = SwathDefinition(lons, lats, attrs=attrs1)
+    attrs2 = dict(meta2="data")
+    swath2 = SwathDefinition(lons, lats, attrs=attrs2)
+    swath = swath1.concatenate(swath2)
+    assert swath.attrs == dict(meta1="data", meta2="data")
+
+
+def test_future_swath_concat_fails_on_different_crs():
+    """Test that future concatenated SwathDefinition must have the same crs."""
+    from pyresample.future.geometry import SwathDefinition
+    lons, lats = _gen_swath_lons_lats()
+    swath1 = SwathDefinition(lons, lats, crs="mycrs")
+    swath2 = SwathDefinition(lons, lats, crs="myothercrs")
+    with pytest.raises(ValueError, match="Incompatible CRSs."):
+        _ = swath1.concatenate(swath2)