Add checksum verification for ENSEMBL sequence and annotation wrappers

bio/reference/ensembl-annotation/meta.yaml

@@ -1,4 +1,5 @@
 name: ensembl-annotation
-description: Download annotation of genomic sites (e.g. transcripts) from ENSEMBL FTP servers, and store them in a single .gtf or .gff3 file.
+description: Download annotation of genomic sites (e.g. transcripts) from ENSEMBL FTP servers, and store them in a single .gtf or .gff3 file. After download compare checksums with ENSEMBL FTP 'CHECKSUMS' file.
 authors:
   - Johannes Köster
+  - Fritjof Lammers

bio/reference/ensembl-annotation/test/Snakefile

@@ -1,6 +1,6 @@
 rule get_annotation:
     output:
-        "refs/annotation.gtf"
+        "refs/annotation.gtf.gz"
     params:
         species="homo_sapiens",
         release="98",

bio/reference/ensembl-annotation/wrapper.py

@@ -2,21 +2,45 @@
 __copyright__ = "Copyright 2019, Johannes Köster"
 __email__ = "johannes.koester@uni-due.de"
 __license__ = "MIT"
+__contributor__= "Johannes Köster, Fritjof Lammers"
+
+import sys
+sys.stderr = open(snakemake.log[0], "w")
 
 from ftplib import FTP
+from io import StringIO
+from subprocess import run
+from os.path import basename
+import csv
 
 species = snakemake.params.species.lower()
 release = snakemake.params.release
 fmt = snakemake.params.fmt
 build = snakemake.params.build
 
+
+def verify_checksum(reader):
+    for fields in csv.reader(reader, sep="\t"):
+        cksum = int(fields[0])
+        filename = fields[2]
+        if filename == basename(snakemake.output[0]):
+            cksum_local = int(run(["sum", snakemake.output[0]], capture_output=True).stdout.strip().split()[0])
+            if cksum_local == cksum:
+                print("CHECKSUM OK: %s" % snakemake.output[0])
+                exit(0)
+            else:
+                print("CHECKSUM FAILED: %s" % snakemake.output[0])
+                exit(1)
+        else:
+            print("No matching file for CHECKSUM test found", file=sys.stderr)
+            continue
+
 suffix = ""
 if fmt == "gtf":
     suffix = "gtf.gz"
 elif fmt == "gff3":
     suffix = "gff3.gz"
 
-
 with FTP("ftp.ensembl.org") as ftp, open(snakemake.output[0], "wb") as out:
     ftp.login()
     ftp.retrbinary(
@@ -30,3 +54,18 @@
         ),
         out.write,
     )
+
+    checksum_reader = StringIO()
+    ftp.retrlines(
+        "RETR pub/release-{release}/{fmt}/{species}/CHECKSUMS".format(
+            release=release,
+            build=build,
+            species=species,
+            fmt=fmt,
+            species_cap=species.capitalize(),
+        ),
+        lambda s, w=checksum_reader.write: w(s + '\n'),
+        # use StringIO instance for callback, add "\n" because ftplib.retrlines omits newlines
+    )
+
+    verify_checksum(checksum_reader)

bio/reference/ensembl-sequence/test/Snakefile

@@ -1,6 +1,6 @@
 rule get_sequence:
     output:
-        "refs/genome.fasta"
+        "refs/genome.fasta.gz"
     params:
         species="drosophila_melanogaster",
         datatype="dna",

bio/reference/ensembl-sequence/wrapper.py

@@ -4,13 +4,37 @@
 __license__ = "MIT"
 
 from ftplib import FTP
+from io import StringIO
+from subprocess import run
+from os.path import basename
 
 species = snakemake.params.species.lower()
 release = snakemake.params.release
 build = snakemake.params.build
 
 suffixes = ""
 datatype = snakemake.params.get("datatype", "")
+
+
+def checksum():
+    lines = r.getvalue().strip().split("\n")
+    for line in lines:
+        fields = line.strip().split()
+        cksum = int(fields[0])
+        filename = fields[2]
+        if filename == basename(snakemake.output[0]):
+            cksum_local = int(run(["sum", snakemake.output[0]], capture_output=True).stdout.strip().split()[0])
+            if cksum_local == cksum:
+                print("CHECKSUM OK: %s" % snakemake.output[0])
+                exit(0)
+            else:
+                print("CHECKSUM FAILED: %s" % snakemake.output[0])
+                exit(1)
+        else:
+            # print("No matching file for CHECKSUM test found")
+            continue
+
+
 if datatype == "dna":
     suffixes = ["dna.primary_assembly.fa.gz", "dna.toplevel.fa.gz"]
 elif datatype == "cdna":
@@ -43,6 +67,19 @@
 
         ftp.retrbinary("RETR " + url, out.write)
         success = True
+        # retrieve CHECKSUMS file
+        r = StringIO()
+        ftp.retrlines(
+            "RETR pub/release-{release}/fasta/{species}/{datatype}/CHECKSUMS".format(
+                release=release,
+                species=species,
+                datatype=datatype
+            ),
+            lambda s, w=r.write: w(s + '\n'),
+        )
+        # use StringIO instance for callback, add "\n" because ftplib.retrlines omits newlines
+
+    checksum()
 
 if not success:
     raise ValueError(