Preserve FASTA-header versions in SequenceData (closes #351)

pyensembl/database.py

@@ -631,6 +631,14 @@ def _load_gtf_as_dataframe(self, usecols=None, features=None):
             # 2.7.0+) renames the GENCODE columns onto the Ensembl names at
             # parse time so the rest of pyensembl can stay Ensembl-centric.
             attribute_aliases=GENCODE_BIOTYPE_ALIASES,
+            # Opt out of gtfparse's Int64 cast for *_version columns: when
+            # any row is missing a version (e.g. start_codon rows that
+            # don't carry transcript_version), pandas' nullable Int64
+            # routes through float on the sqlite write path and stores as
+            # "7.0" text, which then breaks our `int(result[0])` parse.
+            # pyensembl handles the string→int conversion itself on the
+            # `*.version` property side.
+            cast_version_columns=False,
             infer_biotype_column=True,
             usecols=usecols,
             features=features,

pyensembl/fasta.py

@@ -29,7 +29,17 @@
 def _parse_header_id(line):
     """
     Pull the transcript or protein identifier from the header line
-    which starts with '>'
+    which starts with '>'.
+
+    The full versioned form (e.g. ``ENSP00000123456.3``) is returned when
+    the header carries a version. Stripping happens at lookup time via
+    :func:`pyensembl.sequence_data.lookup_sequence_with_version_fallback`,
+    not here, so the FASTA-header version is preserved as the authoritative
+    identity of the sequence.
+
+    Non-ENS IDs (e.g. TAIR ``AT1G01010.1``, where ``.1`` is an isoform
+    suffix rather than a version) are returned verbatim — this function
+    does no surgery on them.
     """
     if type(line) is not bytes:
         raise TypeError(
@@ -47,22 +57,36 @@ def _parse_header_id(line):
     else:
         identifier = line[1:]
 
-    # annoyingly Ensembl83 reformatted the transcript IDs of its
-    # cDNA FASTA to include sequence version numbers
-    # .e.g.
-    # "ENST00000448914.1" instead of "ENST00000448914"
-    # So now we have to parse out the identifier
-
-    # only split name of ENSEMBL naming. In other database, such as TAIR,
-    # the '.1' notation is the isoform not the version.
-    if identifier.startswith(b"ENS"):
-        dot_index = identifier.find(b".")
-        if dot_index >= 0:
-            identifier = identifier[:dot_index]
+    # GENCODE protein/cDNA FASTAs use a pipe-delimited header packing the
+    # versioned protein ID, transcript ID, gene ID, etc., e.g.
+    #   >ENSP00000493376.2|ENST00000641515.2|ENSG00000186092.7|...
+    # The leading field is the canonical identifier for this sequence;
+    # everything after the first '|' is cross-reference metadata.
+    pipe_index = identifier.find(b"|")
+    if pipe_index >= 0:
+        identifier = identifier[:pipe_index]
 
     return identifier.decode("ascii")
 
 
+def _split_ens_version(identifier):
+    """
+    Split an ENS-prefix identifier into ``(bare_id, version_int)``.
+
+    Returns ``(identifier, None)`` for IDs that don't carry a parseable
+    ENS version. Non-ENS IDs (e.g. TAIR ``AT1G01010.1``) are always
+    returned as-is with version ``None`` — the ``.N`` in those is an
+    isoform suffix, not a version.
+    """
+    if not identifier or not identifier.startswith("ENS") or "." not in identifier:
+        return identifier, None
+    bare, _, suffix = identifier.rpartition(".")
+    try:
+        return bare, int(suffix)
+    except ValueError:
+        return identifier, None
+
+
 class FastaParser(object):
     """
     FastaParser object consumes lines of a FASTA file incrementally

pyensembl/protein.py

@@ -19,9 +19,14 @@ class Protein(Serializable):
     Accessed via :attr:`Transcript.protein`.
     """
 
-    def __init__(self, protein_id, protein_version=None):
+    def __init__(self, protein_id, protein_version=None, genome=None):
         self.protein_id = protein_id
         self.protein_version = protein_version
+        # Optional genome reference so ``fasta_version`` can consult the
+        # protein FASTA's :class:`SequenceData`. Defaults to ``None`` for
+        # back-compat (callers constructing ``Protein`` outside of
+        # :class:`Transcript` won't have a meaningful genome).
+        self.genome = genome
 
     @property
     def id(self):
@@ -45,6 +50,25 @@ def versioned_id(self):
         """Alias for :attr:`versioned_protein_id`."""
         return self.versioned_protein_id
 
+    @property
+    def fasta_version(self):
+        """
+        Integer version that the protein FASTA header carried for this
+        ``protein_id``, or ``None`` if the FASTA didn't carry one (older
+        Ensembl releases shipped bare headers) or the genome has no
+        protein FASTA attached.
+
+        Differs from :attr:`protein_version` (which comes from the GTF's
+        ``protein_version`` attribute). When the two disagree, the FASTA
+        version is the authoritative source-of-truth for the bytes
+        returned by :meth:`Transcript.protein_sequence`.
+        """
+        if self.genome is None:
+            return None
+        if not self.genome.requires_protein_fasta:
+            return None
+        return self.genome.protein_sequences.fasta_version(self.protein_id)
+
     def __eq__(self, other):
         return (
             other.__class__ is Protein

pyensembl/sequence_data.py

@@ -18,40 +18,55 @@
 from collections import Counter
 import pickle
 from .common import load_pickle, dump_pickle
-from .fasta import parse_fasta_dictionary
+from .fasta import _split_ens_version, parse_fasta_dictionary
 
 
 logger = logging.getLogger(__name__)
 
 
 def lookup_sequence_with_version_fallback(sequence_data, identifier):
     """
-    Look up ``identifier`` in ``sequence_data``, tolerating an ENS.N version
-    suffix mismatch between the caller's ID and the FASTA's dict keys.
+    Look up ``identifier`` in ``sequence_data``, tolerating ENS ``.N``
+    version-suffix mismatches in either direction.
 
     Both Ensembl and GENCODE work with versioned IDs (e.g.
     ``ENSP00000123456.3``); the two formats just split that information
     differently. Ensembl GTFs keep the bare ID in ``protein_id`` /
     ``transcript_id`` and put the version in a separate ``*_version``
     attribute, while GENCODE GTFs embed the version directly in the ID
-    attribute. pyensembl's FASTA parser strips ``.N`` from ENS-prefix
-    headers, so the FASTA dict keys are always unversioned — meaning a
-    literal lookup of a GENCODE-style versioned ID would miss.
-
-    Strategy: try ``identifier`` as-is first (covers both unversioned IDs
-    and any future FASTA that preserves versions); on miss, strip a
-    trailing ``.N`` suffix and try again, but only for ENS-prefix IDs
-    because for non-Ensembl IDs (e.g. TAIR ``AT1G01010.1``) the ``.N`` is
-    an isoform suffix and stripping would be incorrect.
+    attribute. Whichever convention the FASTA on disk used, the
+    sequence dict now keys on the form that header carried — so either
+    a versioned or a bare caller-supplied ID may need a fallback.
+
+    Resolution order:
+      1. literal lookup of ``identifier``
+      2. if ``identifier`` is a versioned ENS ID, strip ``.N`` and retry
+         (handles caller-supplied versioned ID against a bare FASTA)
+      3. consult ``sequence_data._stripped_index`` for a versioned alias
+         of a bare caller-supplied ID (handles bare ID against a
+         versioned FASTA, the GENCODE case)
+      4. None
+
+    Non-Ensembl IDs (e.g. TAIR ``AT1G01010.1`` where ``.N`` is an
+    isoform suffix, not a version) skip step (2): stripping there is
+    semantically wrong.
     """
     if not identifier:
         return None
     sequence = sequence_data.get(identifier)
     if sequence is not None:
         return sequence
     if identifier.startswith("ENS") and "." in identifier:
-        stripped = identifier.rsplit(".", 1)[0]
-        return sequence_data.get(stripped)
+        bare, version = _split_ens_version(identifier)
+        if version is not None:
+            sequence = sequence_data.get(bare)
+            if sequence is not None:
+                return sequence
+    stripped_index = getattr(sequence_data, "_stripped_index", None)
+    if stripped_index:
+        versioned = stripped_index.get(identifier)
+        if versioned is not None:
+            return sequence_data.get(versioned)
     return None
 
 
@@ -88,6 +103,14 @@ def __init__(self, fasta_paths, cache_directory_path=None):
     def _init_lazy_fields(self):
         self._fasta_dictionary = None
         self._fasta_keys = None
+        # Maps a bare ENS ID -> the versioned form actually keyed in
+        # _fasta_dictionary, so callers passing the unversioned form can
+        # still resolve to a versioned FASTA entry (GENCODE case).
+        self._stripped_index = None
+        # Maps a sequence identifier -> the integer version suffix the
+        # FASTA header carried, or absent for headers without a version.
+        # Lets callers sanity-check FASTA / GTF alignment.
+        self._versions = None
 
     def clear_cache(self):
         self._init_lazy_fields()
@@ -125,9 +148,25 @@ def _add_to_fasta_dictionary(self, fasta_dictionary_tmp):
                 )
                 continue
             self._fasta_dictionary[identifier] = sequence
+            bare, version = _split_ens_version(identifier)
+            if version is not None:
+                if bare in self._stripped_index:
+                    # Two FASTA entries collide on the same bare ENS ID with
+                    # different versions. The newer of the two is kept as
+                    # the canonical alias since version numbers grow
+                    # monotonically.
+                    existing = self._stripped_index[bare]
+                    _, existing_version = _split_ens_version(existing)
+                    if existing_version is None or version > existing_version:
+                        self._stripped_index[bare] = identifier
+                else:
+                    self._stripped_index[bare] = identifier
+                self._versions[identifier] = version
 
     def _load_or_create_fasta_dictionary_pickle(self):
         self._fasta_dictionary = dict()
+        self._stripped_index = dict()
+        self._versions = dict()
         for fasta_path, pickle_path in zip(
             self.fasta_paths, self.fasta_dictionary_pickle_paths
         ):
@@ -169,3 +208,24 @@ def fasta_dictionary(self):
     def get(self, sequence_id):
         """Get sequence associated with given ID or return None if missing"""
         return self.fasta_dictionary.get(sequence_id)
+
+    def fasta_version(self, sequence_id):
+        """
+        Return the integer FASTA-header version for ``sequence_id``, or
+        ``None`` if the header didn't carry a version (e.g. older Ensembl
+        releases) or the ID isn't in the FASTA.
+
+        Accepts either the versioned or bare form of an ENS ID — the
+        bare form is resolved through ``_stripped_index``.
+        """
+        if not sequence_id:
+            return None
+        # ensure lazy load
+        _ = self.fasta_dictionary
+        version = self._versions.get(sequence_id)
+        if version is not None:
+            return version
+        versioned = self._stripped_index.get(sequence_id)
+        if versioned is not None:
+            return self._versions.get(versioned)
+        return None

pyensembl/transcript.py

@@ -591,7 +591,9 @@ def protein(self):
             if result and result[0]:
                 protein_version = int(result[0])
         return Protein(
-            protein_id=self.protein_id, protein_version=protein_version
+            protein_id=self.protein_id,
+            protein_version=protein_version,
+            genome=self.genome,
         )
 
     @memoized_property
@@ -601,3 +603,20 @@ def protein_sequence(self):
         return lookup_sequence_with_version_fallback(
             self.genome.protein_sequences, self.protein_id
         )
+
+    @property
+    def fasta_version(self):
+        """
+        Integer version that the cDNA FASTA header carried for this
+        transcript, or ``None`` if the FASTA didn't carry one (older
+        Ensembl releases shipped bare headers) or no transcript FASTA
+        is attached to the genome.
+
+        Differs from :attr:`transcript_version` (which comes from the
+        GTF's ``transcript_version`` attribute). When the two disagree,
+        the FASTA version is the authoritative source-of-truth for the
+        bytes returned by :attr:`sequence`.
+        """
+        if not self.genome.requires_transcript_fasta:
+            return None
+        return self.genome.transcript_sequences.fasta_version(self.transcript_id)

pyensembl/version.py

@@ -1,4 +1,4 @@
-__version__ = "2.9.8"
+__version__ = "2.10.0"
 
 def print_version():
     print(f"v{__version__}")