feat: persist variant genomes with a deterministic ID

genome_kit/variant_genome.py

@@ -1,13 +1,21 @@
 # Copyright (C) 2016-2023 Deep Genomics Inc. All Rights Reserved.
 from __future__ import absolute_import
 
+import gzip
+import hashlib
+import os
 import pickle
+import re
+import tempfile
 from functools import partial
 from operator import attrgetter
 
 from . import _util
+from . import gk_data
 from . import interval as _interval
 from ._apply_variants import apply_variants, check_variants_list
+from .genome import Genome
+from .variant import Variant
 
 
 class VariantGenome(object):
@@ -21,10 +29,22 @@ class VariantGenome(object):
     will be aligned to the variant genome's coordinate system.
 
     TODO: see explanation of anchors.
+
+    A variant genome can be persisted and recreated later, in another process
+    or on another machine, from its content-derived :py:attr:`id`::
+
+        >>> variant_genome = VariantGenome(genome, variants)
+        >>> variant_genome_id = variant_genome.save()  # keep this id, e.g. in a database
+        >>> VariantGenome.load(variant_genome_id) == variant_genome
+        True
+
+    Saving requires a writable data manager; see :py:meth:`save`.
     """
 
     __slots__ = ('genome', 'variants')
 
+    _FILENAME_FORMAT = "variant_genome.{}.txt.gz"
+
     def __init__(self, reference_genome, variants):
         """Initialize a variant genome.
 
@@ -204,20 +224,83 @@ def __repr__(self):
     def reference_genome(self):
         return self.genome.reference_genome
 
+    def _canonical_variants(self) -> list[Variant]:
+        # Normalized variants in a deterministic order, shared by `__eq__` and `id`.
+        return sorted((vv for v in self.variants for vv in v._normalized_variant),
+                      key=attrgetter("chromosome", "start", "ref", "alt"))
+
+    @property
+    def id(self) -> str:
+        """A content-derived identifier for this variant genome (hex SHA-256 digest of the genome config and the normalized variants).
+
+        Two variant genomes have the same id if and only if they compare equal,
+        i.e. they have the same genome config and the same variants after normalization.
+        """
+        hasher = hashlib.sha256(self.genome.config.encode())
+        for variant in self._canonical_variants():
+            hasher.update(b"\n" + variant.as_variant_string().encode())
+        return hasher.hexdigest()
+
+    def save(self) -> str:
+        """Save this variant genome so that it can be recreated with :py:meth:`load`.
+        Returns the :py:attr:`id` to pass to :py:meth:`load`.
+
+        In the current implementation, a file is uploaded per-variant genome with ``genome_kit.gk_data.data_manager``.
+        Equal variant genomes share a single file.
+
+        The default open source data manager points at GenomeKit's public bucket, which is
+        read-only. To save, bring your own data manager (see :ref:`sharing-data`).
+        """
+        variant_genome_id = self.id
+        filename = self._FILENAME_FORMAT.format(variant_genome_id)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, filename)
+            lines = [self.genome.config] + [v.as_variant_string() for v in self._canonical_variants()]
+            # Fixed gzip header so equal genomes produce identical bytes
+            with open(path, "wb") as raw, gzip.GzipFile(fileobj=raw, mode="wb", filename="", mtime=0) as f:
+                f.write("".join(line + "\n" for line in lines).encode())
+            gk_data.data_manager.upload_file(path, filename)
+        return variant_genome_id
+
+    @classmethod
+    def load(cls, variant_genome_id: str) -> "VariantGenome":
+        """Recreate a variant genome that was stored with :py:meth:`save`.
+        Returns a variant genome equal to the one that was saved.
+
+        In the current implementation, a file is fetched with ``genome_kit.gk_data.data_manager``.
+
+        Parameters
+        ----------
+        variant_genome_id : :py:class:`str`
+            The id returned by :py:meth:`save`.
+
+        Raises
+        ------
+        :py:class:`~genome_kit.data_manager.GKDataFileNotFoundError`
+            No variant genome with this id was found.
+        :py:exc:`ValueError`
+            `variant_genome_id` is not an id produced by :py:meth:`save`, or the stored file does not match it.
+        """
+        if not re.fullmatch("[0-9a-f]{64}", variant_genome_id):
+            raise ValueError("Invalid variant genome id: {!r}".format(variant_genome_id))
+        path = gk_data.data_manager.get_file(cls._FILENAME_FORMAT.format(variant_genome_id))
+        with gzip.open(path, "rt") as f:
+            genome = Genome(f.readline().rstrip("\n"))
+            variants = []
+            for line in f:
+                chrom, pos, ref, alt = line.rstrip("\n").split(":")
+                variants.append(Variant(chrom, int(pos) - 1, ref, alt, genome))
+        variant_genome = cls(genome, variants)
+        if variant_genome.id != variant_genome_id:
+            raise ValueError("Stored variant genome does not match id {}".format(variant_genome_id))
+        return variant_genome
+
     def __eq__(self, other):
         if not isinstance(other, VariantGenome):
             return False
-        if self.genome != other.genome:
-            return False
-        # sort and normalize both sides.
-        # `start` is ok to use here because VariantGenome doesn't support overlapping variants in any case.
-        vself = sorted((vv for n in (v._normalized_variant for v in self.variants) for vv in n),
-                          key=attrgetter('start'))
-        vother = sorted((vv for n in (v._normalized_variant for v in other.variants) for vv in n),
-                          key=attrgetter('start'))
-        if vself != vother:
+        if self.genome.config != other.genome.config:
             return False
-        return True
+        return self._canonical_variants() == other._canonical_variants()
 
     def __getattr__(self, name):
         # Intercept any attribute requests that weren't found on the VariantGenome object itself,

src/gk_assert.h

@@ -6,6 +6,7 @@ Copyright (C) 2016-2023 Deep Genomics Inc. All Rights Reserved.
 #define __GK_ASSERT_H__
 
 #include "defines.h"
+#include <exception>
 #include <format>
 #include <stdexcept>
 #include <string>

tests/test_variant_genome.py

@@ -2,11 +2,34 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+import gzip
+import os
+import shutil
+import tempfile
 import unittest
+from unittest.mock import patch
+from genome_kit import DataManager
 from genome_kit import VariantGenome
 from genome_kit import Interval
 from genome_kit import Variant
-from . import MiniGenome
+from genome_kit import gk_data
+from genome_kit.data_manager import GKDataFileNotFoundError
+from . import MiniGenome, get_test_file
+
+
+class LocalDataManager(DataManager):
+    """Stores uploads in a local directory, serving other files from the test data."""
+
+    def get_file(self, filename):
+        path = os.path.join(self.data_dir, filename)
+        if os.path.exists(path):
+            return path
+        if filename.startswith(VariantGenome._FILENAME_FORMAT.split("{")[0]):
+            raise GKDataFileNotFoundError(filename, "not found")
+        return get_test_file(filename)
+
+    def upload_file(self, filepath, filename, metadata=None):
+        shutil.copyfile(filepath, os.path.join(self.data_dir, filename))
 
 
 class TestVariantGenome(unittest.TestCase):
@@ -265,6 +288,67 @@ def test_string_variant(self):
         with self.assertRaises(TypeError):
             VariantGenome(self.genome, ["chr1:15:G:T"])
 
+    def _variant(self, chrom, start, alt, ref_len=1, genome=None):
+        genome = genome or self.genome
+        ref = genome.dna(Interval(chrom, "+", start, start + ref_len, genome))
+        return Variant(chrom, start, ref, alt, genome)
+
+    def test_eq(self):
+        # Same start on different chromosomes, in different orders
+        a = self._variant("chr1", 10, "T")
+        b = self._variant("chr2", 10, "C")
+        self.assertEqual(VariantGenome(self.genome, [a, b]), VariantGenome(self.genome, [b, a]))
+
+        # Un-normalized variant equals its normalized form
+        ref = self.genome.dna(Interval("chr1", "+", 10, 12, self.genome))
+        # G-padded = "11:GT:GC", non-padded = "12:T:C"
+        padded = Variant("chr1", 10, ref, ref[0] + "C", self.genome)
+        non_padded = self._variant("chr1", 11, "C")
+        self.assertEqual(VariantGenome(self.genome, padded), VariantGenome(self.genome, non_padded))
+
+        # different variants
+        self.assertNotEqual(VariantGenome(self.genome, a), VariantGenome(self.genome, b))
+        # different genome
+        self.assertNotEqual(VariantGenome(self.genome, a), VariantGenome(MiniGenome("hg19"), self._variant("chr1", 10, "T", genome=MiniGenome("hg19"))))
+
+    def test_id(self):
+        a = self._variant("chr1", 10, "T")
+        b = self._variant("chr2", 10, "C")
+        vg = VariantGenome(self.genome, [a, b])
+        self.assertRegex(vg.id, "^[0-9a-f]{64}$")
+        self.assertEqual(vg.id, VariantGenome(self.genome, [b, a]).id)
+        self.assertNotEqual(vg.id, VariantGenome(self.genome, a).id)
+        self.assertNotEqual(vg.id, VariantGenome(self.genome, []).id)
+
+    def test_save_load(self):
+        variants = [self._variant("chr1", 10, "T"), self._variant("chr2", 10, "C"), self._variant("chr1", 20, "", ref_len=2)]
+        vg = VariantGenome(self.genome, variants)
+        with tempfile.TemporaryDirectory() as tmpdir, patch.object(gk_data, "data_manager", LocalDataManager(tmpdir)):
+            vg_id = vg.save()
+            self.assertEqual(vg_id, vg.id)
+            with open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "rb") as f:
+                saved = f.read()
+            VariantGenome(self.genome, variants[::-1]).save()
+            with open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "rb") as f:
+                self.assertEqual(f.read(), saved)
+            loaded = VariantGenome.load(vg_id)
+            self.assertEqual(loaded, vg)
+            self.assertEqual(loaded.id, vg_id)
+            self.assertEqual(loaded.dna(Interval("chr1", "+", 5, 25, loaded)), vg.dna(Interval("chr1", "+", 5, 25, vg)))
+
+            with self.assertRaises(GKDataFileNotFoundError):
+                VariantGenome.load("0" * 64)
+
+            for bad_id in ["", vg_id[:-1], vg_id.upper(), "../" + vg_id, vg_id + "\n"]:
+                with self.assertRaisesRegex(ValueError, "Invalid variant genome id"):
+                    VariantGenome.load(bad_id)
+
+            # Stored content that does not match its id is rejected
+            with gzip.open(os.path.join(tmpdir, VariantGenome._FILENAME_FORMAT.format(vg_id)), "wt") as f:
+                f.write(self.genome.config + "\n")
+            with self.assertRaisesRegex(ValueError, "does not match id"):
+                VariantGenome.load(vg_id)
+
     def test_allow_outside_chromsome(self):
         variant = VariantGenome(
             self.genome, Variant.from_string("chr1:15:G:T", self.genome)