From 81d5919a12a7f18d8e129c1f7bef13cabe63b849 Mon Sep 17 00:00:00 2001
From: Alex Rubinsteyn <alex.rubinsteyn@gmail.com>
Date: Wed, 15 Apr 2026 14:29:07 -0400
Subject: [PATCH] Route predict_peptides_dataframe through canonical schema
 (fixes #193)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

predict_peptides_dataframe (deprecated) returned the legacy
BindingPrediction schema — missing predictor_version, kind, value, and
using prediction_method_name instead of predictor_name. Meanwhile
predict_proteins_dataframe / predict_dataframe emit the canonical
mhctools.pred.COLUMNS schema.

The asymmetry forced downstream consumers that treat predictions from
either path uniformly (e.g. topiary's CachedPredictor design) to
None-fill the missing identity columns.

Fix: predict_peptides_dataframe now delegates to predict_dataframe,
so both batch-dataframe paths emit the same columns.

This is a schema change on an already-deprecated method. Callers that
relied on the old legacy columns (affinity, prediction_method_name,
length) should migrate to the canonical names (value, predictor_name)
and add predictor_version / kind.
---
 mhctools/base_predictor.py            | 13 ++++++--
 tests/test_dataframe_schema_parity.py | 46 +++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_dataframe_schema_parity.py

diff --git a/mhctools/base_predictor.py b/mhctools/base_predictor.py
index 7309ac7..cf96d60 100644
--- a/mhctools/base_predictor.py
+++ b/mhctools/base_predictor.py
@@ -197,8 +197,17 @@ def predict_peptides(self, peptides):
             "%s must implement predict_peptides" % (self.__class__.__name__,))
 
     def predict_peptides_dataframe(self, peptides):
-        """Deprecated: use predict_dataframe() instead."""
-        return self.predict_peptides(peptides).to_dataframe()
+        """Deprecated: use predict_dataframe() instead.
+
+        Emits the canonical prediction schema (see ``mhctools.pred.COLUMNS``)
+        for parity with ``predict_proteins_dataframe`` / ``predict_dataframe``.
+        Previously this returned the legacy BindingPrediction schema
+        (``source_sequence_name, offset, peptide, allele, score, affinity,
+        percentile_rank, prediction_method_name, length``). That schema lacked
+        ``predictor_version``, ``kind``, ``value`` and used
+        ``prediction_method_name`` instead of ``predictor_name`` — see #193.
+        """
+        return self.predict_dataframe(peptides)
 
     def _check_peptide_lengths(self, peptide_lengths=None):
         """
diff --git a/tests/test_dataframe_schema_parity.py b/tests/test_dataframe_schema_parity.py
new file mode 100644
index 0000000..5319f0c
--- /dev/null
+++ b/tests/test_dataframe_schema_parity.py
@@ -0,0 +1,46 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+
+"""Regression test for openvax/mhctools#193.
+
+predict_peptides_dataframe (legacy/deprecated) and predict_proteins_dataframe
+(current) must emit the same canonical column schema so downstream code can
+treat rows from either path uniformly.
+"""
+
+from mhctools import RandomBindingPredictor
+from mhctools.pred import COLUMNS
+
+
+def test_peptides_and_proteins_dataframes_share_schema():
+    p = RandomBindingPredictor(
+        alleles=["HLA-A*02:01"], default_peptide_lengths=[9],
+    )
+    df_peptides = p.predict_peptides_dataframe(["SIINFEKLA"])
+    df_proteins = p.predict_proteins_dataframe({"src": "MASIINFEKLA"})
+
+    assert list(df_peptides.columns) == list(COLUMNS), (
+        "predict_peptides_dataframe must emit canonical COLUMNS schema "
+        "(was previously missing predictor_version, kind, value and used "
+        "prediction_method_name; see #193)"
+    )
+    assert list(df_proteins.columns) == list(COLUMNS)
+    assert list(df_peptides.columns) == list(df_proteins.columns)
+
+
+def test_peptides_dataframe_has_predictor_identity_columns():
+    """Downstream (e.g. topiary CachedPredictor) needs a stable
+    (predictor_name, predictor_version) identity on every row."""
+    p = RandomBindingPredictor(alleles=["HLA-A*02:01"], default_peptide_lengths=[9])
+    df = p.predict_peptides_dataframe(["SIINFEKLA"])
+    assert "predictor_name" in df.columns
+    assert "predictor_version" in df.columns
+    assert "kind" in df.columns
+    assert "value" in df.columns
+    # legacy name must not reappear
+    assert "prediction_method_name" not in df.columns
+    # canonical name is populated (RandomBindingPredictor sets it)
+    assert df["predictor_name"].iloc[0] != ""