fix selected samples as state for query list #421

cutevariant/gui/plugins/query_list/widgets.py

@@ -7,6 +7,7 @@
 from cutevariant.gui.widgets import CodeEdit
 from cutevariant.gui.settings import Config
 from cutevariant.gui import FIcon, style
+from cutevariant import LOGGER
 
 from PySide6.QtWidgets import *
 from PySide6.QtGui import *
@@ -334,11 +335,76 @@ def _run_query(self):
         query = self.model.data(index, Qt.UserRole)
         query_params = parse_one_vql(query)
 
-        self.mainwindow.set_state_data("fields",query_params.get("fields",[]))
+        # Config
+        config = Config("variables")
+
+        # additional genotype fields
+        additional_genotype_fields = ""
+        if "additional_genotype_fields" in config:
+            additional_genotype_fields=config["additional_genotype_fields"]
+        else:
+            config["additional_genotype_fields"]=""
+            config.save()
+        genotype_fields_existing=self.conn.execute(f"SELECT * FROM sample_has_variant LIMIT 1").fetchone().keys()
+        additional_genotype_fields_clear=[]
+        for additional_genotype_field in additional_genotype_fields.split(","):
+            if additional_genotype_field in genotype_fields_existing:
+                additional_genotype_fields_clear.append(additional_genotype_field)
+            else:
+                LOGGER.warning(f"field '{additional_genotype_field}' not existing")
+        if additional_genotype_fields_clear != []:
+            additional_genotype_fields=",".join(additional_genotype_fields_clear)
+
+        # genotype position
+        genotype_position = ""
+        if "genotype_position" in config:
+            genotype_position=config["genotype_position"]
+        else:
+            config["genotype_position"]=""
+            config.save()
+        if genotype_position not in ("right", "left"):
+            genotype_position="right"
+
+        # fields
+        fields = query_params.get("fields",[])
+
+        # filters
+        filters = query_params.get("filters",{})
+        if not filters:
+            root = "$or"
+            filters["$or"] = []
+        else:
+            root = list(filters.keys())[0]
+            filters[root] = [i for i in filters[root] if not list(i.keys())[0].startswith("samples")]
+
+        # samples selected
+        samples_selected=self.mainwindow.get_state_data("samples_selected")
+        if samples_selected:
+            samples_selected_fields=[]
+            for sample in samples_selected:
+                # fields
+                samples_selected_fields += [f"samples.{sample}.gt"]
+                if additional_genotype_fields:
+                    for additional_genotype_field in additional_genotype_fields.split(","):
+                        samples_selected_fields += [f"samples.{sample}.{additional_genotype_field}"]
+                # filters
+                key = f"samples.{sample}.gt"
+                condition = {key: {"$gte": 1}}
+                filters[root].append(condition)
+
+            if samples_selected_fields != []:
+                if genotype_position == "right":
+                    fields += samples_selected_fields
+                elif genotype_position == "left":
+                    fields = samples_selected_fields + fields
+
+        # set state
+        self.mainwindow.set_state_data("fields",fields)
         self.mainwindow.set_state_data("source",query_params.get("source","variants"))
-        self.mainwindow.set_state_data("filters",query_params.get("filters",[]))
+        self.mainwindow.set_state_data("filters",filters)
         self.mainwindow.set_state_data("order_by",query_params.get("order_by",[]))
 
+        # refresh
         self.mainwindow.refresh_plugins(sender=self)
 
     def _add_query(self):

cutevariant/gui/plugins/samples/widgets.py

@@ -12,6 +12,8 @@
 from cutevariant.gui import FIcon
 from cutevariant.gui.widgets import ChoiceWidget, create_widget_action, SampleDialog
 
+from cutevariant import LOGGER
+
 # from gui.style import SAMPLE_CLASSIFICATION
 
 SAMPLE_STYLE = {
@@ -245,6 +247,7 @@ def on_search(self):
         # self.model.tags_filter = list([i["name"] for i in self.tags_choice.selected_items()])
         self.model.valid_filter = list([i["data"] for i in self.valid_choice.selected_items()])
         self.model.load()
+        self.mainwindow.set_state_data("samples_current",self.model._samples)
 
     def on_edit(self):
         sample = self.model.get_sample(self.view.currentIndex().row())
@@ -255,21 +258,68 @@ def on_edit(self):
             if dialog.exec():
                 # TODO : update only if  necessary
                 self.model.load()
+                self.mainwindow.set_state_data("samples_current",self.model._samples)
 
     def on_run(self):
 
         fields = self.mainwindow.get_state_data("fields")
         fields = [f for f in fields if not f.startswith("samples")]
 
+        # Config
+        config = Config("variables")
+
+        # additional genotype fields
+        additional_genotype_fields = ""
+        if "additional_genotype_fields" in config:
+            additional_genotype_fields=config["additional_genotype_fields"]
+        else:
+            config["additional_genotype_fields"]=""
+            config.save()
+        genotype_fields_existing=self.model.conn.execute(f"SELECT * FROM sample_has_variant LIMIT 1").fetchone().keys()
+        additional_genotype_fields_clear=[]
+        for additional_genotype_field in additional_genotype_fields.split(","):
+            if additional_genotype_field in genotype_fields_existing:
+                additional_genotype_fields_clear.append(additional_genotype_field)
+            else:
+                LOGGER.warning(f"field '{additional_genotype_field}' not existing")
+        if additional_genotype_fields_clear != []:
+            additional_genotype_fields=",".join(additional_genotype_fields_clear)
+
+        # genotype position
+        genotype_position = ""
+        if "genotype_position" in config:
+            genotype_position=config["genotype_position"]
+        else:
+            config["genotype_position"]=""
+            config.save()
+        if genotype_position not in ("right", "left"):
+            genotype_position="right"
+
         # hugly : Create genotype fields
         indexes = self.view.selectionModel().selectedRows()
+        samples_selected=[]
         if indexes:
-            sample_name = indexes[0].siblingAtColumn(0).data()
-
-            fields += [f"samples.{sample_name}.gt"]
-
+            samples_selected_fields=[]
+            for i in indexes:
+                sample_name = i.siblingAtColumn(0).data()
+                samples_selected.append(sample_name)
+                samples_selected_fields += [f"samples.{sample_name}.gt"]
+                if additional_genotype_fields:
+                    for additional_genotype_field in additional_genotype_fields.split(","):
+                        samples_selected_fields += [f"samples.{sample_name}.{additional_genotype_field}"]
+
+            if samples_selected_fields != []:
+                if genotype_position == "right":
+                    fields += samples_selected_fields
+                elif genotype_position == "left":
+                    fields = samples_selected_fields + fields
+
+        # set state
+        self.mainwindow.set_state_data("samples_selected", samples_selected)
         self.mainwindow.set_state_data("fields", fields)
         self.mainwindow.set_state_data("filters", self._create_filters())
+
+        # refresh
         self.mainwindow.refresh_plugins(sender=self)
 
     def on_create_source(self):
@@ -309,6 +359,7 @@ def on_open_project(self, conn: sqlite3.Connection):
         self.model.clear()
         self.model.conn = conn
         self.model.load()
+        self.mainwindow.set_state_data("samples_current",self.model._samples)
 
     def on_refresh(self):
         """This method is called from mainwindow.refresh_plugins()

cutevariant/gui/settings.py

@@ -402,12 +402,52 @@ def __init__(self):
         variant_name_pattern_label.setTextFormat(Qt.RichText)
 
         self.gene_field_edit = QLineEdit()
+        gene_field_label = QLabel(
+            """
+            (Examples: 'ann.gene', 'gnomen')
+            """
+        )
+        gene_field_label.setTextFormat(Qt.RichText)
+
         self.transcript_field_edit = QLineEdit()
+        transcript_field_label = QLabel(
+            """
+            (Examples: 'ann.transcript', 'tnomen')
+            """
+        )
+        transcript_field_label.setTextFormat(Qt.RichText)
+
+        self.additional_genotype_fields_edit = QLineEdit()
+        additional_genotype_fields_label = QLabel(
+            """
+            (Examples: 'dp', 'vaf', 'dp,vaf')
+            """
+        )
+        additional_genotype_fields_label.setTextFormat(Qt.RichText)
+
+        self.genotype_position_combobox = QComboBox()
+        self.genotype_position_choice=["right","left"]
+        self.genotype_position_combobox.addItems(self.genotype_position_choice)
+        genotype_position_label = QLabel(
+            """
+            ('right'/'left': genotype fields in right/left of variant table)
+            """
+        )
+        genotype_position_label.setTextFormat(Qt.RichText)
+
+
+        self.additional_genotype_fieldss = QLineEdit()
         mainLayout = QFormLayout()
         mainLayout.addRow(self.tr("Variant name pattern:"), self.variant_name_pattern_edit)
         mainLayout.addWidget(variant_name_pattern_label)
         mainLayout.addRow(self.tr("Gene field:"), self.gene_field_edit)
-        mainLayout.addRow(self.tr("Transcript field:"), self.transcript_field_edit)
+        mainLayout.addWidget(gene_field_label)
+        mainLayout.addRow(self.tr("Transcript fields:"), self.transcript_field_edit)
+        mainLayout.addWidget(transcript_field_label)
+        mainLayout.addRow(self.tr("Additional genotype fields:"), self.additional_genotype_fields_edit)
+        mainLayout.addWidget(additional_genotype_fields_label)
+        mainLayout.addRow(self.tr("Genotype position:"), self.genotype_position_combobox)
+        mainLayout.addWidget(genotype_position_label)
 
         self.setLayout(mainLayout)
 
@@ -421,9 +461,13 @@ def save(self):
         variant_name_pattern = self.variant_name_pattern_edit.text()
         gene_field = self.gene_field_edit.text()
         transcript_field = self.transcript_field_edit.text()
+        additional_genotype_fields = self.additional_genotype_fields_edit.text()
+        genotype_position = self.genotype_position_combobox.currentText()
         config["variant_name_pattern"] = variant_name_pattern
         config["gene_field"] = gene_field
         config["transcript_field"] = transcript_field
+        config["additional_genotype_fields"] = additional_genotype_fields
+        config["genotype_position"] = genotype_position
         config.save()
 
         # Clear pixmap cache
@@ -442,9 +486,14 @@ def load(self):
         variant_name_pattern = config.get("variant_name_pattern", "{chr}:{pos} - {ref}>{alt}")
         gene_field = config.get("gene_field", "ann.gene")
         transcript_field = config.get("transcript_field", "ann.transcript")
+        additional_genotype_fields = config.get("additional_genotype_fields", "")
+        genotype_position = config.get("genotype_position", "right")
         self.variant_name_pattern_edit.setText(variant_name_pattern)
         self.gene_field_edit.setText(gene_field)
         self.transcript_field_edit.setText(transcript_field)
+        self.additional_genotype_fields_edit.setText(additional_genotype_fields)
+        self.genotype_position_combobox.setCurrentIndex(self.genotype_position_choice.index(genotype_position))
+        # self.styles_combobox.setCurrentIndex(available_styles.index(style_name))
 
 class SettingsDialog(QDialog):
     """Main widget for settings window