Repository: mwtoews/gridit

Issue:
Feature: Grid.from_array filter to accept other criteria
The request is to allow `filter` to accept more than one value from an iterable. This would allow a Grid definition to be created without needing to merge several features beforehand.

Comments:

Patch:
diff --git a/gridit/classmethods.py b/gridit/classmethods.py
index ace5ef1..767c9d8 100644
--- a/gridit/classmethods.py
+++ b/gridit/classmethods.py
@@ -143,8 +143,11 @@ def from_vector(
         Input file, such as a shapefile.
     resolution : float
         A grid resolution, e.g. 250.0 for 250m x 250m
-    filter : dict, optional
-        Property filter criteria.
+    filter : dict, str, optional
+        Property filter criteria. For example ``{"id": 4}`` to select one
+        feature with attribute "id" value 4. Or ``{"id": [4, 7, 19]}`` to
+        select features with several values. A SQL WHERE statement can also be
+        used if Fiona 1.9 or later is installed.
     buffer : float, default 0.0
         Add buffer to extents of vector data.
     layer : int or str, default None
@@ -173,38 +176,13 @@ def from_vector(
     with fiona.open(fname, "r", layer=layer) as ds:
         projection = ds.crs_wkt
         if filter:
-            for f in ds:
-                r = []
-                for k in filter.keys():
-                    r.append(f["properties"].get(k, "") == filter[k])
-                if len(r) > 0 and all(r):
-                    break
-            else:
-                raise ValueError(
-                    f"could not find {filter} in {fname} layer {layer}")
-            geom_type = f["geometry"]["type"]
-            if geom_type == "Polygon":
-                ar = np.array(f["geometry"]["coordinates"])
-                assert ar.ndim == 3
-                assert ar.shape[2] >= 2
-                xcs = ar[:, :, 0]
-                ycs = ar[:, :, 1]
-            elif geom_type == "LineString":
-                ar = np.array(f["geometry"]["coordinates"])
-                assert ar.ndim == 2
-                assert ar.shape[1] >= 2
-                xcs = ar[:, 0]
-                ycs = ar[:, 1]
-            elif geom_type == "Point":
-                ar = np.array(f["geometry"]["coordinates"])
-                assert ar.ndim == 1
-                assert ar.shape[0] >= 2
-                xcs = ar[0]
-                ycs = ar[1]
-            else:
-                raise NotImplementedError(
-                    f"unexpected geometry type {geom_type}")
-            bounds = xcs.min(), ycs.min(), xcs.max(), ycs.max()
+            from gridit.file import fiona_filter_collection
+
+            flt = fiona_filter_collection(ds, filter)
+            if len(flt) == 0:
+                logger.error("no features filtered with %s", filter)
+            bounds = flt.bounds
+            flt.close()
         else:  # full shapefile bounds
             bounds = ds.bounds
     shape, top_left = get_shape_top_left(bounds, resolution, buffer)
diff --git a/gridit/file.py b/gridit/file.py
index 58e0877..cf4b855 100644
--- a/gridit/file.py
+++ b/gridit/file.py
@@ -1,8 +1,16 @@
 """File methods."""
+from collections.abc import Iterable
 from pathlib import Path
 
 import numpy as np
 
+__all__ = [
+    "fiona_filter_collection",
+    "fiona_property_type",
+    "float32_is_also_float64",
+    "write_raster",
+    "write_vector",
+]
 
 def float32_is_also_float64(val):
     """Return True if float32 and float64 values are the same."""
@@ -232,3 +240,69 @@ def write_vector(
     with fiona.open(fname, "w", **kwargs) as ds:
         ds.writerecords(recs)
     grid.logger.info("wrote %d features", idxs.size)
+
+
+def fiona_filter_collection(ds, filter):
+    """Returns Fiona collection with applied filter.
+
+    Parameters
+    ----------
+    ds : fiona.Collection
+        Input data source
+    filter : dict, str
+        Property filter criteria. For example ``{"id": 4}`` to select one
+        feature with attribute "id" value 4. Or ``{"id": [4, 7, 19]}`` to
+        select features with several values. A SQL WHERE statement can also be
+        used if Fiona 1.9 or later is installed.
+
+    Returns
+    -------
+    fiona.Collection
+
+    Raises
+    ------
+    ModuleNotFoundError
+        If fiona is not installed.
+    """
+    try:
+        import fiona
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("fiona_filter_collection requires fiona")
+    if not isinstance(ds, fiona.Collection):
+        raise ValueError(f"ds must be fiona.Collection; found {type(ds)}")
+    elif ds.closed:
+        raise ValueError("ds is closed")
+    flt = fiona.io.MemoryFile().open(
+        driver=ds.driver, schema=ds.schema, crs=ds.crs)
+    if isinstance(filter, dict):
+        # check that keys are found in datasource
+        filter_keys = list(filter.keys())
+        ds_attrs = list(ds.schema["properties"].keys())
+        if not set(filter_keys).issubset(ds_attrs):
+            not_found = set(filter_keys).difference(ds_attrs)
+            raise KeyError(
+                f"cannot find filter keys: {not_found}; "
+                f"choose from data source attributes: {ds_attrs}")
+        found = 0
+        for feat in ds:
+            for attr, filt_val in filter.items():
+                feat_val = feat["properties"][attr]
+                if (isinstance(filt_val, Iterable)
+                        and not isinstance(filt_val, str)):
+                    for fv in filt_val:
+                        if feat_val == fv:
+                            found += 1
+                            flt.write(feat)
+                else:
+                    if feat_val == filt_val:
+                        found += 1
+                        flt.write(feat)
+    elif isinstance(filter, str):
+        if fiona.__version__[0:3] < "1.9":
+            raise ValueError(
+                "Fiona 1.9 or later required to use filter str as SQL WHERE")
+        for feat in ds.filter(where=filter):
+            flt.write(feat)
+    else:
+        raise ValueError("filter must be a dict or str")
+    return flt
diff --git a/pyproject.toml b/pyproject.toml
index 00d2fc7..c2a4617 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,6 +27,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+dev = ["gridit[test,optional]"]
 optional = [
     "fiona",
     "flopy",


Test patch:
diff --git a/tests/test_classmethods.py b/tests/test_classmethods.py
index f0a93ee..869753c 100644
--- a/tests/test_classmethods.py
+++ b/tests/test_classmethods.py
@@ -74,10 +74,19 @@ def test_grid_from_vector_point():
     assert grid == expected
 
     # filter
+    grid = Grid.from_vector(points_path, 250, {"id": [5, 9]})
+    expected = Grid(250.0, (19, 7), (1810500.0, 5873750.0), grid.projection)
+    assert grid == expected
+
     grid = Grid.from_vector(points_path, 250, {"id": 5})
     expected = Grid(250.0, (1, 1), (1812000.0, 5869250.0), grid.projection)
     assert grid == expected
 
+    # filter + buffer
+    grid = Grid.from_vector(points_path, 250, {"id": 5}, buffer=240)
+    expected = Grid(250.0, (2, 2), (1811750.0, 5869250.0), grid.projection)
+    assert grid == expected
+
 
 @requires_pkg("fiona")
 def test_grid_from_vector_polygon():
@@ -110,10 +119,27 @@ def test_grid_from_vector_line():
 
     # filter
     grid = Grid.from_vector(lines_path, 250, {"StreamOrde": 5})
-    expected = Grid(250.0, (4, 3), (1811500.0, 5877000.0), grid.projection)
+    expected = Grid(250.0, (19, 14), (1808750.0, 5877000.0), grid.projection)
+    assert grid == expected
+
+    grid = Grid.from_vector(lines_path, 250, {"StreamOrde": [4, 5]})
+    expected = Grid(250.0, (28, 41), (1804750.0, 5877000.0), grid.projection)
     assert grid == expected
 
     # buffer
     grid = Grid.from_vector(lines_path, 250, buffer=500)
     expected = Grid(250.0, (70, 66), (1803000.0, 5878750.0), grid.projection)
     assert grid == expected
+
+
+@requires_pkg("fiona")
+def test_grid_from_vector_filter_sql_where():
+    import fiona
+
+    if fiona.__version__[0:3] < "1.9":
+        pytest.skip("Fiona 1.9 or later required to use SQL WHERE")
+
+    # filter
+    grid = Grid.from_vector(lines_path, 250, "StreamOrde>=5")
+    expected = Grid(250.0, (19, 14), (1808750.0, 5877000.0), grid.projection)
+    assert grid == expected
diff --git a/tests/test_file.py b/tests/test_file.py
index efa3df8..88ed267 100644
--- a/tests/test_file.py
+++ b/tests/test_file.py
@@ -2,8 +2,15 @@
 import pytest
 import numpy as np
 
+from .conftest import datadir
 from gridit import Grid
-from gridit.file import float32_is_also_float64, fiona_property_type
+from gridit.file import (
+    fiona_filter_collection,
+    fiona_property_type,
+    float32_is_also_float64,
+)
+
+points_path = datadir / "waitaku2_points.shp"
 
 
 def test_float32_is_also_float64():
@@ -246,3 +253,60 @@ def test_write_vector(tmp_path, grid_basic, grid_projection):
         grid_basic.write_vector(ar2d.T, "out.shp", "val")
     with pytest.raises(ValueError, match="Unable to detect driver"):
         grid_basic.write_vector(ar2d, "out.nope", "val")
+
+
+def test_fiona_filter_collection():
+    fiona = pytest.importorskip("fiona")
+    expected_schema = {
+        "geometry": "Point",
+        "properties": {"id": "int:10"},
+    }
+    with fiona.open(points_path) as ds:
+        flt = fiona_filter_collection(ds, filter={"id": 0})
+        assert flt.schema == expected_schema
+        assert len(flt) == 0
+        assert flt.bounds == (0.0, 0.0, 0.0, 0.0)
+
+        flt = fiona_filter_collection(ds, filter={"id": 1})
+        assert flt.schema == expected_schema
+        assert len(flt) == 1
+        np.testing.assert_array_almost_equal(
+            flt.bounds,
+            (1814758.4763, 5871013.6156, 1814758.4763, 5871013.6156),
+            decimal=4,
+        )
+
+        flt = fiona_filter_collection(ds, filter={"id": [1, 8]})
+        assert flt.schema == expected_schema
+        assert len(flt) == 2
+        np.testing.assert_array_almost_equal(
+            flt.bounds,
+            (1812243.7372, 5871013.6156, 1814758.4763, 5876813.8657),
+            decimal=4,
+        )
+
+        if fiona.__version__[0:3] >= "1.9":
+            flt = fiona_filter_collection(ds, filter="id=2")
+            assert len(flt) == 1
+            np.testing.assert_array_almost_equal(
+                flt.bounds,
+                (1812459.1405, 5875971.5429, 1812459.1405, 5875971.5429),
+                decimal=4,
+            )
+
+            # errors
+            with pytest.raises(ValueError, match="SQL Expression Parsing Err"):
+                fiona_filter_collection(ds, filter="id==2")
+
+        else:
+            with pytest.raises(ValueError, match="filter str as SQL WHERE"):
+                fiona_filter_collection(ds, filter="id=2")
+
+        with pytest.raises(ValueError, match=r"ds must be fiona\.Collection"):
+            fiona_filter_collection(None, filter={"id": 1})
+
+        with pytest.raises(KeyError, match="cannot find filter keys"):
+            fiona_filter_collection(ds, filter={"ID": 1})
+
+    with pytest.raises(ValueError, match="ds is closed"):
+        fiona_filter_collection(ds, filter={"id": 1})