refactor: overload the flag-driven return unions

[michael-denyer]

Continues the pyrefly burn-down. 135 → 113 entries.

All 22 entries come from one pattern: a function whose return type is a union selected
by a boolean argument. Every caller inherits both arms and has to answer for the one that
cannot occur — so a single un-overloaded callee scatters errors across the files that use it.

File	Errors	Cause
kinship/compute.py	6	_stream_s_full_and_chr (5), compute_loco_kinship_streaming (1)
lmm/likelihood_numpy.py	5	three un-overloaded helpers under the optimizer
lmm/compute_numpy.py	1	unpacking golden_section_optimize_lambda_split_ncvt1_numpy
pipeline.py	1	the last entry in the file, from compute_loco_kinship_streaming
lmm/loco.py	1	same callee
tests	8	downstream of the same four callees

likelihood_numpy.py — finishing a started migration

golden_section_optimize_lambda_numpy was already overloaded on return_pab. The
helpers beneath it never were, so the union just moved one level down. Overloaded to match:

• _batch_reml_at_lambda_numpy — ndarray vs (ndarray, ndarray)
• _batch_reml_at_lambda_split_ncvt1_numpy — same
• _batch_golden_section_numpy — 2- vs 3-tuple, keyed on whether a Pab evaluator is passed.
  Its two callable parameters were bare and untyped; they now have named aliases
  (_BatchLoglFn, _BatchLoglPabFn).
• golden_section_optimize_lambda_split_ncvt1_numpy — 2- vs 3-tuple

kinship/compute.py

_stream_s_full_and_chr returns S_full only when asked, so its ndarray | None first
element forced five unguarded call sites inside compute_loco_kinship_streaming.

compute_loco_kinship_streaming itself returns either the iterator or an
(iterator, cache) pair. Overloading it lets lmm/loco.py drop this:

    cast(Iterator[tuple[str, np.ndarray]], loco_stream),
    # "The streamer's return type is overloaded on return_snp_stats;
    #  with it set we always get the (iterator, cache) pair."

The comment asserted an overload that did not exist. Now it does, so the cast is gone and
the claim is enforced by the type checker instead of by a comment.

Signature changes

return_pab and S_full_accum are now keyword-only. That is what lets two overloads
express the signature exactly instead of needing three (a Literal[True] arm can't follow a
defaulted parameter positionally). Every existing call site — production and tests — already
passed them by keyword, so no call site changed. Verified by grep across the repo.

The one non-signature change

compute_loco_kinship_streaming now returns on if snp_stats_cache is not None: rather than
if return_snp_stats:. Equivalent by construction, and it narrows where the flag doesn't.

I first tried moving the cache construction into the return branch — that was wrong, for
two reasons worth recording: there is a deliberate if not return_snp_stats: del stats that
frees the stats early, and SnpStats.__post_init__ marks the stats arrays read-only in
place, so relocating the construction would change when that happens. The cache is built
where it always was.

Verification

@overload stubs are erased at runtime, cast is the identity, and the narrowing swap is
equivalent — so no behaviour change is possible here. Checked anyway, on the LOCO streaming
path this actually restructures:

Compared (before vs after)	Result
LOCO .assoc.txt	byte-identical
pve, pve_se, n_snps_tested	identical
Per-chromosome kinship matrix sums, both flag values	identical
Returned SnpStatsCache — col_means/col_vars/miss_counts sums, n_snps, sample_scope, and col_means.flags.writeable	identical
pytest tests/	2241 passed, 10 skipped, 3 xfailed
pytest -m "tier2 or slow"	38 passed
ruff check / ruff format --check	clean

The tier2 or slow run is included because the default addopts filter skips it and it
covers the streaming path this touches.

Baseline

Regenerated on Python 3.11.13 / numpy 2.4.6 per #121, verified on both:

Environment	pyrefly check --baseline
Python 3.11.13, numpy 2.4.6	0 errors
Python 3.13.5, numpy 2.5.1	0 errors

Running total

174 (#121) → 152 (#122) → 135 (#123) → 113. pipeline.py, validation/compare.py,
kinship/compute.py and lmm/likelihood_numpy.py are now at zero.

Next largest: tests/test_incremental_writer.py (23, undiagnosed), tests/test_jlinalg_dgemm.py
(8), lmm/io.py (7), lmm/compute_numpy.py (6 — Optional narrowing across a mode dispatch,
plus TypedDict.update).