init fix for querying for datasets

1 files changed, 17 insertions, 15 deletions

diff --git a/wqflask/base/trait.py b/wqflask/base/trait.py
index 1514bca2..dcd81085 100644
--- a/wqflask/base/trait.py
+++ b/wqflask/base/trait.py
@@ -24,13 +24,23 @@ def create_trait(**kw):
 
     assert bool(kw.get('name')), "Needs trait name"
 
+
     if bool(kw.get('dataset')):
         dataset = kw.get('dataset')
+
+
     else:
         if kw.get('dataset_name') != "Temp":
+
+
             dataset = create_dataset(kw.get('dataset_name'))
         else:
-            dataset = create_dataset("Temp", group_name=kw.get('group_name'))
+
+            dataset = create_dataset(
+                    dataset_name="Temp",
+                    dataset_type="Temp",
+                    group_name= kw.get('name').split("_")[2])
+
 
     if dataset.type == 'Publish':
         permissions = check_resource_availability(
@@ -38,8 +48,10 @@ def create_trait(**kw):
     else:
         permissions = check_resource_availability(dataset)
 
+
     if permissions['data'] != "no-access":
-        the_trait = GeneralTrait(**kw)
+        
+        the_trait = GeneralTrait(**dict(kw,dataset=dataset))
         if the_trait.dataset.type != "Temp":
             the_trait = retrieve_trait_info(
                 the_trait,
@@ -59,21 +71,11 @@ class GeneralTrait:
 
     def __init__(self, get_qtl_info=False, get_sample_info=True, **kw):
         # xor assertion
-        assert bool(kw.get('dataset')) != bool(
-            kw.get('dataset_name')), "Needs dataset ob. or name"
+        assert kw.get("dataset"), "Dataset obj is needed as a kwarg"
+
         # Trait ID, ProbeSet ID, Published ID, etc.
         self.name = kw.get('name')
-        if kw.get('dataset_name'):
-            if kw.get('dataset_name') == "Temp":
-                temp_group = self.name.split("_")[2]
-                self.dataset = create_dataset(
-                    dataset_name="Temp",
-                    dataset_type="Temp",
-                    group_name=temp_group)
-            else:
-                self.dataset = create_dataset(kw.get('dataset_name'))
-        else:
-            self.dataset = kw.get('dataset')
+        self.dataset = kw.get("dataset")
         self.cellid = kw.get('cellid')
         self.identification = kw.get('identification', 'un-named trait')
         self.haveinfo = kw.get('haveinfo', False)