Merge pull request #509 from padix-key/application

Amend:  Increase performance of fetching via sra-tools (#504)

Author: padix-key

doc/examples/scripts/sequence/quality_control.py

@@ -32,9 +32,12 @@
 FIG_SIZE = (8.0, 6.0)
 
 
+app = sra.FastqDumpApp("ERR266411")
+app.start()
+app.join()
 # Each run can have multiple reads per spot
 # by selecting index 0 we take only the first read for every spot
-sequences_and_scores = sra.FastqDumpApp.fetch("ERR266411")[0]
+sequences_and_scores = app.get_sequences_and_scores()[0]
 sequence_codes = np.stack([
     sequence.code for sequence, _ in sequences_and_scores.values()
 ])
@@ -102,7 +105,7 @@
 
 ########################################################################
 # This is a typical distribution.
-# 
+#
 # Now we want to see the appearance of each base over the length of the
 # sequence reads.
 # In a random library one would expect, that :math:`p(A) \approx p(T)`
@@ -209,7 +212,7 @@
 ax.bar(
     np.arange(0, len(duplication_level_freq)),
     duplication_level_freq,
-    width=0.6, 
+    width=0.6,
     color=biotite.colors["dimorange"]
 )
 ax.set_xlim(0.5, len(duplication_level_freq) + 0.5)
@@ -218,6 +221,8 @@
 ax.set_ylabel("Sequence percentage (%)")
 fig.tight_layout()
 
+plt.show()
+
 ########################################################################
 # The dataset has quite an unusual repetition profile:
 # Usually one would expect, that most sequences occur only once and the

src/biotite/application/sra/app.py

@@ -7,10 +7,10 @@
 __all__ = ["FastaDumpApp", "FastqDumpApp"]
 
 import abc
+from os.path import join
 from subprocess import Popen, SubprocessError, PIPE, TimeoutExpired
 import glob
-from tempfile import gettempprefix, NamedTemporaryFile
-from ..localapp import cleanup_tempfile
+from tempfile import TemporaryDirectory
 from ..application import Application, AppState, AppStateError, \
                           requires_state
 from ...sequence.seqtypes import NucleotideSequence
@@ -54,13 +54,11 @@ def __init__(self, uid, output_path_prefix=None,
         self._prefetch_path = prefetch_path
         self._fasterq_dump_path = fasterq_dump_path
         self._uid = uid
+        self._sra_dir = TemporaryDirectory(suffix="_sra")
         if output_path_prefix is None:
-            self._prefix = gettempprefix()
+            self._prefix = join(self._sra_dir.name, self._uid)
         else:
             self._prefix = output_path_prefix
-        self._sra_file = NamedTemporaryFile(
-            "w", suffix=".sra", delete=False
-        )
         self._prefetch_process = None
         self._fasterq_dump_process = None
 
@@ -94,10 +92,14 @@ def join(self, timeout=None):
 
 
     def run(self):
+        # Prefetch into a temp directory with file name equaling UID
+        # This ensures that the ID in the header is not the temp prefix
+        sra_file_name = join(self._sra_dir.name, self._uid)
         command = (
-            f"{self._prefetch_path} -q -o {self._sra_file.name} {self._uid}; "
+            f"{self._prefetch_path} -q -O {self._sra_dir.name} "
+            f"{self.get_prefetch_options()} {self._uid}; "
             f"{self._fasterq_dump_path} -q -o {self._prefix}.fastq "
-            f"{self.get_fastq_dump_options()} {self._sra_file.name}"
+            f"{self.get_fastq_dump_options()} {sra_file_name}"
         )
         self._process = Popen(
             command, stdout=PIPE, stderr=PIPE, shell=True, encoding="UTF-8"
@@ -120,8 +122,8 @@ def evaluate(self):
         if exit_code != 0:
             err_msg = self._stderr.replace("\n", " ")
             raise SubprocessError(
-                f"'{self._bin_path}' returned with exit code {exit_code}: "
-                f"{err_msg}"
+                f"'prefetch' or 'fasterq-dump' returned with exit code "
+                f"{exit_code}: {err_msg}"
             )
 
         self._file_names = (
@@ -130,6 +132,8 @@ def evaluate(self):
             # For entries with multiple reads per spot
             glob.glob(self._prefix + "_*.fastq")
         )
+        print(self._prefix)
+        print(self._file_names)
         # Only load FASTQ files into memory when needed
         self._fastq_files = None
 
@@ -142,9 +146,24 @@ def wait_interval(self):
     def clean_up(self):
         if self.get_app_state() == AppState.CANCELLED:
             self._process.kill()
-        cleanup_tempfile(self._sra_file)
+        # Directory with temp files does not need to be deleted,
+        # as temp dir is automatically deleted upon object destruction
 
 
+    @requires_state(AppState.CREATED)
+    def get_prefetch_options(self):
+        """
+        Get additional options for the `prefetch` call.
+
+        PROTECTED: Override when inheriting.
+
+        Returns
+        -------
+        options: str
+            The additional options.
+        """
+        return ""
+
     @requires_state(AppState.CREATED)
     def get_fastq_dump_options(self):
         """
@@ -155,7 +174,7 @@ def get_fastq_dump_options(self):
         Returns
         -------
         options: str
-            The additional options
+            The additional options.
         """
         return ""
 
@@ -359,6 +378,14 @@ def __init__(self, uid, output_path_prefix=None, prefetch_path="prefetch",
         self._fasta_files = None
 
 
+    @requires_state(AppState.CREATED)
+    def get_prefetch_options(self):
+        return
+        # TODO: Use '--eliminate-quals'
+        # when https://github.com/ncbi/sra-tools/issues/883 is resolved
+        # return "--eliminate-quals"
+
+
     @requires_state(AppState.CREATED)
     def get_fastq_dump_options(self):
         return "--fasta"