WIP: First pass at gridded subset by shape

cis/subsetting/subset.py

@@ -136,6 +136,16 @@ def constrain(self, data):
                 data = data.intersection(**intersection_constraint)
             except IndexError:
                 return None
+
+        if _shape is not None:
+            if data.ndim > 2:
+                raise NotImplementedError("Unable to perform shape subset for multidimensional gridded datasets")
+            mask = np.ma.masked_all_like(data.data)
+            mask[np.unravel_index(_get_gridded_subset_region_indices(data, _shape), data.shape)] = False
+            if isinstance(data.data, np.ma.MaskedArray):
+                data.data.mask &= mask
+            else:
+                data.data = np.ma.masked_array(data.data, mask)
         return gridded_data.make_from_cube(data)
 
     def _make_extract_and_intersection_constraints(self, data):

cis/test/unit/subset/test_subset.py

@@ -74,11 +74,11 @@ def test_can_subset_gridded_data_using_multiple_extract_constraints(self):
         assert subset.data.shape == (5, 3, 3, 2)
 
     def test_can_subset_gridded_data_by_shape(self):
-        data = make_from_cube(cis.test.util.mock.make_square_5x3_2d_cube_with_time())
+        data = make_from_cube(cis.test.util.mock.make_square_5x3_2d_cube())
         subset = data.subset(shape=cis.test.util.mock.WKT_DIAMOND)
         # The corner should be masked, but the center not
-        assert subset.data[0, 0, 0].mask
-        assert ~subset.data[2, 1, 0].mask
+        assert subset.data[0, 0].mask
+        assert ~subset.data[2, 1].mask
 
     def test_can_subset_2d_gridded_data_by_longitude_with_wrapping_at_180(self):
         data = make_from_cube(cis.test.util.mock.make_mock_cube(lat_dim_length=5, lon_dim_length=9))