Fix 6 correctness bugs found in code review

mhctools/__init__.py

@@ -85,7 +85,7 @@ def __getattr__(name):
     raise AttributeError(
         "module %r has no attribute %r" % (__name__, name))
 
-__version__ = "3.27.0"
+__version__ = "3.28.0"
 
 __all__ = [
     "Prediction",

mhctools/annotate.py

@@ -69,7 +69,9 @@
     "value": (Kind.pMHC_affinity, "value"),
     "presentation": (Kind.pMHC_presentation, "score"),
     "processing": (Kind.antigen_processing, "score"),
-    "stability": (Kind.pMHC_stability, "value"),
+    # NetMHCstabpan reports half-life (Thalf) in `score`, not `value`
+    # (`parse_netmhcstabpan` sets no ic50), so read it from there.
+    "stability": (Kind.pMHC_stability, "score"),
     "immunogenicity": (Kind.immunogenicity, "score"),
     "tap_transport": (Kind.tap_transport, "score"),
     "erap_trimming": (Kind.erap_trimming, "score"),
@@ -217,17 +219,33 @@ def _build_lookup(results, spec):
     ``by_peptide`` maps ``peptide -> Prediction`` (allele-free predictions,
     e.g. processing predictors). Predictions whose target field is ``None``
     are skipped.
+
+    A kind-agnostic token (``score``/``percentile_rank``/``rank``, i.e.
+    ``spec.kind is None``) matches every kind, so if the predictor emits more
+    than one kind for the same key the result would be an arbitrary
+    last-one-wins pick. That is raised as an error rather than silently
+    resolved — the caller should use a kind-specific token instead.
     """
     by_pair = {}
     by_peptide = {}
     for peptide_result in results:
         for pred in peptide_result.filter(kind=spec.kind):
             if getattr(pred, spec.prediction_field) is None:
                 continue
-            if pred.allele:
-                by_pair[(pred.peptide, pred.allele)] = pred
-            else:
-                by_peptide[pred.peptide] = pred
+            target, key = (
+                (by_pair, (pred.peptide, pred.allele)) if pred.allele
+                else (by_peptide, pred.peptide))
+            existing = target.get(key)
+            if (spec.kind is None and existing is not None
+                    and existing.kind != pred.kind):
+                raise ValueError(
+                    "Ambiguous %r field: predictor emits multiple kinds "
+                    "(%s, %s) for peptide %r allele %r. Use a kind-specific "
+                    "token (e.g. affinity, presentation, immunogenicity) "
+                    "instead of %r."
+                    % (spec.field, existing.kind, pred.kind,
+                       pred.peptide, pred.allele, spec.field))
+            target[key] = pred
     return by_pair, by_peptide
 
 

mhctools/cli/args.py

@@ -54,6 +54,9 @@
     NetMHCcons,
     NetMHCstabpan,
     NetChop,
+    NetCleave,
+    NetCleave_I,
+    NetCleave_II,
     Pepsickle,
     RandomBindingPredictor,
     IedbNetMHCpan,
@@ -155,6 +158,9 @@ def __hash__(self):
     "bigmhc-el": _BigMHC_EL,
     "bigmhc-im": _BigMHC_IM,
     "netchop": NetChop,
+    "netcleave": NetCleave,
+    "netcleave-i": NetCleave_I,
+    "netcleave-ii": NetCleave_II,
     "pepsickle": Pepsickle,
     "random": RandomBindingPredictor,
     # use NetMHCpan via IEDB's web API

mhctools/cli/script.py

@@ -91,6 +91,16 @@ def parse_args(args_list=None):
     return arg_parser.parse_args(args_list)
 
 def _run_single_predictor(predictor, args):
+    # The legacy prediction CLI is built on the BindingPrediction model
+    # (predict_peptides / predict_subsequences). New-model-only predictors
+    # (e.g. bigmhc, calis, deeptap, eramer, netchop, pepsickle) implement only
+    # predict(); route the user to the predict-table subcommand rather than
+    # failing later with an opaque AttributeError.
+    if not hasattr(predictor, "predict_peptides"):
+        raise ValueError(
+            "%s does not support this command (it implements the new "
+            "prediction model only). Use `mhctools predict-table` instead."
+            % type(predictor).__name__)
     if args.input_fasta_file:
         input_dictionary = parse_fasta_dictionary(args.input_fasta_file)
         if not input_dictionary:

mhctools/mixmhc2pred.py

@@ -39,7 +39,7 @@
 
 from .base_predictor import BasePredictor, _check_flank_inputs
 from .cleanup_context import CleanupFiles
-from .pred import COLUMNS, Kind, PeptideResult, Prediction
+from .pred import Kind, PeptideResult, Prediction
 from .process_helpers import run_command
 
 
@@ -167,13 +167,6 @@ def predict(self, peptides, n_flanks=None, c_flanks=None):
             for peptide in peptide_list
         ]
 
-    def predict_dataframe(self, peptides, sample_name=""):
-        """``predict()`` flattened to a DataFrame."""
-        dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)]
-        if not dfs:
-            return pd.DataFrame(columns=COLUMNS)
-        return pd.concat(dfs, ignore_index=True)
-
     def _default_pred_kind(self):
         return Kind.pMHC_presentation
 

mhctools/nettcr.py

@@ -62,15 +62,21 @@ def _suppress_native_stderr():
     the prediction loop, so the window is small.
     """
     sys.stderr.flush()
-    saved_fd = os.dup(2)
-    devnull_fd = os.open(os.devnull, os.O_WRONLY)
+    # Acquire both descriptors inside the try so a failure partway through
+    # (e.g. fd exhaustion at os.open) can't leak the one already taken.
+    saved_fd = None
+    devnull_fd = None
     try:
+        saved_fd = os.dup(2)
+        devnull_fd = os.open(os.devnull, os.O_WRONLY)
         os.dup2(devnull_fd, 2)
         yield
     finally:
-        os.dup2(saved_fd, 2)
-        os.close(devnull_fd)
-        os.close(saved_fd)
+        if saved_fd is not None:
+            os.dup2(saved_fd, 2)
+            os.close(saved_fd)
+        if devnull_fd is not None:
+            os.close(devnull_fd)
 
 
 # BLOSUM50, restricted to the 20 standard amino acids, in NetTCR's column

mhctools/prime.py

@@ -45,7 +45,7 @@
 from .allele_normalization import normalize_allele_name
 from .base_predictor import BasePredictor, _check_flank_inputs
 from .cleanup_context import CleanupFiles
-from .pred import COLUMNS, Kind, PeptideResult, Prediction
+from .pred import Kind, PeptideResult, Prediction
 from .process_helpers import run_command
 
 
@@ -146,13 +146,6 @@ def predict(self, peptides, n_flanks=None, c_flanks=None):
             for peptide in peptide_list
         ]
 
-    def predict_dataframe(self, peptides, sample_name=""):
-        """``predict()`` flattened to a DataFrame."""
-        dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)]
-        if not dfs:
-            return pd.DataFrame(columns=COLUMNS)
-        return pd.concat(dfs, ignore_index=True)
-
     def _default_pred_kind(self):
         return Kind.immunogenicity
 

tests/test_annotate_table.py

@@ -444,6 +444,87 @@ def test_output_field_tokens_include_primary_three():
         assert expected in tokens
 
 
+# --- stability token reads `score` (NetMHCstabpan puts Thalf there) ---------
+
+class _StabilityFixturePredictor:
+    """Allele-bearing predictor emitting pMHC_stability with half-life in
+    ``score`` and no ``value`` — mirroring parse_netmhcstabpan."""
+
+    def __init__(self, alleles):
+        self.alleles = list(alleles) if alleles else []
+
+    def predict(self, peptides):
+        results = []
+        for peptide in peptides:
+            preds = [
+                Prediction(
+                    kind=Kind.pMHC_stability,
+                    peptide=peptide,
+                    allele=allele,
+                    score=12.5,          # half-life (hours), higher = better
+                    value=None,          # NetMHCstabpan sets no ic50
+                    predictor_name="stab-fixture")
+                for allele in self.alleles]
+            results.append(PeptideResult(preds=tuple(preds)))
+        return results
+
+
+def test_stability_token_reads_score_not_value():
+    # Regression: the `stability` token used to map to `value`, which
+    # NetMHCstabpan never populates, so the column was always NaN.
+    spec = parse_annotation_spec("netmhcstabpan:stab:stability")
+    assert (spec.kind, spec.prediction_field) == (Kind.pMHC_stability, "score")
+
+    out = annotate_table(
+        _table(),
+        [AnnotationSpec(lambda a: _StabilityFixturePredictor(a), "stab",
+                        field="stability")],
+        allele_column="hla")
+    assert list(out["stab"]) == [12.5, 12.5]  # not NaN
+
+
+# --- generic score/rank token is rejected on multi-kind predictors ----------
+
+class _MultiKindFixturePredictor:
+    """Emits BOTH affinity and presentation for the same (peptide, allele)."""
+
+    def __init__(self, alleles):
+        self.alleles = list(alleles) if alleles else []
+
+    def predict(self, peptides):
+        results = []
+        for peptide in peptides:
+            preds = []
+            for allele in self.alleles:
+                preds.append(Prediction(
+                    kind=Kind.pMHC_affinity, peptide=peptide, allele=allele,
+                    score=0.2, value=100.0, predictor_name="multi"))
+                preds.append(Prediction(
+                    kind=Kind.pMHC_presentation, peptide=peptide, allele=allele,
+                    score=0.9, predictor_name="multi"))
+            results.append(PeptideResult(preds=tuple(preds)))
+        return results
+
+
+def test_generic_score_token_ambiguous_multikind_raises():
+    with pytest.raises(ValueError, match="[Aa]mbiguous"):
+        annotate_table(
+            _table(),
+            [AnnotationSpec(lambda a: _MultiKindFixturePredictor(a), "s",
+                            field="score")],
+            allele_column="hla")
+
+
+def test_kind_specific_token_unambiguous_on_multikind():
+    # The same multi-kind predictor is fine with a kind-specific token.
+    out = annotate_table(
+        _table(),
+        [AnnotationSpec(lambda a: _MultiKindFixturePredictor(a), "pres",
+                        field="presentation")],
+        allele_column="hla")
+    assert list(out["pres"]) == [0.9, 0.9]
+
+
 # --- integration-flavored smoke test with the real random predictor ---------
 
 def test_random_predictor_smoke():

tests/test_cli_registry.py

@@ -317,3 +317,23 @@ def test_legacy_binding_predictor_rejects_multiple():
     args = _make_args(["random", "random"])
     with pytest.raises(ValueError, match="exactly one"):
         mhc_binding_predictor_from_args(args)
+
+
+def test_netcleave_family_registered():
+    from mhctools import NetCleave, NetCleave_I, NetCleave_II
+    assert mhc_predictors["netcleave"] == NetCleave
+    assert mhc_predictors["netcleave-i"] == NetCleave_I
+    assert mhc_predictors["netcleave-ii"] == NetCleave_II
+
+
+def test_legacy_cli_rejects_new_model_only_predictor():
+    # New-model-only predictors (predict() but no predict_peptides) must give a
+    # clear "use predict-table" error on the legacy CLI, not an AttributeError.
+    from mhctools.cli.script import _run_single_predictor
+
+    class _NewModelOnly:
+        def predict(self, peptides):
+            return []
+
+    with pytest.raises(ValueError, match="predict-table"):
+        _run_single_predictor(_NewModelOnly(), Namespace())

tests/test_mixmhc2pred2.py

@@ -159,3 +159,9 @@ def test_mixmhc2pred_end_to_end():
             assert pred.score is not None
             assert pred.percentile_rank is not None
         assert result.presentation is not None
+
+
+def test_predict_dataframe_inherits_flank_signature():
+    import inspect
+    params = inspect.signature(MixMHC2pred.predict_dataframe).parameters
+    assert "n_flanks" in params and "c_flanks" in params

tests/test_prime.py

@@ -164,3 +164,12 @@ def test_prime_dataframe_end_to_end():
     assert len(df) == 2
     assert set(df["kind"]) == {Kind.immunogenicity}
     assert (df["allele"] == "HLA-A*02:01").all()
+
+
+def test_predict_dataframe_inherits_flank_signature():
+    # Regression: PRIME used to override predict_dataframe with a narrower
+    # signature that dropped n_flanks/c_flanks (a TypeError trap). It should
+    # inherit the base method, which forwards flanks to predict().
+    import inspect
+    params = inspect.signature(PRIME.predict_dataframe).parameters
+    assert "n_flanks" in params and "c_flanks" in params