lookup endpoint, better tag filtering

slice.py

@@ -1,5 +1,4 @@
 import argparse
-import json
 import duckdb
 import geohash
 import geopandas as gpd
@@ -9,12 +8,31 @@ from shapely import wkb
 from shapely.geometry import box
 from tqdm import tqdm  # For progress bar
 from typing import List, Dict, Optional
+import json
 
 # Configure logging
 import logging
 logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+def load_tags_config(tags_file: Path) -> Dict:
+    """Loads the tags configuration from a JSON file."""
+    with open(tags_file, "r") as f:
+        return json.load(f)
+
+def filter_tags(tags: Dict, tags_config: Dict) -> Dict:
+    """Filters tags based on the rules defined in the tags configuration."""
+    filtered_tags = {}
+    for key, value in tags.items():
+        if key in tags_config:
+            if tags_config[key] is True:
+                # Keep all values for this key
+                filtered_tags[key] = value
+            elif isinstance(tags_config[key], list) and value in tags_config[key]:
+                # Keep only specific values for this key
+                filtered_tags[key] = value
+    return filtered_tags
+
 def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
     """Generates all 4-character Geohashes that intersect a bounding box."""
     geohashes = set()
@@ -120,13 +138,21 @@ def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.G
     updated_gdf.to_parquet(geohash_file, index=False)
     return True  # Water polygons added
 
-def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Optional[str], export_parquet: bool, export_geojson: bool, export_cbor: bool):
-    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and optionally adds water polygons."""
+def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Optional[str], tags_file: Optional[str], export_parquet: bool, export_geojson: bool, export_cbor: bool):
+    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles, optionally adds water polygons, and filters tags."""
     input_path, output_path = Path(input_file), Path(output_dir)
     if not input_path.exists():
         raise FileNotFoundError(f"File '{input_file}' not found.")
     output_path.mkdir(parents=True, exist_ok=True)
 
+    # Load tags configuration if provided
+    tags_config = {}
+    if tags_file:
+        tags_path = Path(tags_file)
+        if not tags_path.exists():
+            raise FileNotFoundError(f"File '{tags_file}' not found.")
+        tags_config = load_tags_config(tags_path)
+
     # Load water dataset if provided
     water_gdf = None
     if water_file:
@@ -191,6 +217,12 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Opt
                 filtered_data = filtered_data.dropna(subset=["geometry"])
 
                 if not filtered_data.empty:
+                    # Filter tags if tags configuration is provided
+                    if tags_config:
+                        filtered_data["tags"] = filtered_data["tags"].apply(
+                            lambda tags: filter_tags(tags, tags_config) if isinstance(tags, dict) else {}
+                        )
+
                     gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
 
                     # Export to GeoParquet
@@ -214,10 +246,11 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Opt
     print(f"🎉 Processing complete! Generated {total_tiles} tiles, {tiles_with_water} of which had water polygons added.")
 
 def main():
-    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, optionally add water polygons, and export in desired formats.")
+    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, optionally add water polygons, filter tags, and export in desired formats.")
     parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.")
     parser.add_argument("-o", "--output", required=True, help="Directory to save output files.")
     parser.add_argument("-w", "--water", required=False, help="Path to the water.parquet file.")
+    parser.add_argument("-t", "--tags", required=False, help="Path to the tags.json file.")
     parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet")
     parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON")
     parser.add_argument("--cbor", action="store_true", help="Export as CBOR")
@@ -228,7 +261,7 @@ def main():
         args.parquet = True  # Default to GeoParquet if no options are given
 
     try:
-        slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor)
+        slice_and_split_geoparquet(args.input, args.output, args.water, args.tags, args.parquet, args.geojson, args.cbor)
     except Exception as e:
         logger.error(f"❌ Error: {e}")