fix:add lit for trait list

1 files changed, 28 insertions, 11 deletions

diff --git a/wqflask/wqflask/correlation/correlation_gn3_api.py b/wqflask/wqflask/correlation/correlation_gn3_api.py
index 4949bbe2..d68bb604 100644
--- a/wqflask/wqflask/correlation/correlation_gn3_api.py
+++ b/wqflask/wqflask/correlation/correlation_gn3_api.py
@@ -83,6 +83,23 @@ def tissue_for_trait_lists(corr_results, this_dataset, this_trait):
     return corr_results
 
 
+def lit_for_trait_list(corr_results, this_dataset, this_trait):
+    (this_trait_geneid, geneid_dict, species) = do_lit_correlation(
+            this_trait, this_dataset)
+
+    trait_lists = {list(corr_results)[0]: 1 for corr_result in corr_results}
+
+    geneid_dict = {trait_name: geneid for (trait_name, geneid) in geneid_dict if
+    trait_lists.get(trait_name)}
+
+    conn, _cursor_object = database_connector()
+    correlation_results = compute_all_lit_correlation(
+                conn=conn, trait_lists=list(geneid_dict.items()),
+                species=species, gene_id=this_trait_geneid)
+
+    return correlation_results[0:corr_return_results]
+
+
 def compute_correlation(start_vars, method="pearson"):
     """compute correlation for to call gn3  api"""
     # pylint: disable-msg=too-many-locals
@@ -136,30 +153,30 @@ def compute_correlation(start_vars, method="pearson"):
         conn, _cursor_object = database_connector()
         with conn:
             correlation_results = compute_all_lit_correlation(
-                conn=conn, trait_lists=list(geneid_dict.items()),
-                species=species, gene_id=this_trait_geneid)
+                conn = conn, trait_lists = list(geneid_dict.items()),
+                species = species, gene_id = this_trait_geneid)
 
     return correlation_results[0:corr_return_results]
 
 
 def do_lit_correlation(this_trait, this_dataset):
     """function for fetching lit inputs"""
-    geneid_dict = this_dataset.retrieve_genes("GeneId")
-    species = this_dataset.group.species.lower()
-    trait_geneid = this_trait.geneid
+    geneid_dict=this_dataset.retrieve_genes("GeneId")
+    species=this_dataset.group.species.lower()
+    trait_geneid=this_trait.geneid
     return (trait_geneid, geneid_dict, species)
 
 
 def get_tissue_correlation_input(this_trait, trait_symbol_dict):
     """Gets tissue expression values for the primary trait and target tissues values"""
-    primary_trait_tissue_vals_dict = correlation_functions.get_trait_symbol_and_tissue_values(
-        symbol_list=[this_trait.symbol])
+    primary_trait_tissue_vals_dict=correlation_functions.get_trait_symbol_and_tissue_values(
+        symbol_list = [this_trait.symbol])
     if this_trait.symbol.lower() in primary_trait_tissue_vals_dict:
-        primary_trait_tissue_values = primary_trait_tissue_vals_dict[this_trait.symbol.lower(
+        primary_trait_tissue_values=primary_trait_tissue_vals_dict[this_trait.symbol.lower(
         )]
-        corr_result_tissue_vals_dict = correlation_functions.get_trait_symbol_and_tissue_values(
-            symbol_list=list(trait_symbol_dict.values()))
-        primary_tissue_data = {
+        corr_result_tissue_vals_dict=correlation_functions.get_trait_symbol_and_tissue_values(
+            symbol_list = list(trait_symbol_dict.values()))
+        primary_tissue_data={
             "this_id": this_trait.name,
             "tissue_values": primary_trait_tissue_values