Merge github.com:genenetwork/genenetwork2 into acenteno

1 files changed, 39 insertions, 43 deletions

diff --git a/wqflask/utility/helper_functions.py b/wqflask/utility/helper_functions.py
index e7c04fef..27dd0729 100644
--- a/wqflask/utility/helper_functions.py
+++ b/wqflask/utility/helper_functions.py
@@ -1,72 +1,68 @@
-from __future__ import absolute_import, print_function, division
-
-from base.trait import GeneralTrait
 from base import data_set
+from base.trait import create_trait
 from base.species import TheSpecies
 
 from utility import hmac
 
-from flask import Flask, g
+from flask import g
 
 import logging
-logger = logging.getLogger(__name__ )
+logger = logging.getLogger(__name__)
 
 
 def get_species_dataset_trait(self, start_vars):
-    #assert type(read_genotype) == type(bool()), "Expecting boolean value for read_genotype"
-    if "temp_trait" in start_vars.keys():
-      if start_vars['temp_trait'] == "True":
-        self.dataset = data_set.create_dataset(dataset_name = "Temp", dataset_type = "Temp", group_name = start_vars['group'])
-      else:
-        self.dataset = data_set.create_dataset(start_vars['dataset'])
+    if "temp_trait" in list(start_vars.keys()):
+        if start_vars['temp_trait'] == "True":
+            self.dataset = data_set.create_dataset(
+                dataset_name="Temp",
+                dataset_type="Temp",
+                group_name=start_vars['group'])
+        else:
+            self.dataset = data_set.create_dataset(start_vars['dataset'])
     else:
-      self.dataset = data_set.create_dataset(start_vars['dataset'])
-    logger.debug("After creating dataset")
+        self.dataset = data_set.create_dataset(start_vars['dataset'])
     self.species = TheSpecies(dataset=self.dataset)
-    logger.debug("After creating species")
-    self.this_trait = GeneralTrait(dataset=self.dataset,
+    self.this_trait = create_trait(dataset=self.dataset,
                                    name=start_vars['trait_id'],
                                    cellid=None,
                                    get_qtl_info=True)
-    logger.debug("After creating trait")
-
-    #if read_genotype:
-    #self.dataset.group.read_genotype_file()
-    #self.genotype = self.dataset.group.genotype
-
 
 def get_trait_db_obs(self, trait_db_list):
-    if isinstance(trait_db_list, basestring):
+    if isinstance(trait_db_list, str):
         trait_db_list = trait_db_list.split(",")
 
     self.trait_list = []
     for trait in trait_db_list:
         data, _separator, hmac_string = trait.rpartition(':')
         data = data.strip()
-        assert hmac_string==hmac.hmac_creation(data), "Data tampering?"
-        trait_name, dataset_name = data.split(":")
+        assert hmac_string == hmac.hmac_creation(data), "Data tampering?"
+        trait_name, dataset_name = data.split(":")[:2]
         if dataset_name == "Temp":
-            dataset_ob = data_set.create_dataset(dataset_name=dataset_name, dataset_type="Temp", group_name=trait_name.split("_")[2])
+            dataset_ob = data_set.create_dataset(
+                dataset_name=dataset_name, dataset_type="Temp",
+                group_name=trait_name.split("_")[2])
         else:
             dataset_ob = data_set.create_dataset(dataset_name)
-        trait_ob = GeneralTrait(dataset=dataset_ob,
-                               name=trait_name,
-                               cellid=None)
-        self.trait_list.append((trait_ob, dataset_ob))
-
-def get_species_groups():
+        trait_ob = create_trait(dataset=dataset_ob,
+                                name=trait_name,
+                                cellid=None)
+        if trait_ob:
+            self.trait_list.append((trait_ob, dataset_ob))
 
-    species_query = "SELECT SpeciesId, MenuName FROM Species"
-    species_ids_and_names = g.db.execute(species_query).fetchall()
 
-    species_and_groups = []
-    for species_id, species_name in species_ids_and_names:
-        this_species_groups = {}
-        this_species_groups['species'] = species_name
-        groups_query = "SELECT InbredSetName FROM InbredSet WHERE SpeciesId = %s" % (species_id)
-        groups = [group[0] for group in g.db.execute(groups_query).fetchall()]
-
-        this_species_groups['groups'] = groups
-        species_and_groups.append(this_species_groups)
+def get_species_groups():
+    """Group each species into a group"""
+    _menu = {}
 
-    return species_and_groups
+    for species, group_name in g.db.execute(
+            "SELECT s.MenuName, i.InbredSetName FROM InbredSet i "
+            "INNER JOIN Species s ON s.SpeciesId = i.SpeciesId "
+            "ORDER BY i.SpeciesId ASC, i.Name ASC").fetchall():
+        if species in _menu:
+            if _menu.get(species):
+                _menu = _menu[species].append(group_name)
+            else:
+                _menu[species] = [group_name]
+    return [{"species": key,
+             "groups": value} for key, value in
+            list(_menu.items())]