From 541bb70d1f418563d8df7b87fec9312ad304d7c1 Mon Sep 17 00:00:00 2001 From: igerber Date: Mon, 13 Jul 2026 07:58:33 -0400 Subject: [PATCH] chore(types): mypy triage to enforceable zero + blocking CI job Tier 1 (real fixes): - 19 survey params typed `object` across 13 estimator files -> proper Optional[SurveyDesign]/Optional[ResolvedSurveyDesign] via TYPE_CHECKING imports; linalg's param typed as the SurveyDesign|ResolvedSurveyDesign union its isinstance branches actually handle - solve_ols overload chain unbroken (_validate_weights moved out from between the @overload block and the implementation) - stale return-arity annotations fixed: _aggregate_simple (2->3-tuple), _compute_cluster_psi_sums (2->3-tuple), plus Literal overloads for the flag-dependent returns of _compute_aggregated_se_with_wif and SunAbraham._compute_overall_att - cdh: float-keyed diagnostics dict annotated str-keyed; duplicate inner annotations dropped; a variable reused for two types in the bootstrap renamed; sentinel "_a11_warnings" str key in int-keyed dicts documented with scoped ignores - missing mixin attribute stub (TwoStageDiDBootstrapMixin.pretrends); efficient_did._unit_resolved_survey no longer inferred as type None - ~50 assert-based narrowings at flag-guarded blocks, each documenting a real invariant (all verified: 3,434 targeted tests pass, no assert fires) Tier 2/3 (documented suppressions): - prep_dgp per-module [index] override (seeded-DGP Optional covariate arrays; None-vs-array shape is RNG-stream-load-bearing) - matplotlib/plotly follow_imports=skip for local/CI stub parity - warn_unused_ignores=true; 23 already-stale ignores removed - reasoned inline ignores for mypy's union-explosion limitation CI: lint.yml gains a blocking Mypy job (pinned mypy + numpy/pandas/scipy, no package build); Lint Gate needs extended; TestLintWorkflowPinSync now covers the mypy pin. TODO.md Type Annotations note rewritten to the honest enforced-zero state; Large Module Files table regenerated (2026-07-13); override-tightening tracked as an Actionable row. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- .github/workflows/lint.yml | 32 ++++-- CHANGELOG.md | 15 +++ CLAUDE.md | 10 +- CONTRIBUTING.md | 15 +-- TODO.md | 99 ++++++++++++------- diff_diff/bacon.py | 7 +- diff_diff/bootstrap_chunking.py | 21 +++- diff_diff/bootstrap_utils.py | 6 +- diff_diff/business_report.py | 2 +- diff_diff/chaisemartin_dhaultfoeuille.py | 26 +++-- .../chaisemartin_dhaultfoeuille_bootstrap.py | 16 +-- diff_diff/continuous_did.py | 16 ++- diff_diff/diagnostic_report.py | 10 +- diff_diff/efficient_did.py | 16 ++- diff_diff/efficient_did_bootstrap.py | 9 +- diff_diff/efficient_did_covariates.py | 2 + diff_diff/estimators.py | 10 +- diff_diff/had.py | 2 +- diff_diff/had_pretests.py | 16 ++- diff_diff/imputation.py | 30 ++++-- diff_diff/linalg.py | 88 ++++++++++------- diff_diff/prep_dgp.py | 8 +- diff_diff/pretrends.py | 2 +- diff_diff/spillover.py | 28 ++++-- diff_diff/staggered.py | 15 ++- diff_diff/staggered_aggregation.py | 57 +++++++++-- diff_diff/staggered_bootstrap.py | 8 ++ diff_diff/staggered_triple_diff.py | 14 ++- diff_diff/sun_abraham.py | 51 ++++++++-- diff_diff/survey.py | 23 ++++- diff_diff/synthetic_control_results.py | 2 + diff_diff/synthetic_did.py | 11 ++- diff_diff/triple_diff.py | 8 +- diff_diff/twfe.py | 9 +- diff_diff/two_stage.py | 24 +++-- diff_diff/two_stage_bootstrap.py | 1 + diff_diff/utils.py | 2 +- diff_diff/wooldridge.py | 4 +- diff_diff/wooldridge_results.py | 4 +- pyproject.toml | 20 ++++ tests/test_openai_review.py | 37 +++++-- 41 files changed, 564 insertions(+), 212 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 9c30628ca..ffff12415 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -1,7 +1,7 @@ name: Lint -# Check-only lint enforcement (ruff + black). This workflow never modifies -# code — it reads the tree and passes or fails. +# Check-only lint + type enforcement (ruff + black + mypy). This workflow +# never modifies code — it reads the tree and passes or fails. # # Deliberately UNGATED (no ready-for-ci label check; precedent: # ai_pr_review.yml): these are cheap static checks that should give feedback @@ -45,15 +45,33 @@ jobs: - name: Black run: black --check diff_diff tests + typecheck: + name: Mypy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - name: Set up Python + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 + with: + python-version: '3.14' + # Runtime deps MUST be installed: with ignore_missing_imports=true an + # absent numpy silently becomes Any and most errors vanish — CI would + # go green while local runs stay red. Pinning them freezes stub drift. + # The package itself is NOT installed (mypy analyzes source directly; + # no maturin/Rust build ever enters this workflow). + - name: Install mypy + runtime deps for their type stubs (pinned) + run: pip install mypy==2.1.0 numpy==2.4.5 pandas==3.0.3 scipy==1.17.1 + - name: Mypy + run: mypy diff_diff + # Single stable required-check name so branch protection is configured once. - # Later additions (e.g. a mypy job) extend `needs` without any protection - # change. `if: always()` is load-bearing: without it a failed upstream job - # would leave this job SKIPPED, and GitHub treats a skipped required check - # as satisfied. + # New jobs extend `needs` without any protection change. `if: always()` is + # load-bearing: without it a failed upstream job would leave this job + # SKIPPED, and GitHub treats a skipped required check as satisfied. lint-gate: name: Lint Gate runs-on: ubuntu-latest - needs: [format-lint] + needs: [format-lint, typecheck] if: always() steps: - name: All lint jobs must succeed diff --git a/CHANGELOG.md b/CHANGELOG.md index b9669d72e..d895c27bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **Internal: mypy enforced at zero errors.** Triaged the 184 pre-existing + `mypy diff_diff` errors to an enforceable zero and added a blocking Mypy job to + the Lint CI workflow (pinned `mypy==2.1.0` + pinned numpy/pandas/scipy for stub + stability; the package itself is not installed). Real fixes include: 19 survey + parameters typed `object` across 13 estimator files (now properly typed via + `TYPE_CHECKING` imports), four stale return-arity annotations, a helper that + broke the `solve_ols` overload chain, a float-keyed diagnostics dict annotated + str-keyed, duplicate conflicting local annotations, a missing bootstrap-mixin + attribute stub, and ~50 explicit `assert ... is not None` narrowings at + flag-guarded blocks (documenting real invariants). Residual noise is contained + by documented suppressions (`prep_dgp` per-module `[index]` override, + `matplotlib`/`plotly` `follow_imports = "skip"`, a handful of reasoned inline + ignores) with `warn_unused_ignores = true` to keep them from going stale — 23 + already-stale ignores were removed. Tightening tracked in TODO.md. No public + API or numerical behavior change. - **Internal: rdrobust sharp-RD bandwidth-selection port (`diff_diff/_rdrobust_port.py`).** Step-1 machinery for the upcoming `RegressionDiscontinuity` estimator: a faithful pure-Python NumPy/SciPy port of R `rdrobust` 4.0.0's `rdbwselect` sharp-RD path (all 10 data-driven diff --git a/CLAUDE.md b/CLAUDE.md index b3edf413d..5ffa3bf31 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -32,11 +32,11 @@ mypy diff_diff ``` Lint/format/type tool versions are **pinned exactly** in the `dev` extra of -`pyproject.toml`; the `ruff`/`black` pins are mirrored in -`.github/workflows/lint.yml` (the ungated `Lint Gate` check runs `ruff check` -+ `black --check` on every PR push; sync enforced by `TestLintWorkflowPinSync`), -while `mypy` is pinned in `pyproject.toml` only until the type-check CI job -lands — update both surfaces together on any bump. Refresh local tools with `pip install -e ".[dev]"` +`pyproject.toml` and mirrored in `.github/workflows/lint.yml` (the ungated +`Lint Gate` check runs `ruff check` + `black --check` + `mypy diff_diff` on +every PR push; sync enforced by `TestLintWorkflowPinSync`) — update both +surfaces together on any bump. Mypy is enforced at ZERO errors; new code +must type-check cleanly. Refresh local tools with `pip install -e ".[dev]"` (the pinned tools need Python >= 3.10; the library floor stays 3.9). Some ruff rules are deliberately ignored per-file (`[tool.ruff.lint.per-file-ignores]`) — don't "fix" those patterns. One-time setup so `git blame` skips the 2026-07 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 62e344b7b..eb613c83a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,16 +8,17 @@ push, and its `Lint Gate` job is a required status check: ```bash ruff check diff_diff tests black diff_diff tests +mypy diff_diff ``` - Tool versions are **pinned exactly** in the `dev` extra of `pyproject.toml` - (`black`, `ruff`, `mypy`). The Lint workflow currently mirrors the `ruff` - and `black` pins (`.github/workflows/lint.yml`; the sync is enforced by - `TestLintWorkflowPinSync`) — `mypy` is pinned in `pyproject.toml` only, - until the planned type-check CI job lands. A version bump is a deliberate - PR updating both surfaces together. Refresh with `pip install -e ".[dev]"`. - The pinned tools require Python >= 3.10 (dev tooling only; the library - itself still supports Python 3.9). + (`black`, `ruff`, `mypy`) and mirrored in `.github/workflows/lint.yml` + (all three; the sync is enforced by `TestLintWorkflowPinSync`). `mypy + diff_diff` is enforced at zero errors by the Lint workflow's Mypy job. + A version bump is a deliberate PR updating both surfaces together. + Refresh with `pip install -e ".[dev]"`. The pinned tools require + Python >= 3.10 (dev tooling only; the library itself still supports + Python 3.9). - `[tool.ruff.lint.per-file-ignores]` entries are deliberate (trop logger-before-imports E402, honest_did math-notation `l` E741, `__init__` re-export F401, conftest import ordering E402). Do not "fix" those patterns diff --git a/TODO.md b/TODO.md index fab9a1923..2db92bae2 100644 --- a/TODO.md +++ b/TODO.md @@ -48,6 +48,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | ChangesInChanges/QDiD tutorial notebook (2x2 distributional walkthrough: QTE grid, interior range, uniform bands, CiC-vs-QDiD comparison) - deferred from the implementation PR as a documented decision. | `docs/tutorials/` | #682 | Mid | Low | +| Tighten the mypy suppressions that back the enforced-zero posture: burn down `prep_dgp`'s per-module `[index]` override (needs a None-vs-array restructure that preserves the seeded RNG stream), and evaluate re-enabling the globally disabled codes (`arg-type`, `return-value`, `var-annotated`, `assignment`) one at a time — `assignment` alone hid several real annotation drifts found during the 2026-07 triage. | `pyproject.toml` `[tool.mypy]`, `diff_diff/prep_dgp.py` | lint-CI | Mid | Low | | `practitioner_next_steps()` dedicated handler for `ChangesInChangesResults` (currently falls back to `_handle_generic`, which is safe; a dedicated handler is the established full-integration step, cf. HAD Phase 5). | `diff_diff/practitioner.py` | #682 | Quick | Low | --- @@ -173,46 +174,53 @@ For survey-specific limitations (NotImplementedError paths), see the Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates for splitting, 2000-3000 are monitored, 1000-2000 are accepted as a cohesion / scope trade-off. Updated -2026-05-15. +2026-07-13. | File | Lines | Action | |------|-------|--------| -| `chaisemartin_dhaultfoeuille.py` | 8636 | Consider splitting (per-path / placebos / survey IF / aggregation) | -| `had_pretests.py` | 4951 | Consider splitting (Stute / Yatchew / QUG / joint pretests) | -| `had.py` | 4593 | Consider splitting (continuous / mass-point / event-study / survey paths) | -| `staggered.py` | 3963 | Consider splitting — grew through survey + aggregation features | -| `linalg.py` | 3601 | Consider splitting (vcov surfaces) only if cohesion preserved — unified backend; vcov / solver paths tightly coupled | -| `diagnostic_report.py` | 3380 | Consider splitting (per-method renderers + provenance) | -| `power.py` | 3196 | Consider splitting (power analysis + MDE + sample size) | -| `synthetic_did.py` | 2819 | Monitor — variance methods + survey paths | -| `honest_did.py` | 2785 | Monitor | -| `business_report.py` | 2653 | Monitor — per-method narrative renderers | -| `imputation.py` | 2475 | Monitor | -| `survey.py` | 2466 | Monitor — grew with Phase 6 features | -| `utils.py` | 2396 | Monitor | -| `prep_dgp.py` | 2057 | Monitor | -| `triple_diff.py` | 2053 | Monitor | -| `estimators.py` | 1991 | Acceptable | -| `two_stage.py` | 1985 | Acceptable | -| `chaisemartin_dhaultfoeuille_results.py` | 1981 | Acceptable | -| `prep.py` | 1876 | Acceptable | -| `efficient_did.py` | 1793 | Acceptable | -| `sun_abraham.py` | 1713 | Acceptable | -| `continuous_did.py` | 1682 | Acceptable | -| `results.py` | 1676 | Acceptable | -| `staggered_triple_diff.py` | 1619 | Acceptable | -| `_nprobust_port.py` | 1412 | Acceptable | -| `practitioner.py` | 1402 | Acceptable | -| `trop_global.py` | 1350 | Acceptable | -| `trop_local.py` | 1339 | Acceptable | -| `local_linear.py` | 1332 | Acceptable | -| `wooldridge.py` | 1305 | Acceptable | +| `chaisemartin_dhaultfoeuille.py` | 8812 | Consider splitting (per-path / placebos / survey IF / aggregation) | +| `linalg.py` | 5424 | Consider splitting (vcov surfaces) only if cohesion preserved — unified backend; vcov / solver paths tightly coupled | +| `staggered.py` | 4992 | Consider splitting — grew through survey + aggregation features | +| `had.py` | 4748 | Consider splitting (continuous / mass-point / event-study / survey paths) | +| `had_pretests.py` | 4664 | Consider splitting (Stute / Yatchew / QUG / joint pretests) | +| `diagnostic_report.py` | 4135 | Consider splitting (per-method renderers + provenance) | +| `spillover.py` | 3655 | Consider splitting | +| `two_stage.py` | 3512 | Consider splitting | +| `power.py` | 3488 | Consider splitting (power analysis + MDE + sample size) | +| `utils.py` | 3483 | Consider splitting | +| `synthetic_control_results.py` | 3294 | Consider splitting | +| `honest_did.py` | 3068 | Consider splitting | +| `imputation.py` | 2898 | Monitor | +| `synthetic_did.py` | 2826 | Monitor — variance methods + survey paths | +| `business_report.py` | 2728 | Monitor — per-method narrative renderers | +| `survey.py` | 2681 | Monitor — grew with Phase 6 features | +| `synthetic_control.py` | 2526 | Monitor | +| `prep_dgp.py` | 2524 | Monitor | +| `estimators.py` | 2441 | Monitor | +| `continuous_did.py` | 2431 | Monitor | +| `sun_abraham.py` | 2314 | Monitor | +| `triple_diff.py` | 2231 | Monitor | +| `wooldridge.py` | 2192 | Monitor | +| `efficient_did.py` | 2083 | Monitor | +| `chaisemartin_dhaultfoeuille_results.py` | 2004 | Monitor | +| `results.py` | 1948 | Acceptable | +| `pretrends.py` | 1879 | Acceptable | +| `prep.py` | 1878 | Acceptable | +| `efficient_did_covariates.py` | 1818 | Acceptable | +| `staggered_triple_diff.py` | 1680 | Acceptable | +| `trop_local.py` | 1662 | Acceptable | +| `lpdid.py` | 1607 | Acceptable | +| `stacked_did.py` | 1589 | Acceptable | +| `practitioner.py` | 1511 | Acceptable | +| `_nprobust_port.py` | 1425 | Acceptable | +| `bacon.py` | 1376 | Acceptable | +| `local_linear.py` | 1325 | Acceptable | +| `trop_global.py` | 1298 | Acceptable | +| `staggered_aggregation.py` | 1204 | Acceptable | | `chaisemartin_dhaultfoeuille_bootstrap.py` | 1175 | Acceptable | -| `bacon.py` | 1144 | Acceptable | -| `pretrends.py` | 1133 | Acceptable | -| `stacked_did.py` | 1050 | Acceptable | -| `conley.py` | 1006 | Acceptable | -| `visualization/` | 4316 | Subpackage (split across 7 files) — OK | +| `conley.py` | 1140 | Acceptable | +| `trop.py` | 1026 | Acceptable | +| `profile.py` | 1001 | Acceptable | ### Standard Error Consistency @@ -233,8 +241,23 @@ IF/GMM estimators are tracked in ### Type Annotations -Mypy reports 0 errors. All mixin `attr-defined` errors resolved via `TYPE_CHECKING`-guarded -method stubs in the bootstrap mixin classes. +`mypy diff_diff` is **enforced at zero errors** by the Lint CI workflow's Mypy job at the +pinned `mypy==2.1.0` (ungated, every PR push; pins synced with the `dev` extra via +`TestLintWorkflowPinSync`). Zero is reached with documented suppressions — tightening them +is tracked in Actionable Backlog → Testing / docs: + +- Global `disable_error_code = ["arg-type", "return-value", "var-annotated", "assignment"]` + (numpy-stub compatibility, long-standing). +- Per-module override: `diff_diff.prep_dgp` disables `[index]` (seeded-DGP Optional + covariate arrays; see the comment in `pyproject.toml`). +- `matplotlib.*` / `plotly.*` use `follow_imports = "skip"` so local and CI runs see + identical `Any` surfaces regardless of what plotting stubs are installed. +- A handful of reasoned inline `# type: ignore[code]` comments, each with a why-comment; + `warn_unused_ignores = true` prevents them from going stale. + +Mixin cross-class attribute access uses `TYPE_CHECKING`-guarded attribute/method stubs in +the bootstrap mixin classes; keep stubs in sync with implementations (stub drift shows up +as `[misc]` unpack-arity errors). ### Test Coverage diff --git a/diff_diff/bacon.py b/diff_diff/bacon.py index 28b2b5c98..eaba1014e 100644 --- a/diff_diff/bacon.py +++ b/diff_diff/bacon.py @@ -12,7 +12,7 @@ import warnings from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -21,6 +21,9 @@ from diff_diff.utils import pre_demean_norms, snap_absorbed_regressors from diff_diff.utils import within_transform as _within_transform_util +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + @dataclass class Comparison2x2: @@ -1253,7 +1256,7 @@ def bacon_decompose( time: str, first_treat: str, weights: str = "exact", - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> BaconDecompositionResults: """ Convenience function for Goodman-Bacon decomposition. diff --git a/diff_diff/bootstrap_chunking.py b/diff_diff/bootstrap_chunking.py index 0b873ed1b..d0cc0997b 100644 --- a/diff_diff/bootstrap_chunking.py +++ b/diff_diff/bootstrap_chunking.py @@ -43,13 +43,26 @@ from __future__ import annotations from collections import abc -from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Callable, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, + Union, +) import numpy as np from diff_diff._backend import HAS_RUST_BACKEND, _rust_bootstrap_weights from diff_diff.bootstrap_utils import generate_bootstrap_weights_batch_numpy +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign + # Byte ceiling for a single ``(B, n_units)`` float64 weight block. 256 MB keeps # the live intermediate small at millions of units while staying large enough # that the per-block matmuls remain BLAS-efficient and chunk overhead (a handful @@ -118,7 +131,7 @@ def iter_weight_blocks( def iter_survey_multiplier_weight_blocks( n_bootstrap: int, - resolved_survey: object, + resolved_survey: Optional["ResolvedSurveyDesign"], weight_type: str, rng: np.random.Generator, *, @@ -147,10 +160,12 @@ def iter_survey_multiplier_weight_blocks( if block_size < 1: raise ValueError(f"block_size must be >= 1, got {block_size}") + # Callers only reach the survey-multiplier path with a resolved design. + assert resolved_survey is not None psu = getattr(resolved_survey, "psu", None) strata = getattr(resolved_survey, "strata", None) if psu is None: - n_psu = len(resolved_survey.weights) # type: ignore[attr-defined] + n_psu = len(resolved_survey.weights) psu_ids = np.arange(n_psu) else: psu_ids = np.unique(psu) diff --git a/diff_diff/bootstrap_utils.py b/diff_diff/bootstrap_utils.py index 9e09036c2..c94b38b9d 100644 --- a/diff_diff/bootstrap_utils.py +++ b/diff_diff/bootstrap_utils.py @@ -774,9 +774,9 @@ def apply_stratum_centering( # the contribution is zero by construction). Skip to # avoid a divide-by-zero on sqrt(n_h / 0). continue - slc = [slice(None)] * tensor.ndim - slc[psu_axis] = mask_h - slc = tuple(slc) + slc_list: list = [slice(None)] * tensor.ndim + slc_list[psu_axis] = mask_h + slc = tuple(slc_list) tensor[slc] -= tensor[slc].mean(axis=psu_axis, keepdims=True) tensor[slc] *= np.sqrt(n_h / (n_h - 1)) else: diff --git a/diff_diff/business_report.py b/diff_diff/business_report.py index f8331ac7c..19c967146 100644 --- a/diff_diff/business_report.py +++ b/diff_diff/business_report.py @@ -1428,7 +1428,7 @@ def _describe_assumption(estimator_name: str, results: Any = None) -> Dict[str, "event's post-window, so future-treated units can serve " "as controls for earlier events)" ) - block: Dict[str, Any] = { + block = { "parallel_trends_variant": "stacked_sub_experiment", "no_anticipation": True, "description": ( diff --git a/diff_diff/chaisemartin_dhaultfoeuille.py b/diff_diff/chaisemartin_dhaultfoeuille.py index 4b0ab116b..c78a1923f 100644 --- a/diff_diff/chaisemartin_dhaultfoeuille.py +++ b/diff_diff/chaisemartin_dhaultfoeuille.py @@ -2239,7 +2239,7 @@ def fit( set_ids=set_ids_arr, ) # Surface A11 warnings from multi-horizon computation - mh_a11 = multi_horizon_dids.pop("_a11_warnings", None) + mh_a11 = multi_horizon_dids.pop("_a11_warnings", None) # type: ignore[call-overload] # sentinel str key in int-keyed dict if mh_a11: warnings.warn( f"Multi-horizon control-availability violations in " @@ -2301,6 +2301,8 @@ def fit( # tensor is None and the bootstrap has no cell-level path # to take). _mh_pp_cache: Dict[int, np.ndarray] = {} + # multi_horizon_if is computed unconditionally in this branch. + assert multi_horizon_if is not None # Compute inference for ALL horizons 1..L_max (including l=1) # so the event_study_effects dict uses a consistent estimand # (per-group DID_{g,l}) across all horizons. @@ -2461,7 +2463,7 @@ def fit( set_ids=set_ids_arr, ) # Surface placebo A11 warnings - pl_a11 = multi_horizon_placebos.pop("_a11_warnings", None) + pl_a11 = multi_horizon_placebos.pop("_a11_warnings", None) # type: ignore[call-overload] # sentinel str key in int-keyed dict if pl_a11: warnings.warn( f"Multi-horizon placebo control-availability " @@ -2490,8 +2492,10 @@ def fit( ) # Per-placebo-horizon analytical SE via cohort recentering # (same pattern as positive-horizon SE at Step 12c). - placebo_horizon_se: Dict[int, float] = {} - placebo_horizon_inference: Dict[int, Dict[str, Any]] = {} + placebo_horizon_se = {} + placebo_horizon_inference = {} + # placebo_horizon_if is computed unconditionally in this branch. + assert placebo_horizon_if is not None singleton_baseline_set_pl = set(singleton_baseline_groups) eligible_mask_pl = np.array( [g not in singleton_baseline_set_pl for g in all_groups], @@ -3951,7 +3955,7 @@ def fit( # predict_het + survey continues to work via the existing # _compute_path_heterogeneity_test → _compute_heterogeneity_test # forward path. - _path_het_placebo = L_max if self.placebo else 0 + _path_het_placebo = L_max if (self.placebo and L_max is not None) else 0 if _path_het_placebo > 0 and _obs_survey_info is not None: _path_het_placebo = 0 path_heterogeneity_effects = _compute_path_heterogeneity_test( @@ -4758,7 +4762,7 @@ def _compute_covariate_residualization( n_groups, n_periods = Y_mat.shape n_covariates = X_cell.shape[2] Y_resid = Y_mat.copy() - diagnostics: Dict[str, Any] = {} + diagnostics: Dict[float, Dict[str, Any]] = {} failed_baselines: set = set() # Pre-compute observation validity masks for first-differencing. @@ -5078,6 +5082,8 @@ def _compute_heterogeneity_test( # the explicit raise. use_survey = obs_survey_info is not None and group_ids_order is not None if use_survey: + # The flag definition above guarantees these (mypy can't track it). + assert obs_survey_info is not None and group_ids_order is not None from diff_diff.survey import ( compute_replicate_if_variance, compute_survey_if_variance, @@ -5375,6 +5381,8 @@ def _compute_heterogeneity_test( else: df_s_local = min(int(df_s), int(n_valid_het) - 1) else: + # Survey heterogeneity path: obs-level info is present. + assert obs_survey_info is not None obs_tids = np.asarray(obs_survey_info["time_ids"]) periods_arr = np.asarray(obs_survey_info["periods"]) psi_obs = np.zeros(len(obs_w_raw), dtype=np.float64) @@ -5781,7 +5789,7 @@ def _compute_multi_horizon_dids( # Attach A11 warnings to the results for the caller to surface if a11_multi_warnings: - results["_a11_warnings"] = a11_multi_warnings # type: ignore[assignment] + results["_a11_warnings"] = a11_multi_warnings # type: ignore[index] # sentinel str key in int-keyed dict return results @@ -7272,7 +7280,7 @@ def _compute_multi_horizon_placebos( } if a11_placebo_warnings: - results["_a11_warnings"] = a11_placebo_warnings # type: ignore[assignment] + results["_a11_warnings"] = a11_placebo_warnings # type: ignore[index] # sentinel str key in int-keyed dict return results @@ -8340,6 +8348,8 @@ def _survey_se_from_group_if( use_cell_allocator = False if use_cell_allocator: + # The flag definition above guarantees these (mypy can't track it). + assert U_centered_per_period is not None and time_ids is not None and periods is not None tids_eff = np.asarray(time_ids)[pos_mask] # Map row's group to an index in eligible_groups (−1 when the # group is ineligible — singleton-baseline exclusion drops it). diff --git a/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py b/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py index a63aad078..a37b57069 100644 --- a/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py +++ b/diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py @@ -812,16 +812,16 @@ def _compute_dcdh_bootstrap( for path_key, horizon_inputs in path_bootstrap_inputs.items(): bs_ses_for_path = results.path_ses.get(path_key, {}) - valid_horizons = [] + valid_horizon_stats = [] for l_h, (u_h, n_h, eff_h, _u_pp_h) in sorted(horizon_inputs.items()): if u_h.size == 0 or n_h <= 0: continue bs_se = bs_ses_for_path.get(l_h, np.nan) if not np.isfinite(bs_se) or bs_se <= 0: continue - valid_horizons.append((l_h, u_h, n_h, eff_h, bs_se)) + valid_horizon_stats.append((l_h, u_h, n_h, eff_h, bs_se)) - if len(valid_horizons) < 2: + if len(valid_horizon_stats) < 2: continue # All horizons within a path use the same n_eligible @@ -829,7 +829,7 @@ def _compute_dcdh_bootstrap( # _collect_path_bootstrap_inputs's use of # eligible_mask_var for cohort-recentering); use the # first valid horizon's IF size as the shared dim. - n_dim = valid_horizons[0][1].size + n_dim = valid_horizon_stats[0][1].size map_path = _map_for_target( n_dim, group_id_to_psu_code, @@ -844,12 +844,12 @@ def _compute_dcdh_bootstrap( group_to_psu_map=map_path, ) es_dists_path = [] - for _l_h, u_h, n_h, eff_h, _bs_se in valid_horizons: + for _l_h, u_h, n_h, eff_h, _bs_se in valid_horizon_stats: deviations = (shared_weights @ u_h) / n_h es_dists_path.append(eff_h + deviations) boot_matrix = np.asarray(es_dists_path) - effects_vec = np.array([v[3] for v in valid_horizons]) - ses_vec = np.array([v[4] for v in valid_horizons]) + effects_vec = np.array([v[3] for v in valid_horizon_stats]) + ses_vec = np.array([v[4] for v in valid_horizon_stats]) t_stats = np.abs((boot_matrix - effects_vec[:, None]) / ses_vec[:, None]) sup_t_dist = np.max(t_stats, axis=0) finite_mask = np.isfinite(sup_t_dist) @@ -861,7 +861,7 @@ def _compute_dcdh_bootstrap( continue path_cband_crits[path_key] = crit_p - path_cband_n_valid[path_key] = len(valid_horizons) + path_cband_n_valid[path_key] = len(valid_horizon_stats) results.path_cband_crit_values = path_cband_crits results.path_cband_n_valid_horizons = path_cband_n_valid diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index ead130097..547190d5e 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -10,7 +10,7 @@ """ import warnings -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -42,6 +42,10 @@ ) from diff_diff.utils import safe_inference +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign + + __all__ = ["ContinuousDiD", "ContinuousDiDResults", "DoseResponseCurve"] @@ -281,7 +285,7 @@ def fit( first_treat: str, dose: str, aggregate: Optional[str] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> ContinuousDiDResults: """ Fit the continuous DiD estimator. @@ -1472,7 +1476,7 @@ def _compute_dose_response_gt( degree: int, dvals: np.ndarray, survey_weights: Optional[np.ndarray] = None, - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, levels: Optional[np.ndarray] = None, ) -> Optional[Dict[str, Any]]: """Compute dose-response for a single (g,t) cell. @@ -1930,7 +1934,7 @@ def _compute_analytical_se( dvals: np.ndarray, agg_att_d: np.ndarray, agg_acrt_d: np.ndarray, - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, ) -> Dict[str, Any]: """Compute analytical SEs using influence functions.""" n_units = precomp["n_units"] @@ -2130,7 +2134,7 @@ def _run_bootstrap( original_att_d: np.ndarray, original_acrt_d: np.ndarray, event_study_effects: Optional[Dict[int, Dict]], - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, ) -> Dict[str, Any]: """Run multiplier bootstrap inference.""" if self.n_bootstrap < 50: @@ -2178,6 +2182,8 @@ def _run_bootstrap( generate_survey_multiplier_weights_batch, ) + # The survey bootstrap branch always has a resolved design. + assert unit_resolved is not None psu_weights, psu_ids = generate_survey_multiplier_weights_batch( self.n_bootstrap, unit_resolved, self.bootstrap_weights, rng ) diff --git a/diff_diff/diagnostic_report.py b/diff_diff/diagnostic_report.py index 7b8cf2f64..c9a4261d3 100644 --- a/diff_diff/diagnostic_report.py +++ b/diff_diff/diagnostic_report.py @@ -2873,7 +2873,7 @@ def _scalars_from_mapping(mapping: Any) -> List[float]: return [] else: try: - values = list(mapping) # type: ignore[arg-type] + values = list(mapping) except Exception: # noqa: BLE001 return [] for val in values: @@ -3787,6 +3787,8 @@ def _render_overall_interpretation(schema: Dict[str, Any], labels: Dict[str, str "or a survey-design collapse before interpreting." ) elif val_finite: + # val_finite includes the isinstance check (mypy can't track flags). + assert isinstance(val, (int, float)) direction = "increased" if val > 0 else "decreased" if val < 0 else "did not change" # Use the headline's own alpha rather than hardcoding 95 so prose # stays consistent with the rendered interval when alpha != 0.05. @@ -3800,7 +3802,11 @@ def _render_overall_interpretation(schema: Dict[str, Any], labels: Dict[str, str and len(ci) == 2 and all(isinstance(v, (int, float)) and np.isfinite(v) for v in ci) ) - ci_str = f" ({ci_level}% CI: {ci[0]:.3g} to {ci[1]:.3g})" if ci_finite else "" + ci_str = ( + f" ({ci_level}% CI: {ci[0]:.3g} to {ci[1]:.3g})" + if ci_finite and isinstance(ci, (list, tuple)) + else "" + ) p_str = f", p = {p:.3g}" if isinstance(p, (int, float)) and np.isfinite(p) else "" sentences.append( f"On {est}, {treatment} {direction} {outcome} by {val:.3g}{ci_str}{p_str}." diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 29861c97a..35e97648b 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -23,7 +23,7 @@ """ import warnings -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -32,6 +32,9 @@ EDiDBootstrapResults, EfficientDiDBootstrapMixin, ) + +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign from diff_diff.efficient_did_covariates import ( OMEGA_RIDGE_DEFAULT, _silverman_bandwidth, @@ -343,7 +346,7 @@ def __init__( self.omega_ridge = omega_ridge self.is_fitted_ = False self.results_: Optional[EfficientDiDResults] = None - self._unit_resolved_survey = None + self._unit_resolved_survey: Optional["ResolvedSurveyDesign"] = None self._validate_params() def _validate_params(self) -> None: @@ -744,6 +747,7 @@ def fit( if resolved_survey is not None: # Use the resolved survey's weights (already normalized per weight_type) # subset to unit level via _unit_first_panel_row (aligned to all_units) + assert self._unit_resolved_survey is not None unit_level_weights = self._unit_resolved_survey.weights self._unit_level_weights = unit_level_weights @@ -1057,7 +1061,7 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any "y1_col": effective_p1_col, } ) - group_time_effects[(g, t)] = None # type: ignore[assignment] + group_time_effects[(g, t)] = None continue # ----- Legacy (omega_ridge=0) covariate path ----- @@ -1433,6 +1437,8 @@ def _compute_survey_eif_se(self, eif_vals: np.ndarray) -> float: once in ``fit()``, ensuring consistent unit-level arrays and avoiding repeated subsetting of panel-level survey data. """ + # Built once in fit() before any call lands here (see docstring). + assert self._unit_resolved_survey is not None if self._unit_resolved_survey.uses_replicate_variance: from diff_diff.survey import compute_replicate_if_variance @@ -1595,6 +1601,8 @@ def _aggregate_overall( # the survey-weighted WIF term). Dispatch replicate vs TSL. if self._unit_resolved_survey is not None: uw = self._unit_level_weights + # Set together with _unit_resolved_survey in fit(). + assert uw is not None total_w = float(np.sum(uw)) psi_total = uw * agg_eif / total_w + wif / total_w @@ -1713,6 +1721,8 @@ def _aggregate_event_study( if self._unit_resolved_survey is not None: uw = self._unit_level_weights + # Set together with _unit_resolved_survey in fit(). + assert uw is not None total_w = float(np.sum(uw)) psi_total = uw * agg_eif / total_w + wif_e / total_w diff --git a/diff_diff/efficient_did_bootstrap.py b/diff_diff/efficient_did_bootstrap.py index 5e4cca75e..9f0089d02 100644 --- a/diff_diff/efficient_did_bootstrap.py +++ b/diff_diff/efficient_did_bootstrap.py @@ -8,7 +8,7 @@ import warnings from dataclasses import dataclass, field -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple import numpy as np @@ -23,6 +23,9 @@ compute_effect_bootstrap_stats as _compute_effect_bootstrap_stats_func, ) +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign + @dataclass class EDiDBootstrapResults: @@ -66,7 +69,7 @@ def _run_multiplier_bootstrap( cohort_fractions: Dict[float, float], cluster_indices: Optional[np.ndarray] = None, n_clusters: Optional[int] = None, - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, unit_level_weights: Optional[np.ndarray] = None, ) -> EDiDBootstrapResults: """Run multiplier bootstrap on stored EIF values. @@ -117,6 +120,8 @@ def _run_multiplier_bootstrap( ) if _use_survey_bootstrap: + # The flag definition above guarantees this (mypy can't track it). + assert resolved_survey is not None # PSU-level multiplier weights, generated and expanded one draw-block # at a time (unstratified designs tile the generation; stratified # designs have few PSUs and fall back to full generation + slicing). diff --git a/diff_diff/efficient_did_covariates.py b/diff_diff/efficient_did_covariates.py index 2611c0ac0..f228d06ee 100644 --- a/diff_diff/efficient_did_covariates.py +++ b/diff_diff/efficient_did_covariates.py @@ -644,6 +644,8 @@ def estimate_inverse_propensity_sieve( # Unweighted: (Psi_gp' Psi_gp) beta = Psi_all.sum(axis=0) # Weighted: (Psi_gp' W_group Psi_gp) beta = (w_all * Psi_all).sum(axis=0) if w_group is not None: + # w_group and unit_weights are supplied together (survey path). + assert unit_weights is not None A = Psi_gp.T @ (w_group[:, None] * Psi_gp) b = (unit_weights[:, None] * basis_all).sum(axis=0) else: diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 2fb01ba68..3e01827cb 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -630,6 +630,8 @@ def fit( # binary `time` column is integer 0/1 by convention and is unaffected # by the normalization. if _fit_vcov_type == "conley": + # Validated by the conley front-door (_validate_conley_estimator_inputs). + assert self.conley_coords is not None _conley_coords_arr: Optional[np.ndarray] = np.column_stack( [ data[self.conley_coords[0]].values.astype(np.float64), @@ -736,6 +738,8 @@ def _refit_did_absorb(w_r): if survey_metadata and survey_metadata.df_survey else 0 # rank-deficient replicate → NaN inference ) + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None if _n_valid_rep < resolved_survey.n_replicates: _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 if survey_metadata is not None: @@ -1934,6 +1938,8 @@ def fit( # type: ignore[override] # coords + time/unit). When vcov_type != "conley", these are silently # ignored downstream (Phase 1 / 2 convention). if _fit_vcov_type == "conley": + # Validated by the conley front-door (_validate_conley_estimator_inputs). + assert self.conley_coords is not None _conley_coords_arr: Optional[np.ndarray] = np.column_stack( [ data[self.conley_coords[0]].values.astype(np.float64), @@ -1953,7 +1959,7 @@ def fit( # type: ignore[override] # Note: Wild bootstrap for multi-period effects is complex (multiple coefficients) # For now, we use analytical inference even if inference="wild_bootstrap" - coefficients, residuals, fitted, vcov = solve_ols( + coefficients, residuals, fitted, vcov = solve_ols( # type: ignore[call-overload, misc] # mypy gives up on the Optional-arg union explosion ("Not all union combinations were tried") X, y, return_fitted=True, @@ -2149,6 +2155,8 @@ def _refit_mp_absorb(w_r): df = resolved_survey.df_survey # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 if _uses_replicate_mp: + # The flag definition above guarantees this (mypy can't track it). + assert resolved_survey is not None if resolved_survey.df_survey is None: df = 0 # rank-deficient replicate → NaN inference if _n_valid_rep_mp is not None and _n_valid_rep_mp < resolved_survey.n_replicates: diff --git a/diff_diff/had.py b/diff_diff/had.py index 4f173e74e..b1207dcd8 100644 --- a/diff_diff/had.py +++ b/diff_diff/had.py @@ -3220,7 +3220,7 @@ def fit( # narrow via isinstance (or by the aggregate they passed) to # access aggregate-specific fields. if aggregate == "event_study": - return self._fit_event_study( # type: ignore[return-value] + return self._fit_event_study( data=data, outcome_col=outcome_col, dose_col=dose_col, diff --git a/diff_diff/had_pretests.py b/diff_diff/had_pretests.py index 30570ac4d..e1008c25d 100644 --- a/diff_diff/had_pretests.py +++ b/diff_diff/had_pretests.py @@ -4104,10 +4104,14 @@ def _compose_verdict_overall_survey( no rejection -> fail-to-reject string. All branches end with `_QUG_DEFERRED_SUFFIX`. """ - stute_ok = stute is not None and bool(np.isfinite(stute.p_value)) - yatchew_ok = yatchew is not None and bool(np.isfinite(yatchew.p_value)) - stute_rej = stute_ok and bool(stute.reject) - yatchew_rej = yatchew_ok and bool(yatchew.reject) + if stute is not None and bool(np.isfinite(stute.p_value)): + stute_ok, stute_rej = True, bool(stute.reject) + else: + stute_ok, stute_rej = False, False + if yatchew is not None and bool(np.isfinite(yatchew.p_value)): + yatchew_ok, yatchew_rej = True, bool(yatchew.reject) + else: + yatchew_ok, yatchew_rej = False, False reasons = [] if stute_rej: @@ -4541,6 +4545,8 @@ def did_had_pretest_workflow( # starting with "inconclusive" even when all_pass=True). verdict = _compose_verdict_event_study_survey(pretrends_joint, homogeneity_joint) else: + # use_survey_path=False branch: qug_res was computed above. + assert qug_res is not None qug_ok = bool(np.isfinite(qug_res.p_value)) all_pass = bool( qug_ok @@ -4639,6 +4645,8 @@ def did_had_pretest_workflow( # R7 P1 fix: explicit survey-aware verdict composer. verdict = _compose_verdict_overall_survey(stute_res, yatchew_res) else: + # use_survey_path=False branch: qug_res was computed above. + assert qug_res is not None qug_conclusive = bool(np.isfinite(qug_res.p_value)) any_reject = qug_res.reject or stute_res.reject or yatchew_res.reject all_pass = bool(qug_conclusive and linearity_conclusive and not any_reject) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 5275411f9..52692b67d 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -24,7 +24,7 @@ """ import warnings -from typing import Any, Callable, Dict, List, NamedTuple, Optional, Set, Tuple +from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Set, Tuple import numpy as np import pandas as pd @@ -45,6 +45,9 @@ snap_absorbed_regressors, ) +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + class _UntreatedProjection(NamedTuple): """Cached, target-invariant pieces of the untreated imputation projection @@ -322,7 +325,7 @@ def fit( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> ImputationDiDResults: """ Fit the imputation DiD estimator. @@ -583,6 +586,8 @@ def fit( if resolved_survey.psu is not None and survey_metadata is not None: from diff_diff.survey import compute_survey_metadata + # resolved_survey non-None implies survey_design was passed. + assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) if survey_design.weights @@ -784,14 +789,15 @@ def fit( _cohorts_treated = df.loc[omega_1_mask, first_treat].values # Derive keys from actual outputs (excludes filtered/Prop5/ref) + _es_effects = event_study_effects or {} + _grp_effects = group_effects or {} _sorted_rel_times = sorted( e - for e in (event_study_effects or {}).keys() - if np.isfinite(event_study_effects[e]["effect"]) - and event_study_effects[e].get("n_obs", 1) > 0 + for e in _es_effects.keys() + if np.isfinite(_es_effects[e]["effect"]) and _es_effects[e].get("n_obs", 1) > 0 ) _sorted_groups = sorted( - g for g in (group_effects or {}).keys() if np.isfinite(group_effects[g]["effect"]) + g for g in _grp_effects.keys() if np.isfinite(_grp_effects[g]["effect"]) ) _n_es = len(_sorted_rel_times) @@ -858,8 +864,8 @@ def _refit_imp(w_r): # Build full-sample estimate from actual effects _full_est = [overall_att] - _full_est.extend([event_study_effects[e]["effect"] for e in _sorted_rel_times]) - _full_est.extend([group_effects[g]["effect"] for g in _sorted_groups]) + _full_est.extend([_es_effects[e]["effect"] for e in _sorted_rel_times]) + _full_est.extend([_grp_effects[g]["effect"] for g in _sorted_groups]) _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( _refit_imp, np.array(_full_est), resolved_survey @@ -867,6 +873,8 @@ def _refit_imp(w_r): overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) # Override df if replicates were dropped + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None if _n_valid_rep_imp < resolved_survey.n_replicates: _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 if survey_metadata is not None: @@ -1383,7 +1391,7 @@ def _compute_cluster_psi_sums( kept_cov_mask: Optional[np.ndarray] = None, survey_weights_0: Optional[np.ndarray] = None, proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, - ) -> Tuple[np.ndarray, np.ndarray]: + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Compute cluster-level influence function sums (Theorem 3). @@ -2464,6 +2472,8 @@ def _compute_lead_coefficients( # Zero-pad to full panel length (subpopulation approach): # observations outside Omega_0 contribute zero to the score, # but preserve PSU/strata structure for design-based variance. + # The survey full-design path always supplies the full obs count. + assert n_obs_full is not None n_full_obs = n_obs_full k_vcov = X_for_vcov.shape[1] X_full = np.zeros((n_full_obs, k_vcov), dtype=np.float64) @@ -2817,7 +2827,7 @@ def imputation_did( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, vcov_type: str = "hc1", **kwargs, ) -> ImputationDiDResults: diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 9aeb7ab9a..e541ff2e1 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -35,7 +35,7 @@ import os import warnings from dataclasses import dataclass -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, overload import numpy as np import pandas as pd @@ -63,6 +63,10 @@ _validate_conley_kwargs, ) +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign + + # ============================================================================= # Utility Functions # ============================================================================= @@ -869,6 +873,40 @@ def _nonfinite_vcov_needs_python_rerun( return not bool(np.all(np.isfinite(vcov))) +_VALID_WEIGHT_TYPES = {"pweight", "fweight", "aweight"} + + +def _validate_weights(weights, weight_type, n): + """Validate weights array and weight_type for solve_ols/LinearRegression.""" + if weight_type not in _VALID_WEIGHT_TYPES: + raise ValueError( + f"weight_type must be one of {_VALID_WEIGHT_TYPES}, " f"got '{weight_type}'" + ) + if weights is not None: + weights = np.asarray(weights, dtype=np.float64) + if weights.shape[0] != n: + raise ValueError(f"weights length ({weights.shape[0]}) must match " f"X rows ({n})") + if np.any(np.isnan(weights)): + raise ValueError("Weights contain NaN values") + if np.any(np.isinf(weights)): + raise ValueError("Weights contain Inf values") + if np.any(weights < 0): + raise ValueError("Weights must be non-negative") + if np.sum(weights) <= 0: + raise ValueError( + "Weights sum to zero — no observations have positive weight. " + "Cannot fit a model on an empty effective sample." + ) + if weight_type == "fweight": + fractional = weights - np.round(weights) + if np.any(np.abs(fractional) > 1e-10): + raise ValueError( + "Frequency weights (fweight) must be non-negative integers. " + "Fractional values detected. Use pweight for non-integer weights." + ) + return weights + + @overload def solve_ols( X: np.ndarray, @@ -950,40 +988,6 @@ def solve_ols( ]: ... -_VALID_WEIGHT_TYPES = {"pweight", "fweight", "aweight"} - - -def _validate_weights(weights, weight_type, n): - """Validate weights array and weight_type for solve_ols/LinearRegression.""" - if weight_type not in _VALID_WEIGHT_TYPES: - raise ValueError( - f"weight_type must be one of {_VALID_WEIGHT_TYPES}, " f"got '{weight_type}'" - ) - if weights is not None: - weights = np.asarray(weights, dtype=np.float64) - if weights.shape[0] != n: - raise ValueError(f"weights length ({weights.shape[0]}) must match " f"X rows ({n})") - if np.any(np.isnan(weights)): - raise ValueError("Weights contain NaN values") - if np.any(np.isinf(weights)): - raise ValueError("Weights contain Inf values") - if np.any(weights < 0): - raise ValueError("Weights must be non-negative") - if np.sum(weights) <= 0: - raise ValueError( - "Weights sum to zero — no observations have positive weight. " - "Cannot fit a model on an empty effective sample." - ) - if weight_type == "fweight": - fractional = weights - np.round(weights) - if np.any(np.abs(fractional) > 1e-10): - raise ValueError( - "Frequency weights (fweight) must be non-negative integers. " - "Fractional values detected. Use pweight for non-integer weights." - ) - return weights - - def solve_ols( X: np.ndarray, y: np.ndarray, @@ -1415,8 +1419,10 @@ def solve_ols( # be computed on original X and residuals with weights applied exactly once. if _original_X is not None and _original_y is not None: if return_fitted: + assert len(result) == 4 coefficients, _resid_w, _fitted_w, vcov_out = result else: + assert len(result) == 3 coefficients, _resid_w, vcov_out = result # Handle rank-deficient case: use only identified columns for fitted values @@ -2459,6 +2465,7 @@ def _compute_cr2_bm_vcov_and_dof( # dense pseudoinverse convention verbatim. A_g_small = _cr2_adjustment_matrix(np.eye(n_g) - U_g @ U_g.T) if residuals is not None: + assert cluster_scores is not None cluster_scores[gi] = X_g.T @ (A_g_small @ residuals[idx_g]) A_g_Xbi[g] = A_g_small @ B_g continue @@ -2472,11 +2479,13 @@ def _compute_cr2_bm_vcov_and_dof( gamma[pseudo] = -1.0 / lam[pseudo] UQ = U_g @ Q_g if residuals is not None: + assert cluster_scores is not None u_g = residuals[idx_g] # s_g = X_g' A_g u_g = X_g'u_g + (X_g'UQ) diag(gamma) (UQ'u_g) cluster_scores[gi] = X_g.T @ u_g + (X_g.T @ UQ) @ (gamma * (UQ.T @ u_g)) A_g_Xbi[g] = B_g + UQ @ (gamma[:, np.newaxis] * (UQ.T @ B_g)) if residuals is not None: + assert cluster_scores is not None meat = cluster_scores.T @ cluster_scores vcov = M_U @ meat @ M_U else: @@ -3292,7 +3301,7 @@ def _compute_robust_vcov_numpy( X, residuals, np.asarray(conley_coords, dtype=np.float64), - float(conley_cutoff_km), # type: ignore[arg-type] + float(conley_cutoff_km), conley_metric, conley_kernel, bread_matrix, @@ -4251,7 +4260,7 @@ def __init__( rank_deficient_action: str = "warn", weights: Optional[np.ndarray] = None, weight_type: str = "pweight", - survey_design: object = None, + survey_design: Optional[Union["SurveyDesign", "ResolvedSurveyDesign"]] = None, vcov_type: Optional[str] = None, conley_coords: Optional[np.ndarray] = None, conley_cutoff_km: Optional[float] = None, @@ -4626,6 +4635,9 @@ def fit( ) if _uses_rep: + # The flag definition above guarantees this (mypy can't + # propagate isinstance narrowing through the flag variable). + assert isinstance(_effective_survey_design, _RSD) from diff_diff.survey import compute_replicate_vcov nan_mask = np.isnan(coefficients) @@ -4721,6 +4733,8 @@ def fit( and _fit_vcov_type in ("classical", "hc1") ): _fe_scale = _absorbed_fe_vcov_scale(n_eff_df, self.n_params_effective_, df_adjustment) + # vcov_ was computed above on every path reaching this scale-fix. + assert self.vcov_ is not None if np.isnan(_fe_scale): self.vcov_ = np.full_like(self.vcov_, np.nan) elif _fe_scale != 1.0: diff --git a/diff_diff/prep_dgp.py b/diff_diff/prep_dgp.py index 13c4e0579..812c7d358 100644 --- a/diff_diff/prep_dgp.py +++ b/diff_diff/prep_dgp.py @@ -2031,7 +2031,7 @@ def generate_survey_did_data( elif panel and t == 1: _panel_unit_fe = unit_fe # save for reuse elif panel and t > 1: - unit_fe = _panel_unit_fe # type: ignore[possibly-undefined] + unit_fe = _panel_unit_fe # Cross-section informative sampling: re-rank weights each period if informative_sampling and not panel: @@ -2041,7 +2041,7 @@ def generate_survey_did_data( if conditional_pt != 0.0: x1[unit_cohort > 0] += 1.0 x2 = rng.choice([0, 1], size=n_units) - unit_weight = _base_weight.copy() # type: ignore[possibly-undefined] + unit_weight = _base_weight.copy() y0_t = unit_fe + psu_re[unit_psu] + psu_period_re[unit_psu, t - 1] + 0.5 * t if add_covariates: y0_t = y0_t + _beta1 * x1 + _beta2 * x2 @@ -2064,8 +2064,8 @@ def generate_survey_did_data( x1 = None x2 = None if not informative_sampling and panel and t > 1 and add_covariates: - x1 = _panel_x1 # type: ignore[possibly-undefined] - x2 = _panel_x2 # type: ignore[possibly-undefined] + x1 = _panel_x1 + x2 = _panel_x2 elif not informative_sampling and panel and t == 1 and add_covariates: _panel_x1 = x1 _panel_x2 = x2 diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index fbc68b092..0cb20cdef 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -64,7 +64,7 @@ def _compute_nis_acceptance_prob( accept_prob: float try: accept_prob = float( - stats.multivariate_normal.cdf( # type: ignore[arg-type] + stats.multivariate_normal.cdf( upper, lower_limit=lower, mean=np.zeros(len(weights)), diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 3f737b878..fc2a1965d 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -26,7 +26,7 @@ """ import warnings -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union import numpy as np import pandas as pd @@ -43,6 +43,9 @@ from diff_diff.two_stage import _compute_gmm_corrected_meat, _LSMRUnconvergedError from diff_diff.utils import _iterative_fe_solve, safe_inference +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + # Type alias mirroring diff_diff.conley.ConleyMetric so callers can supply # any of the built-in identifiers or a user callable returning a pairwise # distance matrix. @@ -221,7 +224,7 @@ def _compute_nearest_treated_distance_static( d_i = _compute_nearest_treated_distance_sparse( all_coords=all_coords, treated_coords=treated_coords, - metric=metric, # type: ignore[arg-type] + metric=metric, cutoff_km=float(cutoff_km), ) else: @@ -454,7 +457,7 @@ def _compute_nearest_treated_distance_staggered( dists_to_cohort = _compute_nearest_treated_distance_sparse( all_coords=all_coords, treated_coords=treated_coords, - metric=metric, # type: ignore[arg-type] + metric=metric, cutoff_km=float(cutoff_km), ) else: @@ -2090,9 +2093,9 @@ def _validate_spillover_inputs( "is all 0/NaN). SpilloverDiD requires at least one treated unit." ) else: - ft_finite = np.isfinite(data[first_treat].astype(float).values) # type: ignore[arg-type] + ft_finite = np.isfinite(data[first_treat].astype(float).values) n_treated_units = int( - pd.Series(ft_finite & (data[first_treat].astype(float).values != 0)).any() # type: ignore[index] + pd.Series(ft_finite & (data[first_treat].astype(float).values != 0)).any() ) if not n_treated_units: raise ValueError( @@ -2151,7 +2154,7 @@ def fit( treatment: Optional[str] = None, first_treat: Optional[str] = None, covariates: Optional[List[str]] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> SpilloverDiDResults: """Fit the two-stage Gardner DiD with ring-indicator covariates. @@ -2968,7 +2971,11 @@ def fit( # `df_survey`, and meat — even though the same labels passed # via `survey_design.psu=` would be rejected by the panel- # survey validator at `survey.py:1015`. - if self.cluster is not None and resolved_survey.psu is None: + if ( + self.cluster is not None + and resolved_survey is not None + and resolved_survey.psu is None + ): cluster_arr = np.asarray(effective_cluster_ids) unit_arr_full = np.asarray(data[unit].values) # Wave E.3: cluster_arr and the validation unit array @@ -3043,9 +3050,11 @@ def fit( # n_control below). On no-survey path this is bit-identical # to pre-E.3 since survey_weights_fit is None. if survey_weights_fit is not None: + # survey_weights_fit derives from survey_weights (same resolution). + assert survey_weights is not None # Project survey_finite_mask back into the fit-sample (finite_mask) frame survey_finite_in_fit = ( - survey_finite_mask[finite_mask] if n_nan > 0 else (survey_weights > 0) + survey_finite_mask[finite_mask] if (n_nan or 0) > 0 else (survey_weights > 0) ) X_2_fit_active = X_2_fit[survey_finite_in_fit] event_study_meta["n_obs_per_col"] = ( @@ -3081,7 +3090,7 @@ def fit( if survey_weights_fit is not None: solve_kwargs["weights"] = survey_weights_fit solve_kwargs["weight_type"] = "pweight" - coef, residuals, _ = solve_ols(X_2_fit, y_tilde_fit, **solve_kwargs) # type: ignore[misc] + coef, residuals, _ = solve_ols(X_2_fit, y_tilde_fit, **solve_kwargs) # Wave D: Gardner GMM first-stage uncertainty correction. # @@ -3158,6 +3167,7 @@ def fit( if survey_weights is not None and n_nan_or_zero > n_nan: # Project survey_finite_mask into the fit-sample (finite_mask) # frame: True for fit-sample rows that ALSO have weight > 0. + assert survey_weights_fit is not None survey_finite_in_fit = survey_finite_mask[finite_mask] X_2_kept_gamma = X_2_kept[survey_finite_in_fit] eps_2_fit_gamma = eps_2_fit[survey_finite_in_fit] diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 1cd89d595..8f5e44204 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -7,7 +7,7 @@ import bisect import warnings -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -36,6 +36,9 @@ ) from diff_diff.utils import safe_inference, safe_inference_batch +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + # Re-export for backward compatibility __all__ = [ "CallawaySantAnna", @@ -1794,7 +1797,7 @@ def fit( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> CallawaySantAnnaResults: """ Fit the Callaway-Sant'Anna estimator. @@ -2108,11 +2111,15 @@ def fit( compute_survey_metadata, ) + # This branch is only reached with a survey design present. + assert survey_design is not None effective_survey_design = replace(survey_design, psu=self.cluster) resolved_survey = _inject_cluster_as_psu(resolved_survey, cluster_ids) # Recompute survey_metadata to reflect the new effective PSU # (so n_psu / df_survey on Results reflect the injected # cluster, not the no-PSU pre-injection state). + # resolved_survey non-None implies survey_design was passed. + assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) if survey_design.weights @@ -2987,6 +2994,8 @@ def _outcome_regression( sw_t_sum = float(np.sum(sw_treated)) sw_t_norm = sw_treated / sw_t_sum att = float(np.sum(sw_t_norm * treated_residuals)) + # sw_treated/sw_control are supplied together by all callers. + assert sw_control is not None W_c = sw_control xbar_t = np.sum(sw_treated[:, None] * X_treated, axis=0) / sw_t_sum inf_treated = sw_t_norm * (treated_residuals - att) @@ -3891,6 +3900,8 @@ def _compute_att_gt_rc( has_covariates = covariates is not None and obs_covariates is not None if has_covariates: + # The flag definition above guarantees this (mypy can't track it). + assert obs_covariates is not None X_gt = obs_covariates[treated_t] X_gs = obs_covariates[treated_s] X_ct = obs_covariates[control_t] diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 96338abc0..7062106bf 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -5,7 +5,7 @@ group-time average treatment effects into summary measures. """ -from typing import Any, Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, overload import numpy as np import pandas as pd @@ -70,7 +70,7 @@ def _aggregate_simple( df: pd.DataFrame, unit: str, precomputed: Optional["PrecomputedData"] = None, - ) -> Tuple[float, float]: + ) -> Tuple[float, float, Optional[int]]: """ Compute simple weighted average of ATT(g,t). @@ -86,10 +86,10 @@ def _aggregate_simple( in the overall ATT. Pre-treatment effects are computed for parallel trends assessment but are not aggregated into the overall ATT. """ - effects = [] + effects_list: List[Any] = [] weights_list = [] gt_pairs = [] - groups_for_gt = [] + groups_list: List[Any] = [] # Fixed per-cohort aggregation weights (R's did::aggte pg = n_g / N), # preferring the unit-level RC mass so allow_unbalanced_panel weights the @@ -101,7 +101,7 @@ def _aggregate_simple( # Pre-treatment effects are for parallel trends, not overall ATT if t < g - self.anticipation: continue - effects.append(data["effect"]) + effects_list.append(data["effect"]) # Use fixed cohort-level survey weight sum for aggregation. # For RCS, data["agg_weight"] holds the fixed cohort mass; # for panel, fallback to data["n_treated"]. @@ -110,10 +110,10 @@ def _aggregate_simple( else: weights_list.append(data.get("agg_weight", data["n_treated"])) gt_pairs.append((g, t)) - groups_for_gt.append(g) + groups_list.append(g) # Guard against empty post-treatment set - if len(effects) == 0: + if len(effects_list) == 0: import warnings warnings.warn( @@ -124,9 +124,9 @@ def _aggregate_simple( ) return np.nan, np.nan, None - effects = np.array(effects) + effects = np.array(effects_list) weights = np.array(weights_list, dtype=float) - groups_for_gt = np.array(groups_for_gt) + groups_for_gt = np.array(groups_list) # Exclude NaN effects from aggregation (R's aggte() convention). # No warning here — fit() emits a consolidated skip warning covering @@ -486,6 +486,8 @@ def _compute_combined_influence_function( # With survey weights: pg[g] = sum(sw_g) / sum(sw_all) group_sizes = {} if survey_w is not None: + # Survey weights come from precomputed, so it is present here. + assert precomputed is not None # Survey-weighted group sizes precomputed_cohorts = precomputed["unit_cohorts"] for g in unique_groups: @@ -493,6 +495,8 @@ def _compute_combined_influence_function( group_sizes[g] = float(np.sum(survey_w[mask_g])) total_weight = float(np.sum(survey_w)) elif _is_rcs: + # The RCS path always builds precomputed (obs-level bookkeeping). + assert precomputed is not None # RCS without survey: count observations per cohort precomputed_cohorts = precomputed["unit_cohorts"] for g in unique_groups: @@ -541,6 +545,8 @@ def _compute_combined_influence_function( unit_groups_array = np.full(n_units, -1, dtype=np.float64) if _is_rcs: + # The RCS path always builds precomputed (obs-level bookkeeping). + assert precomputed is not None # RCS: direct vectorized assignment — obs indices are positions precomputed_cohorts = precomputed["unit_cohorts"] for g in unique_groups: @@ -638,6 +644,35 @@ def _compute_combined_influence_function( return psi_total, all_units + @overload + def _compute_aggregated_se_with_wif( + self, + gt_pairs: List[Tuple[Any, Any]], + weights: np.ndarray, + effects: np.ndarray, + groups_for_gt: np.ndarray, + influence_func_info: Dict, + df: pd.DataFrame, + unit: str, + precomputed: Optional["PrecomputedData"] = None, + return_psi: Literal[False] = False, + ) -> Tuple[float, Optional[int]]: ... + + @overload + def _compute_aggregated_se_with_wif( + self, + gt_pairs: List[Tuple[Any, Any]], + weights: np.ndarray, + effects: np.ndarray, + groups_for_gt: np.ndarray, + influence_func_info: Dict, + df: pd.DataFrame, + unit: str, + precomputed: Optional["PrecomputedData"] = None, + *, + return_psi: Literal[True], + ) -> Tuple[float, np.ndarray, Optional[int]]: ... + def _compute_aggregated_se_with_wif( self, gt_pairs: List[Tuple[Any, Any]], @@ -909,6 +944,8 @@ def _aggregate_event_study( # reference cells contribute nothing to the variance but their cohort # weight dilutes the real cells, matching R's dynamic aggregation. groups_for_gt = np.array([g for (g, t) in gt_pairs]) + # The wif-SE path requires the fit-time frame (callers pass it). + assert df is not None and unit is not None agg_se, psi_e, eff_df = self._compute_aggregated_se_with_wif( gt_pairs, weights, @@ -963,7 +1000,7 @@ def _aggregate_event_study( df=df_survey_val, ) - event_study_effects = {} + event_study_effects: Dict[Any, Dict[str, Any]] = {} for idx, e in enumerate(agg_periods): event_study_effects[e] = { "effect": agg_effects_list[idx], diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 1b9acf8cd..326e950d7 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -360,6 +360,8 @@ def _agg_mass(gt): _bootstrap_cluster_variance_unidentified = False if _use_survey_bootstrap: + # The flag definition above guarantees this (mypy can't track it). + assert resolved_survey_unit is not None # PSU-level multiplier weights, generated AND expanded one draw-block # at a time so the (n_bootstrap, n_units) matrix is never built in # full. This is the dominant allocation at large n_units, including @@ -499,7 +501,10 @@ def _make_weight_iter( overall_col = len(columns) columns.append([(None, overall_combined_if)]) es_col0 = len(columns) + # rel_periods is non-empty only when event-study info was built. + assert event_study_info is not None or not rel_periods for e in rel_periods: + assert event_study_info is not None columns.append([(None, event_study_info[e]["combined_if"])]) perturbations = tiled_if_matmul(weight_stream, self.n_bootstrap, n_units, columns) @@ -684,10 +689,13 @@ def _make_weight_iter( gt_cis = {gt: (np.nan, np.nan) for gt in gt_cis} if gt_cis else gt_cis gt_p_values = {gt: np.nan for gt in gt_p_values} if gt_p_values else gt_p_values if event_study_ses: + # ses/cis/p_values are populated together upstream. + assert event_study_cis is not None and event_study_p_values is not None event_study_ses = {k: np.nan for k in event_study_ses} event_study_cis = {k: (np.nan, np.nan) for k in event_study_cis} event_study_p_values = {k: np.nan for k in event_study_p_values} if group_effect_ses: + assert group_effect_cis is not None and group_effect_p_values is not None group_effect_ses = {k: np.nan for k in group_effect_ses} group_effect_cis = {k: (np.nan, np.nan) for k in group_effect_cis} group_effect_p_values = {k: np.nan for k in group_effect_p_values} diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index f3492745b..a838ce693 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -10,7 +10,7 @@ """ import warnings -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -29,6 +29,9 @@ from diff_diff.staggered_triple_diff_results import StaggeredTripleDiffResults from diff_diff.utils import safe_inference +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + __all__ = [ "StaggeredTripleDifference", "StaggeredTripleDiffResults", @@ -203,7 +206,7 @@ def fit( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> StaggeredTripleDiffResults: """ Fit the staggered triple difference estimator. @@ -761,6 +764,11 @@ def fit( if event_study_effects and bootstrap_results.event_study_ses: for e_key in event_study_effects: if e_key in bootstrap_results.event_study_ses: + # ses/cis/p_values are populated together. + assert ( + bootstrap_results.event_study_cis is not None + and bootstrap_results.event_study_p_values is not None + ) event_study_effects[e_key]["se"] = bootstrap_results.event_study_ses[ e_key ] @@ -1320,6 +1328,8 @@ def _run_pairwise_did( # Build covariate matrix with intercept for the pair covX = None if has_covariates: + # The flag definition above guarantees this (mypy can't track it). + assert covariate_matrix is not None X_pair = covariate_matrix[pair_idx] covX = np.column_stack([np.ones(n_pair), X_pair]) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index c14286739..a366f3966 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -11,12 +11,15 @@ import warnings from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, overload import numpy as np import pandas as pd from diff_diff.bootstrap_utils import compute_effect_bootstrap_stats + +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign from diff_diff.linalg import LinearRegression from diff_diff.results import _format_survey_block, _get_significance_stars from diff_diff.utils import ( @@ -620,7 +623,7 @@ def fit( time: str, first_treat: str, covariates: Optional[List[str]] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> SunAbrahamResults: """ Fit the Sun-Abraham estimator using saturated regression. @@ -890,6 +893,8 @@ def fit( resolved_survey = _inject_cluster_as_psu(resolved_survey, effective_cluster_ids) if resolved_survey.psu is not None and survey_metadata is not None: + # resolved_survey non-None implies survey_design was passed. + assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) if survey_design.weights @@ -1180,6 +1185,8 @@ def _refit_sa(w_r): ) # Override df if replicates dropped + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None if _n_valid_rep_sa < resolved_survey.n_replicates: _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 if survey_metadata is not None: @@ -1339,7 +1346,7 @@ def _fit_saturated_regression( cluster_var: Optional[str], survey_weights: Optional[np.ndarray] = None, survey_weight_type: str = "pweight", - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, vcov_type: str = "hc1", ) -> Tuple[ Dict[Tuple[Any, int], float], @@ -1698,6 +1705,35 @@ def _compute_iw_effects( return event_study_effects, cohort_weights + @overload + def _compute_overall_att( + self, + df: pd.DataFrame, + first_treat: str, + event_study_effects: Dict[int, Dict[str, Any]], + cohort_effects: Dict[Tuple[Any, int], float], + cohort_weights: Dict[int, Dict[Any, float]], + vcov_cohort: np.ndarray, + coef_index_map: Dict[Tuple[Any, int], int], + survey_weight_col: Optional[str] = None, + return_overall_weights: Literal[False] = False, + ) -> Tuple[float, float]: ... + + @overload + def _compute_overall_att( + self, + df: pd.DataFrame, + first_treat: str, + event_study_effects: Dict[int, Dict[str, Any]], + cohort_effects: Dict[Tuple[Any, int], float], + cohort_weights: Dict[int, Dict[Any, float]], + vcov_cohort: np.ndarray, + coef_index_map: Dict[Tuple[Any, int], int], + survey_weight_col: Optional[str] = None, + *, + return_overall_weights: Literal[True], + ) -> Tuple[float, float, Optional[Dict[Tuple[Any, int], float]]]: ... + def _compute_overall_att( self, df: pd.DataFrame, @@ -1709,7 +1745,10 @@ def _compute_overall_att( coef_index_map: Dict[Tuple[Any, int], int], survey_weight_col: Optional[str] = None, return_overall_weights: bool = False, - ) -> Tuple[float, float]: + ) -> Union[ + Tuple[float, float], + Tuple[float, float, Optional[Dict[Tuple[Any, int], float]]], + ]: """ Compute overall ATT as weighted average of post-treatment effects. @@ -1808,7 +1847,7 @@ def _run_bootstrap( cluster_var: Optional[str], original_event_study: Dict[int, Dict[str, Any]], original_overall_att: float, - resolved_survey: object = None, + resolved_survey: Optional["ResolvedSurveyDesign"] = None, survey_weights: Optional[np.ndarray] = None, survey_weight_type: str = "pweight", survey_weight_col: Optional[str] = None, @@ -2010,7 +2049,7 @@ def _run_rao_wu_bootstrap( cluster_var: Optional[str], original_event_study: Dict[int, Dict[str, Any]], original_overall_att: float, - resolved_survey: object, + resolved_survey: "ResolvedSurveyDesign", survey_weight_type: str, survey_weight_col: Optional[str], rng: np.random.Generator, diff --git a/diff_diff/survey.py b/diff_diff/survey.py index 4443636ed..04f3ef33c 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -1232,7 +1232,16 @@ def _extract_unit_survey_weights(data, unit_col, survey_design, unit_order): return np.array([unit_w[u] for u in unit_order], dtype=np.float64) -def _resolve_survey_for_fit(survey_design, data, inference_mode="analytical"): +def _resolve_survey_for_fit( + survey_design: Optional["SurveyDesign"], + data: pd.DataFrame, + inference_mode: str = "analytical", +) -> Tuple[ + Optional["ResolvedSurveyDesign"], + Optional[np.ndarray], + str, + Optional["SurveyMetadata"], +]: """ Shared helper: validate and resolve a SurveyDesign for an estimator fit() call. @@ -1759,6 +1768,8 @@ def _finalize(meat_scalar: float) -> float: return meat_scalar if scaffolding.mode == "no_strata_no_psu": + # Mode invariant: the scaffolding builder fills this field for this mode. + assert scaffolding.adjustment_direct is not None if scaffolding.n < 2: return float("nan") psi_mean = psi.mean() @@ -1767,6 +1778,8 @@ def _finalize(meat_scalar: float) -> float: return _finalize(meat) if scaffolding.mode == "psu_only": + # Mode invariant: the scaffolding builder fills these fields for this mode. + assert scaffolding.n_psu_only is not None and scaffolding.adjustment_only is not None if scaffolding.n_psu_only < 2: if scaffolding.legitimate_zero_count > 0: return 0.0 @@ -1783,6 +1796,8 @@ def _finalize(meat_scalar: float) -> float: S = len(scaffolding.n_psu_per_stratum) P = len(scaffolding.psu_stratum) + # Mode invariant: the stratified scaffolding builder fills these fields. + assert scaffolding.n_psu_per_stratum is not None and scaffolding.adjustment_h is not None psu_sums = np.bincount(scaffolding.psu_codes, weights=psi, minlength=P) sum_by_h = np.bincount(scaffolding.psu_stratum, weights=psu_sums, minlength=S) sum2_by_h = np.bincount(scaffolding.psu_stratum, weights=psu_sums * psu_sums, minlength=S) @@ -2237,6 +2252,8 @@ def compute_replicate_vcov( from diff_diff.linalg import solve_ols rep_weights = resolved.replicate_weights + # Replicate-variance entry points are only reached on replicate designs. + assert rep_weights is not None method = resolved.replicate_method R = resolved.n_replicates k = X.shape[1] @@ -2380,6 +2397,8 @@ def compute_replicate_if_variance( """ psi = np.asarray(psi, dtype=np.float64).ravel() rep_weights = resolved.replicate_weights + # Replicate-variance entry points are only reached on replicate designs. + assert rep_weights is not None method = resolved.replicate_method R = resolved.n_replicates @@ -2518,6 +2537,8 @@ def compute_replicate_refit_variance( full_sample_estimate = np.asarray(full_sample_estimate, dtype=np.float64).ravel() k = len(full_sample_estimate) rep_weights = resolved.replicate_weights + # Replicate-variance entry points are only reached on replicate designs. + assert rep_weights is not None method = resolved.replicate_method R = resolved.n_replicates diff --git a/diff_diff/synthetic_control_results.py b/diff_diff/synthetic_control_results.py index f728b5f2d..74e2a3057 100644 --- a/diff_diff/synthetic_control_results.py +++ b/diff_diff/synthetic_control_results.py @@ -2963,6 +2963,8 @@ def conformal_average_effect( assert snap_cov is not None cal_cov = list(snap_cov.pre_periods) + list(snap_cov.post_periods) x1_rows, X0_rows = self._conformal_covariate_rows(covariates, cal_cov) + # covariates is non-None on this branch, so rows were built. + assert x1_rows is not None and X0_rows is not None # Covariate rows collapse with the SAME T*-block structure so the # collapsed panel remains a coherent Z (each block-averaged # covariate row enters the proxy like a block-averaged outcome). diff --git a/diff_diff/synthetic_did.py b/diff_diff/synthetic_did.py index c56f87f2d..5772a3aa2 100644 --- a/diff_diff/synthetic_did.py +++ b/diff_diff/synthetic_did.py @@ -924,8 +924,8 @@ def fit( # type: ignore[override] _strata_control_eff: np.ndarray = np.zeros(len(control_units), dtype=np.int64) _strata_treated_eff: np.ndarray = np.zeros(len(treated_units), dtype=np.int64) else: - _strata_control_eff = _strata_control # type: ignore[assignment] - _strata_treated_eff = _strata_treated # type: ignore[assignment] + _strata_control_eff = _strata_control + _strata_treated_eff = _strata_treated # Fit-time feasibility guard for stratified-permutation placebo # (per `feedback_front_door_over_retry_swallow.md`). Case B / Case C @@ -1280,6 +1280,8 @@ def fit( # type: ignore[override] treated_post_trajectory=treated_post_trajectory, time_weights_array=time_weights, ) + # results_ was just constructed above. + assert self.results_ is not None # Explicit LOO granularity flag for ``get_loo_effects_df``. The # non-survey and pweight-only jackknife paths run unit-level LOO # (one estimate per unit, matching ``control_unit_ids + @@ -1507,11 +1509,14 @@ def _bootstrap_se( # rw is the constant per-control survey weight. if _use_rao_wu: boot_rw = generate_rao_wu_weights(resolved_survey, rng) + assert boot_rw is not None rw_control_full = boot_rw[:n_control] rw_treated_full = boot_rw[n_control:] rw_control_draw = rw_control_full[boot_idx[boot_is_control]] rw_treated_draw = rw_treated_full[boot_idx[~boot_is_control] - n_control] elif _pweight_only: + # pweight-only flag implies control weights exist. + assert w_control is not None rw_control_draw = w_control[boot_idx[boot_is_control]] rw_treated_draw = ( w_treated[boot_idx[~boot_is_control] - n_control] @@ -2079,6 +2084,8 @@ def _placebo_variance_se_survey( pseudo_control_idx = np.where(pseudo_control_mask)[0] # Pseudo-panel + # This survey pseudo-panel branch requires control weights. + assert w_control is not None Y_pre_pseudo_control = Y_pre_control[:, pseudo_control_idx] Y_post_pseudo_control = Y_post_control[:, pseudo_control_idx] pseudo_w_tr = w_control[pseudo_treated_idx] diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index d070e9204..571bacecb 100644 --- a/diff_diff/triple_diff.py +++ b/diff_diff/triple_diff.py @@ -29,7 +29,7 @@ import warnings from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -38,6 +38,10 @@ from diff_diff.results import _format_survey_block, _get_significance_stars from diff_diff.utils import safe_inference +if TYPE_CHECKING: + from diff_diff.survey import SurveyDesign + + _MIN_CELL_SIZE = 10 # ============================================================================= @@ -2134,7 +2138,7 @@ def triple_difference( rank_deficient_action: str = "warn", epv_threshold: float = 10, pscore_fallback: str = "error", - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> TripleDifferenceResults: """ Estimate Triple Difference (DDD) treatment effect. diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 33c628f29..fc83dc6f3 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: from diff_diff.bacon import BaconDecompositionResults + from diff_diff.survey import SurveyDesign from diff_diff.estimators import DifferenceInDifferences from diff_diff.linalg import LinearRegression @@ -131,7 +132,7 @@ def fit( # type: ignore[override] time: str, unit: str, covariates: Optional[List[str]] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> DiDResults: """ Fit Two-Way Fixed Effects model. @@ -450,6 +451,8 @@ def fit( # type: ignore[override] resolved_survey = _inject_cluster_as_psu(resolved_survey, survey_cluster_ids) if resolved_survey.psu is not None and survey_metadata is not None: + # resolved_survey non-None implies survey_design was passed. + assert survey_design is not None raw_w = ( data[survey_design.weights].values.astype(np.float64) if survey_design.weights @@ -484,6 +487,8 @@ def fit( # type: ignore[override] # on demeaned scores but the kernel grid uses the original space # (coords) and time/unit indexing. if _fit_vcov_type == "conley": + # Validated by the conley front-door (_validate_conley_estimator_inputs). + assert self.conley_coords is not None _conley_coords_arr: Optional[np.ndarray] = np.column_stack( [ data[self.conley_coords[0]].values.astype(np.float64), @@ -665,6 +670,8 @@ def _refit_twfe(w_r): if survey_metadata and survey_metadata.df_survey else 0 # rank-deficient replicate → NaN inference ) + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None if _n_valid_rep_twfe < resolved_survey.n_replicates: _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 if survey_metadata is not None: diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 0fb92ae83..6f317d9bf 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -48,7 +48,7 @@ # Forward reference for the Wave E.1 survey-design path. Imported under # TYPE_CHECKING to keep the runtime cost zero and avoid any future # circular-import surprises with diff_diff.survey. - from diff_diff.survey import ResolvedSurveyDesign + from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign # Maximum number of elements before falling back to per-column sparse aggregation. # 10M float64 elements ≈ 80 MB peak allocation. Above this, per-column .getcol() @@ -331,7 +331,7 @@ def _compute_gmm_corrected_meat( _validate_conley_kwargs( conley_coords, conley_cutoff_km, - conley_metric, # type: ignore[arg-type] # validator raises ValueError if None + conley_metric, # validator raises ValueError if None conley_kernel, n_for_conley, time=conley_time, @@ -1435,7 +1435,7 @@ def fit( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, ) -> TwoStageDiDResults: """ Fit the two-stage DiD estimator. @@ -1699,6 +1699,8 @@ def fit( if resolved_survey.psu is not None and survey_metadata is not None: from diff_diff.survey import compute_survey_metadata + # resolved_survey non-None implies survey_design was passed. + assert survey_design is not None raw_w = ( np.asarray(data[survey_design.weights].values, dtype=np.float64) if survey_design.weights @@ -1879,21 +1881,21 @@ def fit( from diff_diff.survey import compute_replicate_refit_variance # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + _es_effects_ts = event_study_effects or {} + _grp_effects_ts = group_effects or {} _sorted_es_periods_ts = sorted( - e - for e in (event_study_effects or {}).keys() - if np.isfinite(event_study_effects[e]["effect"]) + e for e in _es_effects_ts.keys() if np.isfinite(_es_effects_ts[e]["effect"]) ) _sorted_groups_ts = sorted( - g for g in (group_effects or {}).keys() if np.isfinite(group_effects[g]["effect"]) + g for g in _grp_effects_ts.keys() if np.isfinite(_grp_effects_ts[g]["effect"]) ) _n_es_ts = len(_sorted_es_periods_ts) _n_grp_ts = len(_sorted_groups_ts) # Build full-sample estimate from actual outputs _full_est_ts = [overall_att] - _full_est_ts.extend([event_study_effects[e]["effect"] for e in _sorted_es_periods_ts]) - _full_est_ts.extend([group_effects[g]["effect"] for g in _sorted_groups_ts]) + _full_est_ts.extend([_es_effects_ts[e]["effect"] for e in _sorted_es_periods_ts]) + _full_est_ts.extend([_grp_effects_ts[g]["effect"] for g in _sorted_groups_ts]) def _refit_ts(w_r): # Wave E.3 parity (PR #482 SpilloverDiD precedent): the main fit @@ -2011,6 +2013,8 @@ def _refit_ts(w_r): overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) # Override df if replicates were dropped + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None if _n_valid_rep_ts < resolved_survey.n_replicates: _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 if survey_metadata is not None: @@ -3440,7 +3444,7 @@ def two_stage_did( covariates: Optional[List[str]] = None, aggregate: Optional[str] = None, balance_e: Optional[int] = None, - survey_design: object = None, + survey_design: Optional["SurveyDesign"] = None, vcov_type: str = "hc1", **kwargs, ) -> TwoStageDiDResults: diff --git a/diff_diff/two_stage_bootstrap.py b/diff_diff/two_stage_bootstrap.py index 7039cc667..b63d9ac43 100644 --- a/diff_diff/two_stage_bootstrap.py +++ b/diff_diff/two_stage_bootstrap.py @@ -43,6 +43,7 @@ class TwoStageDiDBootstrapMixin: alpha: float seed: Optional[int] horizon_max: Optional[int] + pretrends: bool if TYPE_CHECKING: from scipy import sparse diff --git a/diff_diff/utils.py b/diff_diff/utils.py index dbe29ea94..a250873fb 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -2896,7 +2896,7 @@ def _demean_map_rust( for start, stop in zip(bounds, bounds[1:]): x_mat = np.ascontiguousarray(np.column_stack(x_cols[start:stop]), dtype=np.float64) try: - out, iters = _rust_demean_map( # type: ignore[misc] + out, iters = _rust_demean_map( x_mat, codes_mat, n_groups, diff --git a/diff_diff/wooldridge.py b/diff_diff/wooldridge.py index 1869e80ee..d0399f76b 100644 --- a/diff_diff/wooldridge.py +++ b/diff_diff/wooldridge.py @@ -529,7 +529,7 @@ def _validate_constructor_args( def results_(self) -> WooldridgeDiDResults: if not self.is_fitted_: raise RuntimeError("Call fit() before accessing results_") - return self._results # type: ignore[return-value] + return self._results def get_params(self) -> Dict[str, Any]: """Return estimator parameters (sklearn-compatible).""" @@ -1252,7 +1252,7 @@ def _fit_ols( else: _cl_coords = _cl_time = _cl_unit = None - coefs, resids, vcov = solve_ols( + coefs, resids, vcov = solve_ols( # type: ignore[call-overload] # mypy union-explosion limitation at this many Optional args X, y, cluster_ids=cluster_ids, diff --git a/diff_diff/wooldridge_results.py b/diff_diff/wooldridge_results.py index 68fb12e00..f380592bb 100644 --- a/diff_diff/wooldridge_results.py +++ b/diff_diff/wooldridge_results.py @@ -629,7 +629,7 @@ def summary(self, aggregation: str = "simple") -> str: lines.append("-" * 70) def _fmt_row(label: str, att: float, se: float, t: float, p: float, ci: Tuple) -> str: - from diff_diff.results import _get_significance_stars # type: ignore + from diff_diff.results import _get_significance_stars stars = _get_significance_stars(p) if not np.isnan(p) else "" ci_lo = f"{ci[0]:.4f}" if not np.isnan(ci[0]) else "NaN" @@ -816,7 +816,7 @@ def plot_event_study(self, weights: str = "cell", **kwargs) -> None: # cohort_share→cell (or any cross-scheme) stale-cache bugs. self.aggregate("event", weights=weights) - from diff_diff.visualization import plot_event_study # type: ignore + from diff_diff.visualization import plot_event_study effects = {k: v["att"] for k, v in (self.event_study_effects or {}).items()} if weights == "cohort_share": diff --git a/pyproject.toml b/pyproject.toml index b1c63f557..be534c754 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,3 +160,23 @@ strict_optional = true implicit_reexport = true # Disable errors for numpy type compatibility issues disable_error_code = ["arg-type", "return-value", "var-annotated", "assignment"] +# Stale-ignore hygiene: versions are pinned (dev extra + lint.yml), so +# unused-ignore warnings can't churn with tool upgrades. +warn_unused_ignores = true + +# Third-party plotting stubs: matplotlib ships py.typed but is not installed +# in the typecheck CI job (and plotly has partial stubs). Skip both so local +# and CI runs see identical Any surfaces regardless of what is installed. +[[tool.mypy.overrides]] +module = ["matplotlib.*", "plotly.*"] +follow_imports = "skip" + +# prep_dgp builds seeded synthetic panels with Optional covariate arrays that +# are constructed and consumed under the same `add_covariates` flag. The +# None-vs-array shape is load-bearing for the seeded RNG stream (drawing +# unconditionally would shift every downstream draw), so per-site narrowing +# is not worth the regression risk in a test-data generator. Tightening is +# tracked in TODO.md (Testing / docs). +[[tool.mypy.overrides]] +module = ["diff_diff.prep_dgp"] +disable_error_code = ["index"] diff --git a/tests/test_openai_review.py b/tests/test_openai_review.py index 682a0f4e9..9fd197080 100644 --- a/tests/test_openai_review.py +++ b/tests/test_openai_review.py @@ -4677,26 +4677,43 @@ class TestLintWorkflowPinSync: LINT_WORKFLOW = pathlib.Path(__file__).resolve().parent.parent / ".github/workflows/lint.yml" PYPROJECT = pathlib.Path(__file__).resolve().parent.parent / "pyproject.toml" + TOOLS = ("ruff", "black", "mypy") + def _workflow_pins(self) -> "dict[str, str]": text = self.LINT_WORKFLOW.read_text() install_lines = [line for line in text.splitlines() if "pip install" in line] assert install_lines, "lint.yml must contain a pip install step for the linters" - pins = dict(re.findall(r"\b(ruff|black)==([0-9][\w.]*)", "\n".join(install_lines))) - assert set(pins) == { - "ruff", - "black", - }, f"lint.yml must pip-install pinned ruff and black; found pins for {sorted(pins)}" + pins = dict(re.findall(r"\b(ruff|black|mypy)==([0-9][\w.]*)", "\n".join(install_lines))) + assert set(pins) == set( + self.TOOLS + ), f"lint.yml must pip-install pinned {self.TOOLS}; found pins for {sorted(pins)}" return pins def _pyproject_pins(self) -> "dict[str, str]": text = self.PYPROJECT.read_text() - pins = dict(re.findall(r'"(ruff|black)==([0-9][\w.]*)"', text)) - assert set(pins) == { - "ruff", - "black", - }, f"pyproject dev extra must pin ruff and black exactly; found {sorted(pins)}" + pins = dict(re.findall(r'"(ruff|black|mypy)==([0-9][\w.]*)"', text)) + assert set(pins) == set( + self.TOOLS + ), f"pyproject dev extra must pin {self.TOOLS} exactly; found {sorted(pins)}" return pins + def test_mypy_job_stub_deps_exact_pinned(self): + """The Mypy job installs numpy/pandas/scipy so mypy sees real stubs + (an absent numpy silently becomes Any under ignore_missing_imports and + CI would go green while local runs stay red). Those stub pins are + CI-only (pyproject keeps runtime floors), so exactness is asserted + here rather than synced to pyproject.""" + text = self.LINT_WORKFLOW.read_text() + install_lines = [line for line in text.splitlines() if "pip install" in line] + stub_pins = dict( + re.findall(r"\b(numpy|pandas|scipy)==([0-9][\w.]*)", "\n".join(install_lines)) + ) + assert set(stub_pins) == {"numpy", "pandas", "scipy"}, ( + f"lint.yml's Mypy job must exact-pin numpy, pandas, and scipy for " + f"stub stability; found pins for {sorted(stub_pins)}. See the " + f"comment above the install step in lint.yml." + ) + def test_lint_workflow_pins_match_pyproject(self): wf = self._workflow_pins() py = self._pyproject_pins()