From b4e87ee9719003215771ac330f4e0cbee34073c6 Mon Sep 17 00:00:00 2001
From: Dan Snow <dan@sno.ws>
Date: Thu, 31 Oct 2024 04:27:53 -0400
Subject: [PATCH] Setup mypy and fix type errors

---
 .pre-commit-config.yaml     | 10 ++++++++++
 data/src/calculate_times.py |  6 +++---
 data/src/utils/__init__.py  |  0
 data/src/utils/census.py    |  9 ++++++---
 pyproject.toml              | 23 ++++++-----------------
 5 files changed, 25 insertions(+), 23 deletions(-)
 create mode 100644 data/src/utils/__init__.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7a08af3..d2ced78 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,6 +4,16 @@ repos:
     hooks:
       - id: ruff
       - id: ruff-format
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.13.0
+    hooks:
+      - id: mypy
+        args: [--ignore-missing-imports]
+        additional_dependencies: [
+          "pandas-stubs>=2.2.3.241009",
+          "types-PyYAML>=6.0.12.20240917",
+          "types-requests>=2.32.0.20241016"
+        ]
   - repo: https://github.com/adrienverge/yamllint.git
     rev: v1.29.0
     hooks:
diff --git a/data/src/calculate_times.py b/data/src/calculate_times.py
index 73b2de4..998e664 100644
--- a/data/src/calculate_times.py
+++ b/data/src/calculate_times.py
@@ -7,7 +7,7 @@
 from pathlib import Path
 
 import pandas as pd
-import valhalla
+import valhalla  # type: ignore
 import yaml
 from utils.utils import format_time, get_md5_hash
 
@@ -321,7 +321,7 @@ def create_write_path(key: str, out_type: str, output_dict: dict) -> str:
 
     origins = (
         pd.read_parquet(input["files"]["origins_file"])
-        .loc[:, od_cols.keys()]
+        .loc[:, list(od_cols.keys())]
         .rename(columns=od_cols)
         .sort_values(by="id")
     )
@@ -333,7 +333,7 @@ def create_write_path(key: str, out_type: str, output_dict: dict) -> str:
 
     destinations = (
         pd.read_parquet(input["files"]["destinations_file"])
-        .loc[:, od_cols.keys()]
+        .loc[:, list(od_cols.keys())]
         .rename(columns=od_cols)
         .sort_values(by="id")
     )
diff --git a/data/src/utils/__init__.py b/data/src/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/data/src/utils/census.py b/data/src/utils/census.py
index 8ba4def..1088895 100644
--- a/data/src/utils/census.py
+++ b/data/src/utils/census.py
@@ -34,11 +34,14 @@ def weighted_mean(
             return group[value_col].mean()
         return (group[value_col] * group[weight_col]).sum() / total_weight
 
+    def calculate_group_weighted_means(group: pd.DataFrame) -> pd.Series:
+        return pd.Series(
+            {col: weighted_mean(group, weight_col, col) for col in value_cols}
+        )
+
     grouped = df.groupby(group_cols)
     weighted_means = grouped.apply(
-        lambda x: pd.Series(
-            {col: weighted_mean(x, weight_col, col) for col in value_cols}
-        )
+        calculate_group_weighted_means
     ).reset_index()
 
     return weighted_means
diff --git a/pyproject.toml b/pyproject.toml
index d3c3b31..cd841df 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,14 +49,18 @@ site = [
   "pyyaml==6.0.2",
   "requests==2.32.3"
 ]
-test = [
+dev = [
   "setuptools>=61.0",
+  "pandas-stubs>=2.2.3.241009",
+  "pre-commit>=4.0.1",
   "pytest>=7.3.0",
   "pytest-cov>=4.1.0",
+  "types-PyYAML>=6.0.12.20240917",
+  "types-requests>=2.32.0.20241016"
 ]
 
 [build-system]
-requires = ["setuptools >= 61.0"]
+requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
@@ -71,18 +75,3 @@ line-length = 79
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
-
-[tool.tox]
-legacy_tox_ini = """
-[tox]
-min_version = 4.0
-envlist = py39, py38, py310, py311, py312
-
-[gh]
-python =
-    3.12 = py312
-    3.11 = py311
-    3.10 = py310
-    3.9 = py39
-    3.8 = py38
-"""