Repository: EBI-Metagenomics/emg-toolkit (language: python)

Issue:
Bulk Download metadata storage.
With this changeset the bulk download app will store
some of the downloaded files metadata in a tsv in the
study folder.
This will make it easier to interpret the files the tool downloaded,
for example further filter down the files (i.e. find the contigs in
assemblies)


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index 83a9a0e..116b88a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -24,7 +24,7 @@ before_script:
 
 script:
   - cd $TRAVIS_BUILD_DIR/tests
-  - bash tests.sh
+  - bash run_tests.sh
 
 notifications:
   slack:
diff --git a/README.md b/README.md
index 2881dfd..b139fd7 100644
--- a/README.md
+++ b/README.md
@@ -7,9 +7,13 @@ metadata for a given study or sequence to a single csv file.
 
 Install metagenomics toolkit
 ============================
+Through pip
 
     pip install -U mg-toolkit
 
+Or using conda
+
+    conda install -c bioconda mg-toolkit
 
 Usage
 =====
@@ -43,7 +47,7 @@ Download metadata:
 
 Search non-redundant protein database using HMMER and fetch metadata:
 
-    $ mg-toolkit sequence_search -seq test.fasta -db full evalue -incE 0.02
+    $ mg-toolkit sequence_search -seq test.fasta -out test.csv -db full evalue -incE 0.02
 
     Databases:
     - full - Full length sequences (default)
@@ -115,6 +119,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
     <td align="center"><a href="https://github.com/olatarkowska"><img src="https://avatars3.githubusercontent.com/u/1065155?v=4" width="100px;" alt=""/><br /><sub><b>Ola Tarkowska</b></sub></a><br /><a href="https://github.com/EBI-Metagenomics/emg-toolkit/commits?author=olatarkowska" title="Code">💻</a><a href="https://github.com/EBI-Metagenomics/EMG-docs/commits/master?author=olatarkowska">📖</a></td>
     <td align="center"><a href="https://github.com/mscheremetjew"><img src="https://avatars3.githubusercontent.com/u/1681284?v=4" width="100px;" alt=""/><br /><sub><b>Maxim Scheremetjew</b></sub></a><br /><a href="https://github.com/EBI-Metagenomics/emg-toolkit/commits?author=mscheremetjew" title="Code">💻</a><a href="https://github.com/EBI-Metagenomics/EMG-docs/commits/master?author=mscheremetjew">📖</a></td>
     <td align="center"><a href="https://github.com/mberacochea"><img src="https://avatars3.githubusercontent.com/u/1123897?v=4" width="100px;" alt=""/><br /><sub><b>Martin Beracochea</b></sub></a><br /><a href="https://github.com/EBI-Metagenomics/emg-toolkit/commits?author=mberacochea" title="Code">💻</a></td>
+    <td align="center"><a href="https://github.com/emilhaegglund"><img src="https://avatars.githubusercontent.com/u/15671004?s=400&u=1715b96c4879b924eab88f48461ffd9f40e9d37a&v=4" width="100px;" alt=""/><br /><sub><b>Emil Hägglund</b></sub></a><br /><a href="https://github.com/EBI-Metagenomics/emg-toolkit/commits?author=emilhaegglund" title="Code">💻</a></td>
   </tr>
 </table>
 
diff --git a/mg_toolkit/__main__.py b/mg_toolkit/__main__.py
index d068c0f..bca8fea 100644
--- a/mg_toolkit/__main__.py
+++ b/mg_toolkit/__main__.py
@@ -76,6 +76,12 @@ def main():
         nargs="+",
         help="Provide path to fasta file.",
     )
+    sequence_search_parser.add_argument(
+        "-out",
+        "--output",
+        required=False,
+        help="Output csv results file (default: <query_id>_sequence_search.csv",
+    )
     sequence_search_parser.add_argument(
         "-db",
         "--database",
diff --git a/mg_toolkit/bulk_download.py b/mg_toolkit/bulk_download.py
index 1bde8fe..1c22548 100644
--- a/mg_toolkit/bulk_download.py
+++ b/mg_toolkit/bulk_download.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-# Copyright 2020 EMBL - European Bioinformatics Institute
+# Copyright 2021 EMBL - European Bioinformatics Institute
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 import logging
 import os
 import platform
+import csv
 from pathlib import Path
 
 from requests.adapters import HTTPAdapter
@@ -292,6 +293,9 @@ def _process_download_page(self, analysis, download_response):
                     project_id=self.project_id,
                     dest_dir=self.output_path,
                 )
+            # store the metadata for the analysis
+            self.store_metadata(analysis, response_json)
+
             next_page_url = response_json.get("links", {}).get("next")
             if next_page_url:
                 next_page_respose = self.http.get(
@@ -319,3 +323,65 @@ def process_page(self, response_data, progress_bar):
             progress_bar.update(1)
 
         return processed_counter
+
+    def store_metadata(self, analysis, response_json):
+        """
+        Store the API response json in a tsv file called <analysis>_metadata.tsv
+        This file can be used to make it easier to interpret the downloaded files.
+        """
+        metadata_file_name = "{}_metadata.tsv".format(self.project_id)
+        output_file = os.path.join(
+            self.output_path, self.project_id, metadata_file_name
+        )
+        mode = "a" if os.path.exists(output_file) else "w"
+        experyment_type = analysis.get("attributes").get("experiment-type")
+
+        col_names = [
+            "analysis_id",
+            "name",
+            "group_type",
+            "description",
+            "download_url",
+            "pipeline_version",
+            "experiment_type",
+            # TODO: enable when released for the pipeline
+            # "checksum",
+            # "checksum_algorithm",
+        ]
+
+        with open(output_file, mode) as metada_fd:
+            writer = csv.writer(metada_fd, delimiter="\t")
+            if mode == "w":
+                writer.writerow(col_names)
+
+            for entry in response_json.get("data", []):
+                download_attr = entry.get("attributes")
+                alias = download_attr.get("alias")
+                group_type = download_attr.get("group-type")
+                desc_label = download_attr.get("description").get("label")
+
+                download_url = entry.get("links").get("self")
+
+                pipeline_version = (
+                    entry.get("relationships").get("pipeline").get("data").get("id")
+                )
+
+                # TODO: enable when released for the pipeline
+                # checksum = download_attr.get("file-checksum").get("checksum")
+                # checksum_algorithm = download_attr.get("file-checksum").get(
+                #     "checksum-algorithm"
+                # )
+
+                writer.writerow(
+                    [
+                        analysis["id"],
+                        alias,
+                        group_type,
+                        desc_label,
+                        download_url,
+                        pipeline_version,
+                        experyment_type,
+                        # checksum,
+                        # checksum_algorithm,
+                    ]
+                )
diff --git a/mg_toolkit/search.py b/mg_toolkit/search.py
index b9f9c16..7f9b842 100644
--- a/mg_toolkit/search.py
+++ b/mg_toolkit/search.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-# Copyright 2020 EMBL - European Bioinformatics Institute
+# Copyright 2021 EMBL - European Bioinformatics Institute
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -29,17 +29,38 @@
 logger = logging.getLogger(__name__)
 
 
-def sequence_search(args):
+def parse_fasta_file(file_path):
+    """
+    Parse fasta file
+    """
+    sequences = {}
+    query_id = None
+    with open(file_path, "r") as f:
+        for line in f:
+            if line.startswith(">"):
+                query_id = line[1:].strip()
+                sequences[query_id] = ""
+            else:
+                sequences[query_id] += line.strip()
+    return sequences
+
 
+def sequence_search(args):
     """
     Process given fasta file
     """
     args = vars(args)
+    out_df = DataFrame()
     for s in args.pop("sequence"):
-        with open(s) as f:
-            sequence = f.read()
+        sequences = parse_fasta_file(s)
+        for sequence_data in sequences.items():
+            query_id = sequence_data[0]
+            sequence = sequence_data[1]
+            print("Proccessing: {}".format(query_id))
+            # Search MgnifyDB
             seq = SequenceSearch(
                 sequence,
+                query_id,
                 database=args.pop("database", "full"),
                 seq_evalue_threshold=args.pop("seq_evalue_threshold", None),
                 hit_evalue_threshold=args.pop("hit_evalue_threshold", None),
@@ -58,14 +79,29 @@ def sequence_search(args):
                     "report_hit_bitscore_threshold", None
                 ),
             )
-            results = seq.analyse_sequence()
-            job_uuid = results["results"]["uuid"]
-            logger.debug("Job %s" % job_uuid)
-            seq.save_to_csv(seq.fetch_results(results), filename=job_uuid)
+            response = seq.analyse_sequence()
+            if not response:
+                logger.warning("No results to report for %s" % query_id)
+                return
+            # Only process results when the request returned data
+            results = response.get("results")
+            if results:
+                job_uuid = results["uuid"]
+                logger.debug("Job %s" % job_uuid)
+                out_df = out_df.append(
+                    seq.results_to_df(seq.fetch_results(results), job_uuid)
+                )
+            else:
+                logger.warning("No results to report for %s" % query_id)
 
+    output_file = job_uuid + "_sequence_search.csv"
+    if args["output"]:
+        output_file = args["output"]
+
+    out_df.to_csv(output_file, index=False)
 
-class SequenceSearch(object):
 
+class SequenceSearch(object):
     """
     Helper tool allowing to search non-redundant protein database using HMMER
     and fetch environmental metadata.
@@ -74,8 +110,9 @@ class SequenceSearch(object):
     sequence = None
     database = "full"
 
-    def __init__(self, sequence, database="full", *args, **kwargs):
+    def __init__(self, sequence, query_id, database="full", *args, **kwargs):
         self.sequence = sequence
+        self.query_id = query_id
         self.database = database
         self.seq_evalue_threshold = kwargs.pop("seq_evalue_threshold", None)
         self.hit_evalue_threshold = kwargs.pop("hit_evalue_threshold", None)
@@ -121,7 +158,12 @@ def analyse_sequence(self):
             "Content-Type": "application/x-www-form-urlencoded",
         }
         logger.debug("POST: %r" % data)
-        return requests.post(MG_SEQ_URL, data=data, headers=headers).json()
+        request_data = requests.post(MG_SEQ_URL, data=data, headers=headers)
+        # Check if data was returned
+        if request_data:
+            return request_data.json()
+        else:
+            return False
 
     def make_request(self, accession):
         if accession is None:
@@ -184,7 +226,6 @@ def prepare_rows(self, hit):
         return {
             "kg": hit.get("kg", ""),
             "taxid": hit.get("taxid", ""),
-            "name": hit.get("name", ""),
             "desc": hit.get("desc", ""),
             "pvalue": hit.get("pvalue", ""),
             "species": hit.get("species", ""),
@@ -195,8 +236,12 @@ def prepare_rows(self, hit):
         }
 
     def fetch_results(self, results):
+        """
+        Complete the HMMER hits with MGnify metadata from the API
+        TODO: This could be parallelized for better performance
+        """
         csv_rows = dict()
-        for hit in results["results"]["hits"]:
+        for hit in results.get("hits", []):
             _row = self.prepare_rows(hit)
             mgnify = hit.get("mgnify", [])
             for res in mgnify.get("samples") or []:
@@ -219,9 +264,34 @@ def fetch_results(self, results):
                 csv_rows[uuid].update(_meta)
         return csv_rows
 
-    def save_to_csv(self, csv_rows, filename):
+    def results_to_df(self, csv_rows, uuid):
+        """
+        Convert search results to dataframe
+        """
+
         df = DataFrame(csv_rows).T
-        df.index.name = "name"
-        df = df.reindex(columns=sorted(df.columns))
-        filename = "{}_sequence_search.csv".format(filename)
-        df.to_csv(filename)
+        df.reset_index(inplace=True)
+        df.rename(columns={"index": "name"}, inplace=True)
+        # Split index to subject_id and accession
+        df[["subject_id", "accession"]] = df["name"].str.split(" ", 1, expand=True)
+
+        # Put query_id, subject_id, and accession as first three columns
+        subject_id = df.subject_id
+        accession = df.accession
+        df.drop(["name", "subject_id", "accession"], axis=1, inplace=True)
+        df.insert(0, "query_id", self.query_id)
+        df.insert(1, "subject_id", subject_id)
+        df.insert(1, "accession", accession)
+
+        def _clean_column(col):
+            return (
+                col.replace(")", "")
+                .replace("(", "")
+                .replace("/", "_")
+                .replace(",", "_")
+            )
+
+        # Clean columns from (),\
+        df.rename(_clean_column, axis="columns", inplace=True)
+
+        return df


Test patch:
diff --git a/tests/bulk_download/ERR169332.5_8S_rRNA.RF00002.fa.md5 b/tests/integration/bulk_download/ERR169332.5_8S_rRNA.RF00002.fa.md5
similarity index 100%
rename from tests/bulk_download/ERR169332.5_8S_rRNA.RF00002.fa.md5
rename to tests/integration/bulk_download/ERR169332.5_8S_rRNA.RF00002.fa.md5
diff --git a/tests/original_metadata/ERP001178.csv.md5 b/tests/integration/original_metadata/ERP001178.csv.md5
similarity index 100%
rename from tests/original_metadata/ERP001178.csv.md5
rename to tests/integration/original_metadata/ERP001178.csv.md5
diff --git a/tests/integration/sequence_search/sequence_search.csv.md5 b/tests/integration/sequence_search/sequence_search.csv.md5
new file mode 100644
index 0000000..064333f
--- /dev/null
+++ b/tests/integration/sequence_search/sequence_search.csv.md5
@@ -0,0 +1,1 @@
+881bcbe7f2c0529f9cbef1fb46ce77e1  sequence_search.csv
\ No newline at end of file
diff --git a/tests/sequence_search/test.fasta b/tests/integration/sequence_search/test.fasta
similarity index 100%
rename from tests/sequence_search/test.fasta
rename to tests/integration/sequence_search/test.fasta
diff --git a/tests/tests.sh b/tests/run_tests.sh
similarity index 82%
rename from tests/tests.sh
rename to tests/run_tests.sh
index adfbb88..724f284 100755
--- a/tests/tests.sh
+++ b/tests/run_tests.sh
@@ -2,8 +2,18 @@
 
 set -e
 
+echo "Unit tests"
+
+python -m unittest discover unit -v
+
+echo "------------------------"
+
+echo "Integration"
+
 echo "Testing original_metadata"
 
+cd integration
+
 set -x
 mg-toolkit -d original_metadata -a ERP001178
 set +x
@@ -28,4 +38,4 @@ set +x
 
 mv -f *_sequence_search.csv sequence_search.csv
 md5sum -c sequence_search/sequence_search.csv.md5
-rm *.csv
\ No newline at end of file
+rm *.csv
diff --git a/tests/sequence_search/sequence_search.csv.md5 b/tests/sequence_search/sequence_search.csv.md5
deleted file mode 100644
index a6a2620..0000000
--- a/tests/sequence_search/sequence_search.csv.md5
+++ /dev/null
@@ -1,1 +0,0 @@
-485903588290f109b7f9d14aef6997cb  sequence_search.csv
\ No newline at end of file
diff --git a/tests/unit/fixtures/test.fasta b/tests/unit/fixtures/test.fasta
new file mode 100644
index 0000000..57d8a4d
--- /dev/null
+++ b/tests/unit/fixtures/test.fasta
@@ -0,0 +1,9 @@
+>fasta_one example   
+MSTHPIRVFSEIGKLKKVMLHRPGKELENLQPDY
+LERLLFDD
+>fasta_two example
+EEYLEEANIRGRETKKAIRELLHGIKDNQELVEKT
+>fasta_three example 
+AVSLNHMYADTRNRETLYGKYI
+FKYHPVYGGNVELVYNREEDTRI
+EGGDELVLSKDVLAVGISQRTDAA
diff --git a/tests/unit/test_search.py b/tests/unit/test_search.py
new file mode 100644
index 0000000..f391b28
--- /dev/null
+++ b/tests/unit/test_search.py
@@ -0,0 +1,25 @@
+#!/bin/env python3
+
+import os
+import unittest
+
+from mg_toolkit.search import parse_fasta_file
+
+
+class FastaParseTests(unittest.TestCase):
+    def _build_path(self, folder):
+        return os.path.abspath("/" + os.path.dirname(__file__) + folder)
+
+    def test_parse_fasta_file(self):
+        """Test fasta parsing method"""
+        test_fasta = self._build_path("/fixtures/test.fasta")
+
+        expected = {
+            "fasta_one example": "MSTHPIRVFSEIGKLKKVMLHRPGKELENLQPDYLERLLFDD",
+            "fasta_two example": "EEYLEEANIRGRETKKAIRELLHGIKDNQELVEKT",
+            "fasta_three example": "AVSLNHMYADTRNRETLYGKYIFKYHPVYGGNVELVYNREEDTRIEGGDELVLSKDVLAVGISQRTDAA",
+        }
+
+        fn_result = parse_fasta_file(test_fasta)
+
+        self.assertDictEqual(fn_result, expected)