EVA-3386: Check Ensembl rapid release for supported assemblies

ebi_eva_common_pyutils/assembly/__init__.py

@@ -1 +1 @@
-from ebi_eva_common_pyutils.assembly.assembly import NCBIAssembly
+from ebi_eva_common_pyutils.reference.assembly import NCBIAssembly

ebi_eva_common_pyutils/assembly/assembly.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from datetime import datetime
+from functools import lru_cache
 
 from ebi_eva_common_pyutils.logger import logging_config
 from ebi_eva_common_pyutils.network_utils import json_request
@@ -29,3 +31,39 @@ def get_supported_asm_from_ensembl(tax_id: int) -> str:
     if assembly_accession_attribute in response:
         return str(response.get(assembly_accession_attribute))
     return None
+
+
+@lru_cache(maxsize=None)
+def get_taxonomy_to_assembly_mapping_from_ensembl_rapid_release():
+    """
+    Returns a dict mapping taxonomy ID to assembly accession, choosing the most recently released,
+    lexicographically last, non-alternate haplotype assembly when multiple are present.
+    """
+    list_data = json_request('https://ftp.ensembl.org/pub/rapid-release/species_metadata.json')
+    results = {}
+    for asm_data in list_data:
+        tax_id = asm_data['taxonomy_id']
+        asm_accession = asm_data['assembly_accession']
+        strain = asm_data['strain']
+        release_date = datetime.strptime(asm_data['release_date'], '%Y-%m-%d')
+
+        # If we haven't seen this taxonomy before, just use this assembly
+        if tax_id not in results:
+            results[tax_id] = (asm_accession, release_date)
+            continue
+
+        # Skip alternate haplotype assemblies
+        if strain and strain.lower() == 'alternate haplotype':
+            continue
+        current_assembly, current_date = results[tax_id]
+        # Keep the more recent assembly, or the lexicographically last one if release dates are equal
+        if current_date < release_date or (current_date == release_date and asm_accession > current_assembly):
+            results[tax_id] = (asm_accession, release_date)
+
+    return {key: val[0] for key, val in results.items()}
+
+
+def get_supported_asm_from_ensembl_rapid_release(tax_id: int) -> str:
+    # TODO: Replace with API call once supported
+    rapid_release_data = get_taxonomy_to_assembly_mapping_from_ensembl_rapid_release()
+    return rapid_release_data.get(tax_id, None)

setup.py

@@ -7,7 +7,7 @@
     name='ebi_eva_common_pyutils',
     scripts=[os.path.join(os.path.dirname(__file__), 'ebi_eva_common_pyutils', 'archive_directory.py')],
     packages=find_packages(),
-    version='0.5.8.dev0',
+    version='0.5.8.dev1',
     license='Apache',
     description='EBI EVA - Common Python Utilities',
     url='https://github.com/EBIVariation/eva-common-pyutils',

tests/test_assembly.py

@@ -2,6 +2,7 @@
 import shutil
 from unittest.mock import Mock
 
+from ebi_eva_common_pyutils.assembly.assembly import get_supported_asm_from_ensembl_rapid_release
 from ebi_eva_common_pyutils.reference.assembly import NCBIAssembly
 from tests.test_common import TestCommon
 
@@ -141,3 +142,22 @@ def test_get_ncbi_genome_folder_url_and_content_multi(self):
                    'assembly_status.txt', 'md5checksums.txt']
 
         self.assertEqual(assembly._ncbi_genome_folder_url_and_content, (url, content))
+
+
+class TestAssembly(TestCommon):
+
+    def test_get_supported_asm_from_ensembl_rapid_release(self):
+        assembly = get_supported_asm_from_ensembl_rapid_release(9117)
+        assert assembly == 'GCA_028858705.1'
+
+    def test_get_supported_asm_from_ensembl_rapid_release_none_found(self):
+        assert get_supported_asm_from_ensembl_rapid_release(0) == None
+
+    def test_get_supported_asm_from_ensembl_rapid_release_multiple_found(self):
+        # Two assemblies, released same date and one alternate haplotype
+        assembly = get_supported_asm_from_ensembl_rapid_release(30194)
+        assert assembly == 'GCA_930367275.1'
+
+        # Three assemblies, released same date and none alternate haplotype
+        assembly = get_supported_asm_from_ensembl_rapid_release(69293)
+        assert assembly == 'GCA_006232285.1'