Complete type annotations

.github/workflows/python-build-test.yaml

@@ -37,8 +37,8 @@ jobs:
         run: |
           pip install ".[dev]"
 
-      # - name: Run mypy
-      #   run: mypy . --ignore-missing-imports
+      - name: Run mypy
+        run: mypy --strict ./python/nrel
 
       - name: Validate formatting
         run: |

pyproject.toml

@@ -20,9 +20,24 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 keywords = ["eco routing"]
-dependencies = ["toml"]
+dependencies = [
+    "folium",
+    "pandas",
+    "networkx",
+    "toml",
+]
 [project.optional-dependencies]
-dev = ["black", "pytest", "maturin", "jupyter-book", "sphinx-book-theme"]
+dev = [
+    "black",
+    "jupyter-book",
+    "maturin",
+    "mypy",
+    "pytest",
+    "sphinx-book-theme",
+    "types-requests",
+    "types-setuptools",
+    "types-toml",
+]
 
 [project.urls]
 Homepage = "https://github.com/NREL/routee-compass"

python/nrel/routee/compass/compass_app.py

@@ -4,16 +4,18 @@
 import logging
 
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, cast
 from nrel.routee.compass.routee_compass_py import (
     CompassAppWrapper,
 )
 
 import toml
 
 
+Config = Dict[str, Any]
 Query = Dict[str, Any]
-Result = List[Dict[str, Any]]
+Result = Dict[str, Any]
+Results = List[Result]
 
 
 log = logging.getLogger(__name__)
@@ -30,14 +32,15 @@ def __init__(self, app: CompassAppWrapper):
         self._app = app
 
     @classmethod
-    def get_constructor(cls):
+    def get_constructor(cls) -> type:
         """
         Return the underlying constructor for the application.
         This allows a child class to inherit the CompassApp python class
         and implement its own rust based app constructor, while still using
         the original python methods.
         """
-        return CompassAppWrapper
+        # mypy does not understand that this Rust-defined type *is* a type, so we need to cast
+        return cast(type, CompassAppWrapper)
 
     @classmethod
     def from_config_file(
@@ -48,10 +51,10 @@ def from_config_file(
         Build a CompassApp from a config file
 
         Args:
-            config_file (Union[str, Path]): Path to the config file
+            config_file: Path to the config file
 
         Returns:
-            CompassApp: A CompassApp object
+            app: A CompassApp object
 
         Example:
             >>> from nrel.routee.compass import CompassApp
@@ -66,38 +69,42 @@ def from_config_file(
         return cls.from_dict(toml_config, config_path)
 
     @classmethod
-    def from_dict(cls, config: Dict, working_dir: Optional[Path] = None) -> CompassApp:
+    def from_dict(
+        cls, config: Config, working_dir: Optional[Path] = None
+    ) -> CompassApp:
         """
         Build a CompassApp from a configuration object
 
         Args:
-            config (Dict): Configuration dictionary
-            working_dir (Path): optional path to working directory
+            config: Configuration dictionary
+            working_dir: optional path to working directory
 
         Returns:
-            CompassApp: a CompassApp object
+            app: a CompassApp object
 
         Example:
             >>> from nrel.routee.compass import CompassApp
-            >>> conf = { parallelism: 2 }
+            >>> conf = { "parallelism": 2 }
             >>> app = CompassApp.from_config(conf)
         """
         path_str = str(working_dir.absolute()) if working_dir is not None else ""
         toml_string = toml.dumps(config)
-        app = cls.get_constructor()._from_config_toml_string(toml_string, path_str)
+        # mypy does not know about the attribute on this Rust-defined type
+        app = cls.get_constructor()._from_config_toml_string(toml_string, path_str)  # type: ignore
         return cls(app)
 
     def run(
-        self, query: Union[Query, List[Query]], config: Optional[Dict] = None
-    ) -> Result:
+        self, query: Union[Query, List[Query]], config: Optional[Config] = None
+    ) -> Union[Result, Results]:
         """
         Run a query (or multiple queries) against the CompassApp
 
         Args:
-            query (Union[Dict[str, Any], List[Dict[str, Any]]]): A query or list of queries to run
+            query: A query or list of queries to run
+            config: optional configuration
 
         Returns:
-            List[Dict[str, Any]]: A list of results (or a single result if a single query was passed)
+            results: A list of results (or a single result if a single query was passed)
 
         Example:
             >>> from nrel.routee.compass import CompassApp
@@ -130,7 +137,7 @@ def run(
 
         results_json: List[str] = self._app._run_queries(queries_str, config_str)
 
-        results = list(map(json.loads, results_json))
+        results: Results = list(map(json.loads, results_json))
         if single_query and len(results) == 1:
             return results[0]
         return results
@@ -140,24 +147,24 @@ def graph_edge_origin(self, edge_id: int) -> int:
         get the origin vertex id for some edge
 
         Args:
-            edge_id (int): the id of the edge
+            edge_id: the id of the edge
 
         Returns:
-            int: the vertex id at the source of the edge
+            vertex_id: the vertex id at the source of the edge
         """
-        return self._app.graph_edge_origin(edge_id)
+        return cast(int, self._app.graph_edge_origin(edge_id))
 
     def graph_edge_destination(self, edge_id: int) -> int:
         """
         get the destination vertex id for some edge
 
         Args:
-            edge_id (int): the id of the edge
+            edge_id: the id of the edge
 
         Returns:
-            int: the vertex id at the destination of the edge
+            vertex_id: the vertex id at the destination of the edge
         """
-        return self._app.graph_edge_destination(edge_id)
+        return cast(int, self._app.graph_edge_destination(edge_id))
 
     def graph_edge_distance(
         self, edge_id: int, distance_unit: Optional[str] = None
@@ -166,34 +173,34 @@ def graph_edge_distance(
         get the distance for some edge
 
         Args:
-            edge_id (int): the id of the edge
-            distance_unit (Optional[str]): distance unit, by default meters
+            edge_id: the id of the edge
+            distance_unit: distance unit, by default meters
 
         Returns:
-            int: the distance covered by traversing the edge
+            dist: the distance covered by traversing the edge
         """
-        return self._app.graph_edge_distance(edge_id, distance_unit)
+        return cast(float, self._app.graph_edge_distance(edge_id, distance_unit))
 
     def graph_get_out_edge_ids(self, vertex_id: int) -> List[int]:
         """
         get the list of edge ids that depart from some vertex
 
         Args:
-            vertex_id (int): the id of the vertex
+            vertex_id: the id of the vertex
 
         Returns:
-            List[int]: the edge ids of edges departing from this vertex
+            edges: the edge ids of edges departing from this vertex
         """
-        return self._app.graph_get_out_edge_ids(vertex_id)
+        return cast(List[int], self._app.graph_get_out_edge_ids(vertex_id))
 
     def graph_get_in_edge_ids(self, vertex_id: int) -> List[int]:
         """
         get the list of edge ids that arrive from some vertex
 
         Args:
-            vertex_id (int): the id of the vertex
+            vertex_id: the id of the vertex
 
         Returns:
-            List[int]: the edge ids of edges arriving at this vertex
+            edges: the edge ids of edges arriving at this vertex
         """
-        return self._app.graph_get_in_edge_ids(vertex_id)
+        return cast(List[int], self._app.graph_get_in_edge_ids(vertex_id))

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

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import Callable, Dict, Optional, Union
 from pathlib import Path
 from pkg_resources import resource_filename
@@ -6,45 +7,57 @@
 import logging
 import shutil
 
+import pandas as pd
+import networkx
+from numpy.typing import ArrayLike
+
 from nrel.routee.compass.io.utils import add_grade_to_graph
 
 log = logging.getLogger(__name__)
 
 
+HIGHWAY_TYPE = str
+KM_PER_HR = float
+HIGHWAY_SPEED_MAP = dict[HIGHWAY_TYPE, KM_PER_HR]
+
+# Parameters annotated with this pass through OSMnx, then GeoPandas, then to Pandas,
+# this is a best-effort annotation since the upstream doesn't really have one
+AggFunc = Callable[[ArrayLike], ArrayLike]
+
+
 def generate_compass_dataset(
-    g,
+    g: networkx.MultiDiGraph,
     output_directory: Union[str, Path],
-    hwy_speeds: Optional[Dict] = None,
+    hwy_speeds: Optional[HIGHWAY_SPEED_MAP] = None,
     fallback: Optional[float] = None,
-    agg: Optional[Callable] = None,
+    agg: Optional[AggFunc] = None,
     add_grade: bool = False,
     raster_resolution_arc_seconds: Union[str, int] = 1,
     default_config: bool = True,
-):
+) -> None:
     """
     Processes a graph downloaded via OSMNx, generating the set of input
     files required for running RouteE Compass.
 
     The input graph is assumed to be the direct output of an osmnx download.
 
     Args:
-        g (MultiDiGraph): A network graph.
-        output_directory (Union[str, Path]): Directory path to use for writing new Compass files.
-        hwy_speeds (Optional[Dict], optional): OSM highway types and values = typical speeds (km per
+        g: OSMNx graph used to generate input files
+        output_directory: Directory path to use for writing new Compass files.
+        hwy_speeds: OSM highway types and values = typical speeds (km per
             hour) to assign to edges of that highway type for any edges missing
             speed data. Any edges with highway type not in `hwy_speeds` will be
             assigned the mean preexisting speed value of all edges of that highway
-            type. Defaults to None.
-        fallback (Optional[float], optional): Default speed value (km per hour) to assign to edges whose highway
+            type.
+        fallback: Default speed value (km per hour) to assign to edges whose highway
             type did not appear in `hwy_speeds` and had no preexisting speed
-            values on any edge. Defaults to None.
-        agg (Callable, optional): Aggregation function to impute missing values from observed values.
+            values on any edge.
+        agg: Aggregation function to impute missing values from observed values.
             The default is numpy.mean, but you might also consider for example
-            numpy.median, numpy.nanmedian, or your own custom function. Defaults to numpy.mean.
-        add_grade (bool, optional): If true, add grade information. Defaults to False. See add_grade_to_graph() for more info.
-        raster_resolution_arc_seconds (str, optional): If grade is added, the resolution (in arc-seconds) of the tiles to download (either 1 or 1/3). Defaults to 1.
-        default_config (bool, optional): If true, copy default configuration files into the output directory. Defaults to True.
-        energy_model (str, optional): Which trained RouteE Powertrain should we use? Defaults to "2016_TOYOTA_Camry_4cyl_2WD".
+            numpy.median, numpy.nanmedian, or your own custom function.
+        add_grade: If true, add grade information. See add_grade_to_graph() for more info.
+        raster_resolution_arc_seconds: If grade is added, the resolution (in arc-seconds) of the tiles to download (either 1 or 1/3).
+        default_config: If true, copy default configuration files into the output directory.
 
     Example:
         >>> import osmnx as ox
@@ -95,7 +108,7 @@ def generate_compass_dataset(
     print("processing edges")
     lookup = v.set_index("vertex_uuid")
 
-    def replace_id(vertex_uuid):
+    def replace_id(vertex_uuid: pd.Index) -> pd.Series[int]:
         return lookup.loc[vertex_uuid].vertex_id
 
     e = e.reset_index(drop=False).rename(

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

@@ -8,6 +8,8 @@
 import logging
 from typing import Union
 
+import networkx
+
 log = logging.getLogger(__name__)
 
 
@@ -63,7 +65,9 @@ def _lat_lon_to_tile(coord: tuple[int, int]) -> str:
     return f"{lat_prefix}{abs(lat)}{lon_prefix}{abs(lon)}"
 
 
-def _build_download_link(tile: str, resolution=TileResolution.ONE_ARC_SECOND) -> str:
+def _build_download_link(
+    tile: str, resolution: TileResolution = TileResolution.ONE_ARC_SECOND
+) -> str:
     base_link_fragment = "https://prd-tnm.s3.amazonaws.com/StagedProducts/Elevation/"
     resolution_link_fragment = f"{resolution.value}/TIFF/current/{tile}/"
     filename = f"USGS_{resolution.value}_{tile}.tif"
@@ -75,7 +79,7 @@ def _build_download_link(tile: str, resolution=TileResolution.ONE_ARC_SECOND) ->
 def _download_tile(
     tile: str,
     output_dir: Path = Path("cache"),
-    resolution=TileResolution.ONE_ARC_SECOND,
+    resolution: TileResolution = TileResolution.ONE_ARC_SECOND,
 ) -> Path:
     try:
         import requests
@@ -105,10 +109,10 @@ def _download_tile(
 
 
 def add_grade_to_graph(
-    g,
+    g: networkx.MultiDiGraph,
     output_dir: Path = Path("cache"),
     resolution_arc_seconds: Union[str, int] = 1,
-):
+) -> networkx.MultiDiGraph:
     """
     Adds grade information to the edges of a graph.
     This will download the necessary elevation data from USGS as raster tiles and cache them in the output_dir.
@@ -120,12 +124,12 @@ def add_grade_to_graph(
     * 1/3 arc-second: 350 MB
 
     Args:
-        g (nx.MultiDiGraph): The networkx graph to add grades to.
-        output_dir (Path, optional): The directory to cache the downloaded tiles in. Defaults to Path("cache").
-        resolution_arc_seconds (str, optional): The resolution (in arc-seconds) of the tiles to download (either 1 or 1/3). Defaults to 1.
+        g: The networkx graph to add grades to.
+        output_dir: The directory to cache the downloaded tiles in.
+        resolution_arc_seconds: The resolution (in arc-seconds) of the tiles to download (either 1 or 1/3).
 
     Returns:
-        nx.MultiDiGraph: The graph with grade information added to the edges.
+        g: The graph with grade information added to the edges.
 
     Example:
         >>> import osmnx as ox