diff --git a/fmtm_splitter/splitter.py b/fmtm_splitter/splitter.py index 90372a5..bd58e61 100755 --- a/fmtm_splitter/splitter.py +++ b/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,8 +177,17 @@ 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( @@ -182,10 +195,11 @@ def splitBySquare( # noqa: N802 ) 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,6 +420,43 @@ 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 = [] @@ -407,16 +464,16 @@ def split_by_square( 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) diff --git a/tests/test_splitter.py b/tests/test_splitter.py index ab9452c..589d87d 100644 --- a/tests/test_splitter.py +++ b/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), ]