Add collection export function/endpoint

2 files changed, 40 insertions, 7 deletions

diff --git a/wqflask/wqflask/export_traits.py b/wqflask/wqflask/export_traits.py
index 5de6eb43..105d81eb 100644
--- a/wqflask/wqflask/export_traits.py
+++ b/wqflask/wqflask/export_traits.py
@@ -1,23 +1,48 @@
 import csv
-import xlsxwriter
-import io
 import datetime
+import io
 import itertools
+import re
+import xlsxwriter
 
+from pprint import pformat as pf
 from zipfile import ZipFile, ZIP_DEFLATED
 
 import simplejson as json
 
-from base.trait import create_trait, retrieve_trait_info
+from gn3.computations.gemma import generate_hash_of_string
 
-from pprint import pformat as pf
+from base.trait import create_trait, retrieve_trait_info
 
 from utility.logger import getLogger
 logger = getLogger(__name__)
 
+def export_traits(targs, export_type):
+    if export_type == "collection":
+        return export_collection(targs)
+    else:
+        return export_traitlist(targs)
+
+def export_collection(targs):
+    table_data = json.loads(targs['export_data'])
+    table_rows = table_data['rows']
+
+    buff = io.StringIO()
+    writer = csv.writer(buff)
+    for trait in table_rows:
+        writer.writerow(trait.split(":"))
+
+    csv_data = buff.getvalue()
+    buff.close()
+
+    if 'collection_name_export' in targs:
+        file_name = re.sub('\s+', '_', targs['collection_name_export']) # replace whitespace with underscore
+    else:
+        file_name = generate_hash_of_string("".join(table_rows))
 
-def export_traits_csv(targs):
+    return [file_name, csv_data]
 
+def export_traitlist(targs):
     table_data = json.loads(targs['export_data'])
     table_rows = table_data['rows']
 
diff --git a/wqflask/wqflask/views.py b/wqflask/wqflask/views.py
index 729148ac..c38d46ca 100644
--- a/wqflask/wqflask/views.py
+++ b/wqflask/wqflask/views.py
@@ -73,7 +73,7 @@ from wqflask.ctl.gn3_ctl_analysis import run_ctl
 from wqflask.wgcna.gn3_wgcna import run_wgcna
 from wqflask.snp_browser import snp_browser
 from wqflask.search_results import SearchResultPage
-from wqflask.export_traits import export_traits_csv
+from wqflask.export_traits import export_traits
 from wqflask.gsearch import GSearch
 from wqflask.update_search_results import GSearch as UpdateGSearch
 from wqflask.docs import Docs, update_text
@@ -469,7 +469,7 @@ def export_traits_csv():
     logger.info("In export_traits_csv")
     logger.info("request.form:", request.form)
     logger.info(request.url)
-    file_list = export_traits_csv(request.form)
+    file_list = export_traits(request.form, "metadata")
 
     if len(file_list) > 1:
         now = datetime.datetime.now()
@@ -489,6 +489,14 @@ def export_traits_csv():
                         headers={"Content-Disposition": "attachment;filename=" + file_list[0][0]})
 
 
+@app.route('/export_collection', methods=('POST',))
+def export_collection_csv():
+    """CSV file consisting of trait list so collections can be exported/shared"""
+    out_file = export_traits(request.form, "collection")
+    return Response(out_file[1],
+                    mimetype='text/csv',
+                    headers={"Content-Disposition": "attachment;filename=" + out_file[0] + ".csv"})
+
 @app.route('/export_perm_data', methods=('POST',))
 def export_perm_data():
     """CSV file consisting of the permutation data for the mapping results"""