Move MatchJob dataclass from grouping.py to config.py

MatchJob is a plain data container that belongs with the other
config/data types, not in the grouping module.

Author: jeromekelleher

tests/test_grouping.py

@@ -29,7 +29,7 @@
 import numpy as np
 import pytest
 
-from tsinfer import grouping
+from tsinfer import config, grouping
 
 
 @dataclasses.dataclass
@@ -558,7 +558,7 @@ class TestAssignGroups:
     """Test that assign_groups takes ungrouped MatchJobs and assigns groups."""
 
     def _make_job(self, haplotype_index, time, start=0, end=100):
-        return grouping.MatchJob(
+        return config.MatchJob(
             haplotype_index=haplotype_index,
             source="test",
             sample_id=f"s{haplotype_index}",
@@ -617,7 +617,7 @@ def test_sorted_by_group_then_haplotype_index(self):
             )
 
     def test_preserves_job_fields(self):
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=0,
             source="my_source",
             sample_id="my_sample",

tests/test_matcher_fixtures.py

@@ -30,7 +30,7 @@
 import pytest
 
 from tsinfer import config, matching, vcz
-from tsinfer.grouping import MatchJob
+from tsinfer.config import MatchJob
 
 # ---------------------------------------------------------------------------
 # Helpers

tests/test_matching.py

@@ -25,7 +25,7 @@
 import numpy as np
 import tskit
 
-from tsinfer import config, grouping, matching, vcz
+from tsinfer import config, matching, vcz
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -110,7 +110,7 @@ def _make_job(
     population_id=None,
 ):
     """Create a MatchJob with sensible defaults for testing."""
-    return grouping.MatchJob(
+    return config.MatchJob(
         haplotype_index=haplotype_index,
         source=source,
         sample_id=sample_id,

tests/test_vcz.py

@@ -30,7 +30,7 @@
 import pytest
 import zarr
 
-from tsinfer import config, grouping, vcz
+from tsinfer import config, vcz
 
 # ---------------------------------------------------------------------------
 # Fixtures
@@ -346,7 +346,7 @@ def test_ancestor_haplotype(self):
             schedule=schedule,
         )
         # a0 has genotype [0, 1, 0]
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=1,
             source="ancestors",
             sample_id="a0",
@@ -371,7 +371,7 @@ def test_ancestor_haplotype_second(self):
             schedule=schedule,
         )
         # a1 has genotype [1, 0, 1]
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=2,
             source="ancestors",
             sample_id="a1",
@@ -396,7 +396,7 @@ def test_sample_haplotype_encoding(self):
             schedule=schedule,
         )
         # sample_0: gt [0, 1, 1] → encoded [0=anc, 1=derived, 1=derived]
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=3,
             source="test",
             sample_id="sample_0",
@@ -421,7 +421,7 @@ def test_sample_haplotype_second(self):
             schedule=schedule,
         )
         # sample_1: gt [1, 0, 0] → encoded [1=derived, 0=anc, 0=anc]
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=4,
             source="test",
             sample_id="sample_1",
@@ -467,7 +467,7 @@ def test_missing_genotype_encoded_as_minus_one(self):
             anc_alleles,
             schedule=schedule,
         )
-        job = grouping.MatchJob(
+        job = config.MatchJob(
             haplotype_index=2,
             source="test",
             sample_id="sample_0",
@@ -519,7 +519,7 @@ def test_multiple_sources(self):
         schedule = [("src_a", "sample_0", 0), ("src_b", "sample_0", 0)]
         reader = vcz.HaplotypeReader(sources, positions, anc_alleles, schedule=schedule)
 
-        job_a = grouping.MatchJob(
+        job_a = config.MatchJob(
             haplotype_index=2,
             source="src_a",
             sample_id="sample_0",
@@ -529,7 +529,7 @@ def test_multiple_sources(self):
             end_position=200,
             group=2,
         )
-        job_b = grouping.MatchJob(
+        job_b = config.MatchJob(
             haplotype_index=3,
             source="src_b",
             sample_id="sample_0",
@@ -1475,7 +1475,7 @@ def test_facade_returns_alleles_after_reads(self):
         reader = vcz.HaplotypeReader(sources, positions, anc_alleles, schedule=schedule)
 
         # Read from both sources to trigger allele discovery
-        job_a = grouping.MatchJob(
+        job_a = config.MatchJob(
             haplotype_index=0,
             source="src_a",
             sample_id="sample_0",
@@ -1485,7 +1485,7 @@ def test_facade_returns_alleles_after_reads(self):
             end_position=200,
             group=0,
         )
-        job_b = grouping.MatchJob(
+        job_b = config.MatchJob(
             haplotype_index=1,
             source="src_b",
             sample_id="sample_0",

tsinfer/config.py

@@ -161,6 +161,24 @@ class MatchSourceConfig:
     mismatch: float | None = None
 
 
+@dataclasses.dataclass
+class MatchJob:
+    """One row of the compute-groups output — one per haplotype."""
+
+    haplotype_index: int
+    source: str
+    sample_id: str
+    ploidy_index: int
+    time: float
+    start_position: int
+    end_position: int
+    group: int
+    node_flags: int = 1
+    individual_id: int | None = None
+    population_id: int | None = None
+    sample_chunk: int = 0
+
+
 @dataclasses.dataclass
 class MatchConfig:
     """Configuration for the match step."""

tsinfer/grouping.py

@@ -30,29 +30,13 @@
 import numba
 import numpy as np
 
+from . import config
+
 logging.getLogger("numba").setLevel(logging.WARNING)
 logger = logging.getLogger(__name__)
 
 
-@dataclasses.dataclass
-class MatchJob:
-    """One row of the compute-groups output — one per haplotype."""
-
-    haplotype_index: int
-    source: str
-    sample_id: str
-    ploidy_index: int
-    time: float
-    start_position: int
-    end_position: int
-    group: int
-    node_flags: int = 1
-    individual_id: int | None = None
-    population_id: int | None = None
-    sample_chunk: int = 0
-
-
-def assign_groups(jobs: list[MatchJob]) -> list[MatchJob]:
+def assign_groups(jobs: list[config.MatchJob]) -> list[config.MatchJob]:
     """Assign group indices to MatchJobs, return sorted by (group, haplotype_index)."""
     if len(jobs) == 0:
         return jobs

tsinfer/matching.py

@@ -33,7 +33,7 @@
 
 import _tsinfer
 
-from . import arg_ops, config, grouping, vcz
+from . import arg_ops, config, vcz
 
 logger = logging.getLogger(__name__)
 
@@ -286,7 +286,7 @@ def match(
 def extend_ts(
     ts: tskit.TreeSequence,
     *,
-    paired_results: list[tuple[grouping.MatchJob, MatchResult]],
+    paired_results: list[tuple[config.MatchJob, MatchResult]],
     allele_mapper: vcz.AlleleMapper,
     provenance_record: str | None = None,
 ) -> tskit.TreeSequence:

tsinfer/pipeline.py

@@ -64,7 +64,7 @@ def _build_jobs_with_metadata(cfg):
 
     # Build MatchJobs directly from VCZ stores, assigning individual IDs
     # in the same pass.
-    jobs: list[grouping.MatchJob] = []
+    jobs: list[config.MatchJob] = []
     ind_key_to_id: dict[tuple[str, str], int] = {}
     hap_idx = 0
 
@@ -159,7 +159,7 @@ def _build_jobs_with_metadata(cfg):
 
             for p in range(ploidy):
                 jobs.append(
-                    grouping.MatchJob(
+                    config.MatchJob(
                         haplotype_index=hap_idx,
                         source=source.name,
                         sample_id=sid,
@@ -292,10 +292,10 @@ def compute_groups_json(cfg: config.Config) -> str:
     return json_str
 
 
-def _load_match_jobs(path) -> list[grouping.MatchJob]:
+def _load_match_jobs(path) -> list[config.MatchJob]:
     """Read a match-jobs JSON file and return a list of MatchJob objects."""
     records = json.loads(pathlib.Path(path).read_text())
-    return [grouping.MatchJob(**rec) for rec in records]
+    return [config.MatchJob(**rec) for rec in records]
 
 
 def _setup_workdir(workdir, cfg):