Add hooks to generate_compass_dataset function (#469)

* guard examples with main function; update conversion script

* update examples based on latest changes

* fix search pruning taking out parents

* add hooks to generate_compass_dataset function

* make hook parameter field names more verbose and add better doc string

Author: nreinicke

python/nrel/routee/compass/compass_app.py

@@ -11,6 +11,7 @@
 from nrel.routee.compass.io.generate_dataset import (
     GeneratePipelinePhase,
     generate_compass_dataset,
+    DatasetHook,
 )
 
 if TYPE_CHECKING:
@@ -136,6 +137,7 @@ def from_graph(
         vehicle_models: Optional[List[str]] = None,
         parallelism: Optional[int] = None,
         overwrite: bool = False,
+        hooks: Optional[List[DatasetHook]] = None,
     ) -> CompassApp:
         """
         Build a CompassApp from a networkx graph.
@@ -207,6 +209,7 @@ def from_graph(
                 raster_resolution_arc_seconds=raster_resolution_arc_seconds,
                 default_config=True,
                 vehicle_models=vehicle_models,
+                hooks=hooks,
             )
 
         if not config_path.exists():

python/nrel/routee/compass/io/generate_dataset.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import enum
+from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Union, TYPE_CHECKING
 from pathlib import Path
 
@@ -19,6 +20,7 @@
 if TYPE_CHECKING:
     import networkx
     import pandas as pd
+    import geopandas
 
 log = logging.getLogger(__name__)
 
@@ -32,6 +34,33 @@
 AggFunc = Callable[[Any], Any]
 
 
+@dataclass
+class HookParameters:
+    """
+    Parameters passed to hooks registered with generate_compass_dataset.
+
+    These parameters allow developers to access and modify the road network
+    data before the dataset generation process completes.
+
+    Attributes:
+        output_directory (Path): The directory where the dataset files are being written.
+        vertices (pd.DataFrame): A DataFrame containing the vertex (node) data
+            of the road network, including coordinates and IDs.
+        edges (geopandas.GeoDataFrame): A GeoDataFrame containing the edge (link)
+            data, including geometries and attributes like distance and speed.
+        graph (networkx.MultiDiGraph): The processed NetworkX graph object
+            representing the road network topology and attributes.
+    """
+
+    output_directory: Path
+    vertices: pd.DataFrame
+    edges: geopandas.GeoDataFrame
+    graph: networkx.MultiDiGraph
+
+
+DatasetHook = Callable[[HookParameters], None]
+
+
 class GeneratePipelinePhase(enum.Enum):
     GRAPH = 1
     CONFIG = 2
@@ -78,6 +107,7 @@ def generate_compass_dataset(
     requests_kwds: Optional[Dict[Any, Any]] = None,
     afdc_api_key: str = "DEMO_KEY",
     vehicle_models: Optional[List[str]] = None,
+    hooks: Optional[List[DatasetHook]] = None,
 ) -> None:
     """
     Processes a graph downloaded via OSMNx, generating the set of input
@@ -109,6 +139,9 @@ def generate_compass_dataset(
             ``["2017_CHEVROLET_Bolt", "2016_TOYOTA_Camry_4cyl_2WD"]``).
             Use :func:`list_available_vehicle_models` to see valid names.
             When ``None`` (the default) all available models are included.
+        hooks: Optional list of callables that take a ``HookParameters`` object.
+            These hooks will be called after the dataset has been generated
+            and before the function returns.
     Example:
         >>> import osmnx as ox
         >>> g = ox.graph_from_place("Denver, Colorado, USA")
@@ -364,6 +397,18 @@ def replace_id(vertex_uuid: pd.Index) -> pd.Series[int]:
             compression="gzip",
         )
 
+    # RUN HOOKS
+    if hooks is not None:
+        log.info(f"running {len(hooks)} dataset generation hooks")
+        params = HookParameters(
+            output_directory=output_directory,
+            vertices=v,
+            edges=e,
+            graph=g1,
+        )
+        for hook in hooks:
+            hook(params)
+
 
 def _resolve_required_model_bins(vehicle_models: List[str]) -> set[str]:
     """