Added data extracts to avoid creating tasks with no features

fmtm_splitter/splitter.py

@@ -32,6 +32,7 @@
 from osm_rawdata.postgres import PostgresClient
 from psycopg2.extensions import connection
 from shapely.geometry import Polygon, shape
+from shapely.geometry.geo import mapping
 from shapely.ops import unary_union
 
 from fmtm_splitter.db import (
@@ -153,11 +154,14 @@ def geojson_to_shapely_polygon(
     def splitBySquare(  # noqa: N802
         self,
         meters: int,
+        extract_geojson: Optional[Union[dict, FeatureCollection]] = None,
     ) -> FeatureCollection:
         """Split the polygon into squares.
 
         Args:
             meters (int):  The size of each task square in meters.
+            extract_geojson (dict, FeatureCollection): an OSM extract geojson,
+                containing building polygons, or linestrings.
 
         Returns:
             data (FeatureCollection): A multipolygon of all the task boundaries.
@@ -173,19 +177,29 @@ def splitBySquare(  # noqa: N802
 
         cols = list(np.arange(xmin, xmax + width, width))
         rows = list(np.arange(ymin, ymax + length, length))
-
         polygons = []
+        if extract_geojson:
+            features = (
+                extract_geojson.get("features", extract_geojson)
+                if isinstance(extract_geojson, dict)
+                else extract_geojson.features
+            )
+            extract_geoms = [shape(feature["geometry"]) for feature in features]
+        else:
+            extract_geoms = []
+
         for x in cols[:-1]:
             for y in rows[:-1]:
                 grid_polygon = Polygon(
                     [(x, y), (x + width, y), (x + width, y + length), (x, y + length)]
                 )
                 clipped_polygon = grid_polygon.intersection(self.aoi)
                 if not clipped_polygon.is_empty:
-                    polygons.append(clipped_polygon)
+                    if any(geom.within(clipped_polygon) for geom in extract_geoms):
+                        polygons.append(clipped_polygon)
 
         self.split_features = FeatureCollection(
-            [Feature(geometry=poly) for poly in polygons]
+            [Feature(geometry=mapping(poly)) for poly in polygons]
         )
         return self.split_features
 
@@ -382,6 +396,7 @@ def outputGeojson(  # noqa: N802
 def split_by_square(
     aoi: Union[str, FeatureCollection],
     meters: int = 100,
+    osm_extract: Union[str, FeatureCollection] = None,
     outfile: Optional[str] = None,
 ) -> FeatureCollection:
     """Split an AOI by square, dividing into an even grid.
@@ -391,6 +406,11 @@ def split_by_square(
             GeoJSON string, or FeatureCollection object.
         meters(str, optional): Specify the square size for the grid.
             Defaults to 100m grid.
+        osm_extract (str, FeatureCollection): an OSM extract geojson,
+            containing building polygons, or linestrings.
+            Optional param, if not included an extract is generated for you.
+            It is recommended to leave this param as default, unless you know
+            what you are doing.
         outfile(str): Output to a GeoJSON file on disk.
 
     Returns:
@@ -400,23 +420,60 @@ def split_by_square(
     parsed_aoi = FMTMSplitter.input_to_geojson(aoi)
     aoi_featcol = FMTMSplitter.geojson_to_featcol(parsed_aoi)
 
+    if not osm_extract:
+        config_data = dedent(
+            """
+          query:
+            select:
+            from:
+              - nodes
+              - ways_poly
+              - ways_line
+            where:
+              tags:
+                highway: not null
+                building: not null
+                waterway: not null
+                railway: not null
+                aeroway: not null
+            """
+        )
+        # Must be a BytesIO JSON object
+        config_bytes = BytesIO(config_data.encode())
+
+        pg = PostgresClient(
+            "underpass",
+            config_bytes,
+        )
+        # The total FeatureCollection area merged by osm-rawdata automatically
+        extract_geojson = pg.execQuery(
+            aoi_featcol,
+            extra_params={"fileName": "fmtm_splitter", "useStWithin": False},
+        )
+
+    else:
+        extract_geojson = FMTMSplitter.input_to_geojson(osm_extract)
+    if not extract_geojson:
+        err = "A valid data extract must be provided."
+        log.error(err)
+        raise ValueError(err)
     # Handle multiple geometries passed
     if len(feat_array := aoi_featcol.get("features", [])) > 1:
         features = []
         for index, feat in enumerate(feat_array):
             featcol = split_by_square(
                 FeatureCollection(features=[feat]),
                 meters,
+                None,
                 f"{Path(outfile).stem}_{index}.geojson)" if outfile else None,
             )
-            feats = featcol.get("features", [])
-            if feats:
+            if feats := featcol.get("features", []):
                 features += feats
         # Parse FeatCols into single FeatCol
         split_features = FeatureCollection(features)
     else:
         splitter = FMTMSplitter(aoi_featcol)
-        split_features = splitter.splitBySquare(meters)
+        split_features = splitter.splitBySquare(meters, extract_geojson)
         if not split_features:
             msg = "Failed to generate split features."
             log.error(msg)

tests/test_splitter.py

@@ -61,38 +61,37 @@ def test_init_splitter_types(aoi_json):
     assert str(error.value) == "The input AOI cannot contain multiple geometries."
 
 
-def test_split_by_square_with_dict(aoi_json):
+def test_split_by_square_with_dict(aoi_json, extract_json):
     """Test divide by square from geojson dict types."""
     features = split_by_square(
-        aoi_json.get("features")[0],
-        meters=50,
+        aoi_json.get("features")[0], meters=50, osm_extract=extract_json
     )
-    assert len(features.get("features")) == 54
+    assert len(features.get("features")) == 50
     features = split_by_square(
-        aoi_json.get("features")[0].get("geometry"),
-        meters=50,
+        aoi_json.get("features")[0].get("geometry"), meters=50, osm_extract=extract_json
     )
-    assert len(features.get("features")) == 54
+    assert len(features.get("features")) == 50
 
 
-def test_split_by_square_with_str(aoi_json):
+def test_split_by_square_with_str(aoi_json, extract_json):
     """Test divide by square from geojson str and file."""
     # GeoJSON Dumps
     features = split_by_square(
-        geojson.dumps(aoi_json.get("features")[0]),
-        meters=50,
+        geojson.dumps(aoi_json.get("features")[0]), meters=50, osm_extract=extract_json
     )
-    assert len(features.get("features")) == 54
+    assert len(features.get("features")) == 50
     # JSON Dumps
     features = split_by_square(
         json.dumps(aoi_json.get("features")[0].get("geometry")),
         meters=50,
+        osm_extract=extract_json,
     )
-    assert len(features.get("features")) == 54
+    assert len(features.get("features")) == 50
     # File
     features = split_by_square(
         "tests/testdata/kathmandu.geojson",
         meters=100,
+        osm_extract="tests/testdata/kathmandu_extract.geojson",
     )
     assert len(features.get("features")) == 15
 
@@ -105,26 +104,28 @@ def test_split_by_square_with_file_output():
     outfile = Path(__file__).parent.parent / f"{uuid4()}.geojson"
     features = split_by_square(
         "tests/testdata/kathmandu.geojson",
+        osm_extract="tests/testdata/kathmandu_extract.geojson",
         meters=50,
         outfile=str(outfile),
     )
-    assert len(features.get("features")) == 54
+    assert len(features.get("features")) == 50
     # Also check output file
     with open(outfile, "r") as jsonfile:
         output_geojson = geojson.load(jsonfile)
-    assert len(output_geojson.get("features")) == 54
+    assert len(output_geojson.get("features")) == 50
 
 
-def test_split_by_square_with_multigeom_input(aoi_multi_json):
+def test_split_by_square_with_multigeom_input(aoi_multi_json, extract_json):
     """Test divide by square from geojson dict types."""
     file_name = uuid4()
     outfile = Path(__file__).parent.parent / f"{file_name}.geojson"
     features = split_by_square(
         aoi_multi_json,
         meters=50,
+        osm_extract=extract_json,
         outfile=str(outfile),
     )
-    assert len(features.get("features", [])) == 60
+    assert len(features.get("features", [])) == 50
     for index in [0, 1, 2, 3]:
         assert Path(f"{Path(outfile).stem}_{index}.geojson)").exists()
 
@@ -208,10 +209,22 @@ def test_cli_help(capsys):
 def test_split_by_square_cli():
     """Test split by square works via CLI."""
     infile = Path(__file__).parent / "testdata" / "kathmandu.geojson"
+    extract_geojson = Path(__file__).parent / "testdata" / "kathmandu_extract.geojson"
     outfile = Path(__file__).parent.parent / f"{uuid4()}.geojson"
 
     try:
-        main(["--boundary", str(infile), "--meters", "100", "--outfile", str(outfile)])
+        main(
+            [
+                "--boundary",
+                str(infile),
+                "--meters",
+                "100",
+                "--extract",
+                str(extract_geojson),
+                "--outfile",
+                str(outfile),
+            ]
+        )
     except SystemExit:
         pass
 
@@ -226,6 +239,7 @@ def test_split_by_features_cli():
     infile = Path(__file__).parent / "testdata" / "kathmandu.geojson"
     outfile = Path(__file__).parent.parent / f"{uuid4()}.geojson"
     split_geojson = Path(__file__).parent / "testdata" / "kathmandu_split.geojson"
+    extract_geojson = Path(__file__).parent / "testdata" / "kathmandu_extract.geojson"
 
     try:
         main(
@@ -234,6 +248,8 @@ def test_split_by_features_cli():
                 str(infile),
                 "--source",
                 str(split_geojson),
+                "--extract",
+                str(extract_geojson),
                 "--outfile",
                 str(outfile),
             ]