Adding typehints

.github/workflows/test.yml

@@ -5,9 +5,11 @@ on:
     branches: [main]
   pull_request:
     branches: [main]
+    types: [opened, synchronize, reopened, ready_for_review]
 
 jobs:
   test:
+    if: github.event.pull_request.draft == false
     runs-on: ubuntu-latest
     strategy:
       matrix:

CONTRIBUTING.md

@@ -24,6 +24,7 @@ The following commands are available via the Makefile:
 - **`make clean`** - Clean build artifacts and cache files
 - **`make docs`** - Build documentation
 - **`make docs-clean`** - Clean documentation build artifacts
+- **`make type-check`** - Run a type check
 - **`make help`** - Show all available commands with descriptions
 
 ## Pre-commit Hooks

Makefile

@@ -4,6 +4,9 @@ help:  ## Show this help message
 	@echo "Available commands:"
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "  \033[36m%-15s\033[0m %s\n", $$1, $$2}'
 
+type-check:  ## Run mypy type checker
+	mypy meshiphi
+
 coverage:  ## Generate coverage report (terminal and HTML)
 	pytest --cov=meshiphi --cov-report=term-missing --cov-report=html --cov-report=xml -m "not slow"
 	@echo "\nCoverage report generated. Open htmlcov/index.html to view the HTML report."

meshiphi/cli.py

@@ -1,13 +1,15 @@
+from __future__ import annotations
+
 import argparse
-import json
 import inspect
+import json
 import logging
 
 from meshiphi import __version__ as version
-from meshiphi.utils import setup_logging, timed_call
-from meshiphi.mesh_generation.mesh_builder import MeshBuilder
 from meshiphi.mesh_generation.environment_mesh import EnvironmentMesh
+from meshiphi.mesh_generation.mesh_builder import MeshBuilder
 from meshiphi.test_automation.test_automater import TestAutomater
+from meshiphi.utils import setup_logging, timed_call
 
 logger = logging.getLogger(__name__)
 
@@ -19,7 +21,7 @@ def get_args(
     mesh_arg: bool = False,
     format_arg: bool = False,
     merge_arg: bool = False,
-):
+) -> argparse.Namespace:
     """
     Adds required command line arguments to all CLI entry points.
 
@@ -36,7 +38,7 @@ def get_args(
     ap.add_argument(
         "--version",
         action="version",
-        version="%(prog)s {version}".format(version=version),
+        version=f"%(prog)s {version}",
     )
 
     # Optional arguments used in all CLI entry points
@@ -77,9 +79,7 @@ def get_args(
         )
 
     if merge_arg:
-        ap.add_argument(
-            "merge", help="File location of the environmental mesh to merge with"
-        )
+        ap.add_argument("merge", help="File location of the environmental mesh to merge with")
         ap.add_argument(
             "-d",
             "--directory",
@@ -94,14 +94,14 @@ def get_args(
 
 
 @timed_call
-def rebuild_mesh_cli():
+def rebuild_mesh_cli() -> None:
     """
     CLI entry point for rebuilding the mesh based on its encoded config files.
     """
 
     default_output = "rebuild_mesh.output.json"
     args = get_args(default_output, mesh_arg=True, config_arg=False)
-    logger.info("{} {}".format(inspect.stack()[0][3][:-4], version))
+    logger.info(f"{inspect.stack()[0][3][:-4]} {version}")
 
     mesh_json = json.load(args.mesh)
     config = mesh_json["config"]["mesh_info"]
@@ -110,33 +110,35 @@ def rebuild_mesh_cli():
     rebuilt_mesh = MeshBuilder(config).build_environmental_mesh()
     rebuilt_mesh_json = rebuilt_mesh.to_json()
 
-    logger.info("Saving mesh to {}".format(args.output))
+    logger.info(f"Saving mesh to {args.output}")
 
-    json.dump(rebuilt_mesh_json, open(args.output, "w"), indent=4)
+    with open(args.output, "w") as f:
+        json.dump(rebuilt_mesh_json, f, indent=4)
 
 
 @timed_call
-def create_mesh_cli():
+def create_mesh_cli() -> None:
     """
     CLI entry point for the mesh construction
     """
 
     default_output = "create_mesh.output.json"
     args = get_args(default_output)
-    logger.info("{} {}".format(inspect.stack()[0][3][:-4], version))
+    logger.info(f"{inspect.stack()[0][3][:-4]} {version}")
 
     config = json.load(args.config)
 
     # Discrete Meshing
     cg = MeshBuilder(config).build_environmental_mesh()
 
-    logger.info("Saving mesh to {}".format(args.output))
+    logger.info(f"Saving mesh to {args.output}")
     info = cg.to_json()
-    json.dump(info, open(args.output, "w"), indent=4)
+    with open(args.output, "w") as f:
+        json.dump(info, f, indent=4)
 
 
 @timed_call
-def export_mesh_cli():
+def export_mesh_cli() -> None:
     """
     CLI entry point for exporting a mesh to standard formats.
     Currently supported formats are JSON, GEOJSON, TIF, PNG
@@ -145,17 +147,15 @@ def export_mesh_cli():
     args = get_args("mesh.json", config_arg=False, mesh_arg=True, format_arg=True)
 
     if args.format.upper() == "GEOJSON":
-        args = get_args(
-            "mesh_geo.json", config_arg=False, mesh_arg=True, format_arg=True
-        )
+        args = get_args("mesh_geo.json", config_arg=False, mesh_arg=True, format_arg=True)
 
     elif args.format.upper() == "TIF":
         args = get_args("mesh.tif", config_arg=False, mesh_arg=True, format_arg=True)
 
     elif args.format.upper() == "PNG":
         args = get_args("mesh.png", config_arg=False, mesh_arg=True, format_arg=True)
 
-    logger.info("{} {}".format(inspect.stack()[0][3][:-4], version))
+    logger.info(f"{inspect.stack()[0][3][:-4]} {version}")
 
     mesh = json.load(args.mesh)
     env_mesh = EnvironmentMesh.load_from_json(mesh)
@@ -166,7 +166,7 @@ def export_mesh_cli():
 
 
 @timed_call
-def merge_mesh_cli():
+def merge_mesh_cli() -> None:
     """
     CLI entry point for merging two meshes.
     """
@@ -176,44 +176,41 @@ def merge_mesh_cli():
 
     default_output = "merged_mesh.output.json"
     args = get_args(default_output, config_arg=False, mesh_arg=True, merge_arg=True)
-    logger.info("{} {}".format(inspect.stack()[0][3][:-4], version))
+    logger.info(f"{inspect.stack()[0][3][:-4]} {version}")
 
-    with open(args.mesh.name, "r") as f:
+    with open(args.mesh.name) as f:
         mesh1 = json.load(args.mesh)
     env_mesh1 = EnvironmentMesh.load_from_json(mesh1)
 
     if args.directory:
-        logger.debug(
-            "Merging multiple meshes from directory {} with input mesh".format(
-                args.merge
-            )
-        )
+        logger.debug(f"Merging multiple meshes from directory {args.merge} with input mesh")
 
         merge_dir = args.merge
         merge_meshes = [f for f in listdir(merge_dir) if isfile(join(merge_dir, f))]
 
         for mesh in merge_meshes:
             path = join(merge_dir, mesh)
-            with open(path, "r") as f:
+            with open(path) as f:
                 merge_mesh = json.load(f)
             env_mesh_merge = EnvironmentMesh.load_from_json(merge_mesh)
 
             env_mesh1.merge_mesh(env_mesh_merge)
     else:
-        with open(args.merge, "r") as f:
+        with open(args.merge) as f:
             mesh2 = json.load(f)
         env_mesh2 = EnvironmentMesh.load_from_json(mesh2)
 
         env_mesh1.merge_mesh(env_mesh2)
 
     merged_mesh_json = env_mesh1.to_json()
 
-    logger.info("Saving merged mesh to {}".format(args.output))
-    json.dump(merged_mesh_json, open(args.output, "w"), indent=4)
+    logger.info(f"Saving merged mesh to {args.output}")
+    with open(args.output, "w") as f:
+        json.dump(merged_mesh_json, f, indent=4)
 
 
 @timed_call
-def meshiphi_test_cli():
+def meshiphi_test_cli() -> None:
     """
     CLI Entry point for automated testing. Assumes you have git and pytest installed.
 
@@ -223,7 +220,7 @@ def meshiphi_test_cli():
     """
 
     @setup_logging
-    def get_test_automater_args():
+    def get_test_automater_args() -> argparse.Namespace:
         ap = argparse.ArgumentParser()
         # Add one or two arguments as branches
         # If one, compare current branch to specified one

meshiphi/config_validation/config_validator.py

@@ -1,12 +1,14 @@
-import json
-import jsonschema
 import datetime
+import json
 import re
+from typing import Any, Union
+
+import jsonschema
 
 from meshiphi.config_validation.mesh_schema import mesh_schema
 
 
-def flexi_json_input(config):
+def flexi_json_input(config: Union[str, dict[str, Any]]) -> dict[str, Any]:
     """
     Allows flexible inputs. If a string is parsed, then assume it's a file path
     and read in as a json. If a dict is parsed, then assume it's already a
@@ -23,8 +25,8 @@ def flexi_json_input(config):
     """
     if isinstance(config, str):
         # If str, assume filename
-        with open(config, "r") as fp:
-            config_json = json.load(fp)
+        with open(config) as fp:
+            config_json: dict[str, Any] = json.load(fp)
     elif isinstance(config, dict):
         # If dict, assume it's the config
         config_json = config
@@ -35,7 +37,7 @@ def flexi_json_input(config):
     return config_json
 
 
-def validate_mesh_config(config):
+def validate_mesh_config(config: Union[str, dict[str, Any]]) -> None:
     """
     Validates a mesh config
 
@@ -51,7 +53,7 @@ def validate_mesh_config(config):
         ValidationError: Malformed mesh config
     """
 
-    def assert_valid_time(time_str):
+    def assert_valid_time(time_str: str) -> None:
         """
         Checks if the time strings in the config are correctly formatted
 
@@ -63,9 +65,7 @@ def assert_valid_time(time_str):
         """
         correctly_formatted = False
         # If relative time is parsed
-        if re.match(r"TODAY[+,-]\d+", time_str.replace(" ", "")):
-            correctly_formatted = True
-        elif time_str == "TODAY":
+        if re.match(r"TODAY[+,-]\d+", time_str.replace(" ", "")) or time_str == "TODAY":
             correctly_formatted = True
         # Otherwise check if date is parsed correctly
         else:
@@ -81,7 +81,7 @@ def assert_valid_time(time_str):
         if not correctly_formatted:
             raise ValueError(f"{time_str} is not a valid date format!")
 
-    def assert_valid_cellsize(bound_min, bound_max, cell_size):
+    def assert_valid_cellsize(bound_min: float, bound_max: float, cell_size: float) -> None:
         """
         Ensures that the initial cellbox size can evenly divide the initial
         boundary.
@@ -102,17 +102,17 @@ def assert_valid_cellsize(bound_min, bound_max, cell_size):
     jsonschema.validate(instance=config_json, schema=mesh_schema)
 
     # Check that the dates in the region are valid
-    assert_valid_time(config["region"]["start_time"])
-    assert_valid_time(config["region"]["end_time"])
+    assert_valid_time(config_json["region"]["start_time"])
+    assert_valid_time(config_json["region"]["end_time"])
 
     # Check that cellbox width and height evenly divide boundary
     assert_valid_cellsize(
-        config["region"]["lat_min"],
-        config["region"]["lat_max"],
-        config["region"]["cell_height"],
+        config_json["region"]["lat_min"],
+        config_json["region"]["lat_max"],
+        config_json["region"]["cell_height"],
     )
     assert_valid_cellsize(
-        config["region"]["long_min"],
-        config["region"]["long_max"],
-        config["region"]["cell_width"],
+        config_json["region"]["long_min"],
+        config_json["region"]["long_max"],
+        config_json["region"]["cell_width"],
     )

meshiphi/config_validation/mesh_schema.py

@@ -1,4 +1,6 @@
-region_schema = {
+from typing import Any
+
+region_schema: dict[str, Any] = {
     "type": "object",
     "required": [
         "lat_min",
@@ -23,14 +25,14 @@
     },
 }
 
-dataloader_schema = {
+dataloader_schema: dict[str, Any] = {
     "type": "object",
     "required": ["loader", "params"],
     "additionalProperties": False,
     "properties": {"loader": {"type": "string"}, "params": {"type": "object"}},
 }
 
-splitting_schema = {
+splitting_schema: dict[str, Any] = {
     "type": "object",
     "required": ["split_depth", "minimum_datapoints"],
     "additionalProperties": False,
@@ -40,7 +42,7 @@
     },
 }
 
-mesh_schema = {
+mesh_schema: dict[str, Any] = {
     "type": "object",
     "required": ["region", "data_sources", "splitting"],
     "properties": {

meshiphi/dataloaders/dataloader_interface.py

@@ -1,5 +1,10 @@
+from __future__ import annotations
+
 from abc import ABCMeta, abstractmethod
-from meshiphi.mesh_generation.boundary import Boundary
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from meshiphi.mesh_generation.boundary import Boundary
 
 
 class DataLoaderInterface(metaclass=ABCMeta):
@@ -10,23 +15,22 @@ class DataLoaderInterface(metaclass=ABCMeta):
     """
 
     @classmethod
-    def __subclasshook__(cls, subclass):
+    def __subclasshook__(cls, subclass: type) -> bool:
         return (
             hasattr(subclass, "get_value")
             and callable(subclass.get_value)
             and hasattr(subclass, "get_hom_condition")
             and callable(subclass.get_hom_condition)
-            or NotImplemented
-        )
+        ) or NotImplemented
 
     @abstractmethod
-    def get_value(self, bounds: Boundary, agg_type: str):
+    def get_value(self, bounds: Boundary, agg_type: str) -> dict[str, float]:
         """Reads values within boundary and returns aggregated value as
         a dictionary in the form {data_name (str): value (np.float64)}"""
         raise NotImplementedError
 
     @abstractmethod
-    def get_hom_condition(self, bounds: Boundary, splitting_conds: dict):
+    def get_hom_condition(self, bounds: Boundary, splitting_conds: dict[str, Any]) -> str:
         """Reads values within boundary to determine if data is
         homogeneous or heterogeneous. Return as str"""
         raise NotImplementedError