Reject object-dtype missing cluster IDs across all three surfaces

CI review follow-up: the floating-dtype-only missing-ID guard in
`bias_corrected_local_linear` and `lprobust` let object-dtype
arrays with `None` / object `np.nan` sentinels bypass validation.
The downstream `lprobust_vce` cluster loop would then group on
`np.unique`, treating the sentinel as a real cluster and silently
misstating clustered SE.

Extract the dtype-agnostic `_cluster_has_missing` helper already
used inside `lpbwselect_mse_dpi` and apply it at all three entry
points: the selector, the port-level `lprobust`, and the public
`bias_corrected_local_linear` wrapper.

Regression tests: object-dtype cluster arrays with None and with
np.nan sentinels raise a targeted ValueError at both the wrapper
(test_bias_corrected_lprobust) and the port (test_nprobust_port)
entry points.

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

Author: igerber

diff_diff/_nprobust_port.py

@@ -84,6 +84,33 @@
 _VALID_VCE = ("nn", "hc0", "hc1", "hc2", "hc3")
 
 
+def _cluster_has_missing(cluster: np.ndarray) -> bool:
+    """Detect missing cluster IDs across float / object / string dtypes.
+
+    nprobust::lpbwselect complete-case-filters (x, y, cluster) before
+    dispatch. This port deliberately rejects missingness instead so
+    callers see it rather than silently losing rows. Used by
+    ``lpbwselect_mse_dpi`` and ``lprobust`` (and the public
+    ``bias_corrected_local_linear`` wrapper) so all three surfaces
+    honor the same contract.
+    """
+    if cluster.dtype.kind in ("f", "c"):
+        return bool(np.any(~np.isfinite(cluster)))
+    # Object / string / None-containing arrays: treat None and NaN-like
+    # sentinels as missing.
+    try:
+        if bool(np.any([v is None for v in cluster])):
+            return True
+    except TypeError:
+        pass
+    try:
+        # np.nan comparisons are False; cast to float and check finiteness.
+        cluster_f = cluster.astype(np.float64, copy=False)
+        return bool(np.any(~np.isfinite(cluster_f)))
+    except (TypeError, ValueError):
+        return False
+
+
 # =============================================================================
 # Kernel (W.fun, npfunctions.R:1-7)
 # =============================================================================
@@ -684,25 +711,9 @@ def lpbwselect_mse_dpi(
         # before dispatch; this port deliberately rejects instead so
         # callers see the missingness rather than lose rows silently.
         # The "reject" vs "filter" choice is documented in the module
-        # docstring deviations list.
-        has_missing = False
-        if cluster.dtype.kind in ("f", "c"):
-            has_missing = bool(np.any(~np.isfinite(cluster)))
-        else:
-            # object / string / None-containing arrays: treat None and
-            # NaN-like sentinels as missing.
-            try:
-                has_missing = bool(np.any([x is None for x in cluster]))
-            except TypeError:
-                has_missing = False
-            if not has_missing:
-                try:
-                    # np.nan comparisons are False; use pd-style check.
-                    cluster_f = cluster.astype(np.float64, copy=False)
-                    has_missing = bool(np.any(~np.isfinite(cluster_f)))
-                except (TypeError, ValueError):
-                    pass
-        if has_missing:
+        # docstring deviations list. Dtype-agnostic via
+        # `_cluster_has_missing`.
+        if _cluster_has_missing(cluster):
             raise ValueError(
                 "cluster contains missing values (NaN / None). Unlike "
                 "nprobust::lpbwselect which complete-case-filters "
@@ -1130,13 +1141,17 @@ def lprobust(
             raise ValueError(
                 f"cluster length ({cluster.shape[0]}) does not match x/y ({N})."
             )
-        # Reject NaN cluster IDs (Phase 1b convention: surface missingness
-        # rather than silently drop rows).
-        cluster_float = np.asarray(cluster, dtype=np.float64).ravel() if np.issubdtype(
-            cluster.dtype, np.floating
-        ) else None
-        if cluster_float is not None and np.any(~np.isfinite(cluster_float)):
-            raise ValueError("cluster contains non-finite values (NaN or Inf).")
+        # Dtype-agnostic missingness check. Float NaN/Inf, object None,
+        # and object np.nan all get rejected here (shared with
+        # `lpbwselect_mse_dpi` via `_cluster_has_missing`) so the
+        # downstream `lprobust_vce` cluster grouping on `np.unique`
+        # cannot silently treat a missing sentinel as a real cluster.
+        if _cluster_has_missing(cluster):
+            raise ValueError(
+                "cluster contains missing values (NaN / None). "
+                "Filter your data before the call or drop missing "
+                "observations explicitly."
+            )
 
     # --- vce="nn" setup: sort ascending, precompute dups ---
     dups: Optional[np.ndarray] = None

diff_diff/local_linear.py

@@ -1050,19 +1050,27 @@ def bias_corrected_local_linear(
     d, y = _validate_had_inputs(d, y, boundary)
     n_total = int(d.shape[0])
 
-    # Reject missing cluster IDs up front (Phase 1b convention).
+    # Reject missing cluster IDs up front (Phase 1b convention). Delegates
+    # to the dtype-agnostic `_cluster_has_missing` helper in the port so
+    # wrapper, port-level `lprobust`, and `lpbwselect_mse_dpi` all enforce
+    # the same missing-sentinel contract across float / object / string
+    # dtypes (CI review PR #340 P1 follow-up).
     cluster_arr: Optional[np.ndarray] = None
     if cluster is not None:
+        from diff_diff._nprobust_port import _cluster_has_missing
+
         cluster_arr = np.asarray(cluster).ravel()
         if cluster_arr.shape[0] != n_total:
             raise ValueError(
                 f"cluster length ({cluster_arr.shape[0]}) does not match "
                 f"d/y ({n_total})."
             )
-        if np.issubdtype(cluster_arr.dtype, np.floating) and not np.all(
-            np.isfinite(cluster_arr)
-        ):
-            raise ValueError("cluster contains non-finite values (NaN or Inf).")
+        if _cluster_has_missing(cluster_arr):
+            raise ValueError(
+                "cluster contains missing values (NaN / None). Filter "
+                "your data before the call or drop missing observations "
+                "explicitly."
+            )
 
     # --- Resolve (h, b) ---
     # nprobust's lprobust() with the default rho=1 sets b = h / rho = h

tests/test_bias_corrected_lprobust.py

@@ -311,11 +311,34 @@ def test_weights_raises_not_implemented(self):
             bias_corrected_local_linear(d, y, h=0.3, weights=np.ones(100))
 
     def test_cluster_nan_raises(self):
+        """Float NaN in cluster IDs is rejected."""
         d = np.linspace(0.0, 1.0, 100)
         y = d.copy()
         cluster = np.repeat(np.arange(10), 10).astype(np.float64)
         cluster[5] = np.nan
-        with pytest.raises(ValueError, match="cluster contains non-finite"):
+        with pytest.raises(ValueError, match="cluster contains missing"):
+            bias_corrected_local_linear(d, y, h=0.3, cluster=cluster)
+
+    def test_cluster_object_none_raises(self):
+        """Object-dtype cluster with a ``None`` sentinel is rejected.
+        Float-only checks let this through; the dtype-agnostic helper
+        catches it (CI review PR #340 follow-up P1)."""
+        d = np.linspace(0.0, 1.0, 100)
+        y = d.copy()
+        cluster = np.array(
+            [i // 10 if i != 5 else None for i in range(100)], dtype=object
+        )
+        with pytest.raises(ValueError, match="cluster contains missing"):
+            bias_corrected_local_linear(d, y, h=0.3, cluster=cluster)
+
+    def test_cluster_object_nan_raises(self):
+        """Object-dtype cluster with np.nan is rejected."""
+        d = np.linspace(0.0, 1.0, 100)
+        y = d.copy()
+        cluster = np.array(
+            [i // 10 if i != 5 else np.nan for i in range(100)], dtype=object
+        )
+        with pytest.raises(ValueError, match="cluster contains missing"):
             bias_corrected_local_linear(d, y, h=0.3, cluster=cluster)
 
     def test_unknown_kernel_raises(self):

tests/test_nprobust_port.py

@@ -588,6 +588,37 @@ def test_lprobust_shifted_boundary_dgp5(self, lprobust_golden):
         np.testing.assert_allclose(res.se_cl, g["se_cl"], atol=1e-12, rtol=1e-12)
         np.testing.assert_allclose(res.se_rb, g["se_rb"], atol=1e-12, rtol=1e-12)
 
+    def test_lprobust_cluster_object_none_raises(self):
+        """Port-level: object-dtype cluster with None sentinel is rejected.
+        Mirror of the wrapper test; pins that `_cluster_has_missing` fires
+        inside the direct port entry point too (CI review PR #340
+        follow-up P1)."""
+        from diff_diff._nprobust_port import lprobust
+
+        rng = np.random.default_rng(0)
+        G = 200
+        d = rng.uniform(0.0, 1.0, G)
+        y = d + rng.normal(0, 0.3, G)
+        cluster = np.array(
+            [i // 10 if i != 5 else None for i in range(G)], dtype=object
+        )
+        with pytest.raises(ValueError, match="cluster contains missing"):
+            lprobust(y, d, eval_point=0.0, h=0.3, b=0.3, cluster=cluster)
+
+    def test_lprobust_cluster_object_nan_raises(self):
+        """Port-level: object-dtype cluster with np.nan is rejected."""
+        from diff_diff._nprobust_port import lprobust
+
+        rng = np.random.default_rng(0)
+        G = 200
+        d = rng.uniform(0.0, 1.0, G)
+        y = d + rng.normal(0, 0.3, G)
+        cluster = np.array(
+            [i // 10 if i != 5 else np.nan for i in range(G)], dtype=object
+        )
+        with pytest.raises(ValueError, match="cluster contains missing"):
+            lprobust(y, d, eval_point=0.0, h=0.3, b=0.3, cluster=cluster)
+
     def test_lprobust_h_gt_b_selects_h_window(self):
         """When h > b, the active window is ind.h (lprobust.R:182 conditional
         replacement), not a union of ind.h and ind.b."""