Introduce ConleyMetric type alias; fix public conley_metric signature

Address P2 Code Quality finding from CI Codex review of PR #411 on
0500909. Public signatures advertised `conley_metric: str` but the
runtime + docs + tests + REGISTRY all accept a callable
`(coords1, coords2) -> n×n` for custom (e.g. network) distance metrics.
Static type checkers and IDE autocomplete were therefore lying about the
public API contract.

Introduce a shared `ConleyMetric` type alias in `diff_diff/conley.py`:

    ConleyMetric = Union[
        Literal["haversine", "euclidean"],
        Callable[[np.ndarray, np.ndarray], np.ndarray],
    ]

Sweep all 11 public/internal annotation sites:
- `solve_ols` overloads (4) + impl (1)
- `_solve_ols_numpy` overloads (4) + impl (1)
- `compute_robust_vcov`
- `_compute_robust_vcov_numpy`
- `LinearRegression.__init__`

Internal helpers in `conley.py` (`_pairwise_distance_matrix`,
`_validate_conley_kwargs`, `_compute_conley_vcov`) also annotated.
Moved the `from diff_diff.conley import ...` from late-file (line ~1192
with `noqa: E402`) to top-of-module so the `ConleyMetric` name is in
scope for the early function signatures (linalg.py has no
`from __future__ import annotations`, so types evaluate eagerly).

Verified locally: 72 Conley tests pass, ruff clean, sphinx -W exits 0.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/conley.py

@@ -22,10 +22,22 @@
 from __future__ import annotations
 
 import warnings
-from typing import Optional
+from typing import Callable, Literal, Optional, Union
 
 import numpy as np
 
+# Public type alias for the ``conley_metric`` parameter accepted by
+# ``compute_robust_vcov``, ``solve_ols``, and ``LinearRegression``. The
+# implementation accepts the two named strings as well as a user-supplied
+# callable ``(coords1, coords2) -> n×n distance matrix`` for custom
+# (e.g. network) distance metrics. Exported so the public signatures
+# in :mod:`diff_diff.linalg` can advertise the full accepted type to
+# static checkers and IDEs.
+ConleyMetric = Union[
+    Literal["haversine", "euclidean"],
+    Callable[[np.ndarray, np.ndarray], np.ndarray],
+]
+
 # Earth's mean radius (km), matching R conleyreg's haversine convention
 # (Düsterhöft 2021, conleyreg::haversine_dist in src/distance_functions.cpp,
 # CRAN v0.1.9). WGS-84 equatorial radius is 6378.137 km; the 0.01 km delta
@@ -63,7 +75,7 @@ def _haversine_km(
     return _CONLEY_EARTH_RADIUS_KM * 2.0 * np.arcsin(np.sqrt(a))
 
 
-def _pairwise_distance_matrix(coords: np.ndarray, metric) -> np.ndarray:
+def _pairwise_distance_matrix(coords: np.ndarray, metric: ConleyMetric) -> np.ndarray:
     """Build the dense n×n pairwise distance matrix.
 
     ``metric`` is one of ``"haversine"`` (lat/lon in degrees, distance in km),
@@ -116,7 +128,7 @@ def _uniform_kernel(u: np.ndarray) -> np.ndarray:
 def _validate_conley_kwargs(
     coords: Optional[np.ndarray],
     cutoff: Optional[float],
-    metric,
+    metric: ConleyMetric,
     kernel: str,
     n: int,
 ) -> None:
@@ -197,7 +209,7 @@ def _compute_conley_vcov(
     residuals: np.ndarray,
     coords: np.ndarray,
     cutoff: float,
-    metric,
+    metric: ConleyMetric,
     kernel: str,
     bread_matrix: np.ndarray,
 ) -> np.ndarray:

diff_diff/linalg.py

@@ -49,6 +49,16 @@
     _rust_solve_ols,
 )
 
+# Conley (1999) spatial HAC helpers live in diff_diff.conley to keep this
+# module focused on linear-algebra primitives. Imported at the top so the
+# `ConleyMetric` type alias is in scope for the public function signatures
+# below (which advertise `conley_metric: ConleyMetric`).
+from diff_diff.conley import (
+    ConleyMetric,
+    _compute_conley_vcov,
+    _validate_conley_kwargs,
+)
+
 # =============================================================================
 # Utility Functions
 # =============================================================================
@@ -352,7 +362,7 @@ def solve_ols(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]: ...
 
@@ -374,7 +384,7 @@ def solve_ols(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]: ...
 
@@ -396,7 +406,7 @@ def solve_ols(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Union[
     Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
@@ -454,7 +464,7 @@ def solve_ols(
     vcov_type: str = "hc1",
     conley_coords: Optional[np.ndarray] = None,
     conley_cutoff_km: Optional[float] = None,
-    conley_metric: str = "haversine",
+    conley_metric: ConleyMetric = "haversine",
     conley_kernel: str = "bartlett",
 ) -> Union[
     Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
@@ -818,7 +828,7 @@ def _solve_ols_numpy(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]: ...
 
@@ -838,7 +848,7 @@ def _solve_ols_numpy(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]: ...
 
@@ -858,7 +868,7 @@ def _solve_ols_numpy(
     vcov_type: str = ...,
     conley_coords: Optional[np.ndarray] = ...,
     conley_cutoff_km: Optional[float] = ...,
-    conley_metric: str = ...,
+    conley_metric: ConleyMetric = ...,
     conley_kernel: str = ...,
 ) -> Union[
     Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
@@ -880,7 +890,7 @@ def _solve_ols_numpy(
     vcov_type: str = "hc1",
     conley_coords: Optional[np.ndarray] = None,
     conley_cutoff_km: Optional[float] = None,
-    conley_metric: str = "haversine",
+    conley_metric: ConleyMetric = "haversine",
     conley_kernel: str = "bartlett",
 ) -> Union[
     Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
@@ -1175,11 +1185,8 @@ def resolve_vcov_type(
     return vcov_type
 
 
-# Conley (1999) spatial HAC helpers live in diff_diff.conley to keep linalg.py
-# focused on linear-algebra primitives. Imported here so the dispatch in
-# `_compute_robust_vcov_numpy` can route `vcov_type="conley"` without a
-# late/local import.
-from diff_diff.conley import _compute_conley_vcov, _validate_conley_kwargs  # noqa: E402
+# Conley helpers are imported at module top — see the from-import near the
+# header of this file.
 
 
 def compute_robust_vcov(
@@ -1193,7 +1200,7 @@ def compute_robust_vcov(
     *,
     conley_coords: Optional[np.ndarray] = None,
     conley_cutoff_km: Optional[float] = None,
-    conley_metric: str = "haversine",
+    conley_metric: ConleyMetric = "haversine",
     conley_kernel: str = "bartlett",
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
@@ -1640,7 +1647,7 @@ def _compute_robust_vcov_numpy(
     *,
     conley_coords: Optional[np.ndarray] = None,
     conley_cutoff_km: Optional[float] = None,
-    conley_metric: str = "haversine",
+    conley_metric: ConleyMetric = "haversine",
     conley_kernel: str = "bartlett",
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
@@ -2478,7 +2485,7 @@ def __init__(
         vcov_type: Optional[str] = None,
         conley_coords: Optional[np.ndarray] = None,
         conley_cutoff_km: Optional[float] = None,
-        conley_metric: str = "haversine",
+        conley_metric: ConleyMetric = "haversine",
         conley_kernel: str = "bartlett",
     ):
         self.include_intercept = include_intercept

docs/api/_autosummary/diff_diff.ChaisemartinDHaultfoeuilleResults.rst

@@ -41,6 +41,7 @@
       ~ChaisemartinDHaultfoeuilleResults.p_value
       ~ChaisemartinDHaultfoeuilleResults.path_cumulated_event_study
       ~ChaisemartinDHaultfoeuilleResults.path_effects
+      ~ChaisemartinDHaultfoeuilleResults.path_heterogeneity_effects
       ~ChaisemartinDHaultfoeuilleResults.path_placebo_event_study
       ~ChaisemartinDHaultfoeuilleResults.path_sup_t_bands
       ~ChaisemartinDHaultfoeuilleResults.placebo_event_study