From 30255ac76a3165c15a4cdd7f45e0e3e494587b9a Mon Sep 17 00:00:00 2001
From: Milo Hyben <milo.hyben@gmail.com>
Date: Tue, 28 May 2024 10:31:49 +1000
Subject: [PATCH] Draft idea how to deal with local/dev environment when
 including WORKFLOWS.

---
 main.py          | 55 ++++++------------------------------------------
 workflows_dev.py | 12 +++++++++++
 workflows_prd.py | 54 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 48 deletions(-)
 create mode 100644 workflows_dev.py
 create mode 100644 workflows_prd.py

diff --git a/main.py b/main.py
index bee06eb0a..497eb1e0b 100644
--- a/main.py
+++ b/main.py
@@ -4,61 +4,20 @@
 Entry point to run workflows.
 """
 import os
-
 import click
 import coloredlogs
 
 from cpg_utils import to_path
 from cpg_utils.config import set_config_paths
 from cpg_workflows import defaults_config_path
-from cpg_workflows.stages.aip import CreateAIPHTML, GenerateSeqrFile, ValidateMOI
-from cpg_workflows.stages.cram_qc import CramMultiQC
-from cpg_workflows.stages.exomiser import ExomiserSeqrTSV, RunExomiser
-from cpg_workflows.stages.fastqc import FastQCMultiQC
-from cpg_workflows.stages.fraser import Fraser
-from cpg_workflows.stages.gatk_sv.gatk_sv_multisample_1 import (
-    FilterBatch,
-    GenotypeBatch,
-    MergeBatchSites,
-)
-from cpg_workflows.stages.gatk_sv.gatk_sv_multisample_2 import MtToEsSv
-from cpg_workflows.stages.gatk_sv.gatk_sv_single_sample import CreateSampleBatches
-from cpg_workflows.stages.gcnv import AnnotateCohortgCNV, AnnotateDatasetCNV, MtToEsCNV
-from cpg_workflows.stages.gvcf_qc import GvcfMultiQC
-from cpg_workflows.stages.happy_validation import (
-    ValidationHappyOnVcf,
-    ValidationMtToVcf,
-    ValidationParseHappy,
-)
-from cpg_workflows.stages.large_cohort import AncestryPlots, Frequencies, LoadVqsr
-from cpg_workflows.stages.mito import MitoReport
-from cpg_workflows.stages.outrider import Outrider
-from cpg_workflows.stages.seqr_loader import AnnotateDataset, DatasetVCF, MtToEs
-from cpg_workflows.stages.stripy import Stripy
-from cpg_workflows.workflow import StageDecorator, run_workflow
+from cpg_workflows.workflow import run_workflow
 
-WORKFLOWS: dict[str, list[StageDecorator]] = {
-    'aip': [ValidateMOI, CreateAIPHTML, GenerateSeqrFile],
-    'exomiser': [RunExomiser, ExomiserSeqrTSV],
-    'pre_alignment': [FastQCMultiQC],
-    'seqr_loader': [
-        DatasetVCF,
-        AnnotateDataset,
-        MtToEs,
-        GvcfMultiQC,
-        CramMultiQC,
-        Stripy,
-        MitoReport,
-    ],
-    'validation': [ValidationMtToVcf, ValidationHappyOnVcf, ValidationParseHappy],
-    'large_cohort': [LoadVqsr, Frequencies, AncestryPlots, GvcfMultiQC, CramMultiQC],
-    'gatk_sv_singlesample': [CreateSampleBatches],
-    'gatk_sv_multisample_1': [FilterBatch, GenotypeBatch],
-    'gatk_sv_sandwich': [MergeBatchSites],  # stage to run between FilterBatch & GenotypeBatch
-    'gatk_sv_multisample_2': [MtToEsSv],
-    'rare_disease_rnaseq': [Outrider, Fraser],
-    'gcnv': [AnnotateCohortgCNV, AnnotateDatasetCNV, MtToEsCNV],
-}
+# Importing WORKFLOWS from the correct file based on the environment:
+SM_ENVIRONMENT = os.getenv('PL_ENVIRONMENT', 'production').lower()
+if SM_ENVIRONMENT == 'production':
+    from workflows_prd import WORKFLOWS
+else:
+    from workflows_dev import WORKFLOWS
 
 
 @click.command(no_args_is_help=True)
diff --git a/workflows_dev.py b/workflows_dev.py
new file mode 100644
index 000000000..587f0a237
--- /dev/null
+++ b/workflows_dev.py
@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+"""
+Dev version of workflows.
+"""
+from cpg_workflows.stages.cram_qc import CramMultiQC
+from cpg_workflows.stages.gvcf_qc import GvcfMultiQC
+from cpg_workflows.stages.large_cohort import AncestryPlots, Frequencies, LoadVqsr
+from cpg_workflows.workflow import StageDecorator
+
+WORKFLOWS: dict[str, list[StageDecorator]] = {
+    'large_cohort': [LoadVqsr, Frequencies, AncestryPlots, GvcfMultiQC, CramMultiQC],
+}
\ No newline at end of file
diff --git a/workflows_prd.py b/workflows_prd.py
new file mode 100644
index 000000000..3bf8c0f09
--- /dev/null
+++ b/workflows_prd.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+"""
+Production version of workflows.
+"""
+from cpg_workflows.stages.cram_qc import CramMultiQC
+from cpg_workflows.stages.gvcf_qc import GvcfMultiQC
+from cpg_workflows.stages.large_cohort import AncestryPlots, Frequencies, LoadVqsr
+from cpg_workflows.workflow import StageDecorator
+
+from cpg_workflows.stages.aip import CreateAIPHTML, GenerateSeqrFile, ValidateMOI
+from cpg_workflows.stages.exomiser import ExomiserSeqrTSV, RunExomiser
+from cpg_workflows.stages.fastqc import FastQCMultiQC
+from cpg_workflows.stages.fraser import Fraser
+from cpg_workflows.stages.gatk_sv.gatk_sv_multisample_1 import (
+    FilterBatch,
+    GenotypeBatch,
+    MergeBatchSites,
+)
+from cpg_workflows.stages.gatk_sv.gatk_sv_multisample_2 import MtToEsSv
+from cpg_workflows.stages.gatk_sv.gatk_sv_single_sample import CreateSampleBatches
+from cpg_workflows.stages.gcnv import AnnotateCohortgCNV, AnnotateDatasetCNV, MtToEsCNV
+from cpg_workflows.stages.happy_validation import (
+    ValidationHappyOnVcf,
+    ValidationMtToVcf,
+    ValidationParseHappy,
+)
+from cpg_workflows.stages.mito import MitoReport
+from cpg_workflows.stages.outrider import Outrider
+from cpg_workflows.stages.seqr_loader import AnnotateDataset, DatasetVCF, MtToEs
+from cpg_workflows.stages.stripy import Stripy
+
+
+WORKFLOWS: dict[str, list[StageDecorator]] = {
+    'aip': [ValidateMOI, CreateAIPHTML, GenerateSeqrFile],
+    'exomiser': [RunExomiser, ExomiserSeqrTSV],
+    'pre_alignment': [FastQCMultiQC],
+    'seqr_loader': [
+        DatasetVCF,
+        AnnotateDataset,
+        MtToEs,
+        GvcfMultiQC,
+        CramMultiQC,
+        Stripy,
+        MitoReport,
+    ],
+    'validation': [ValidationMtToVcf, ValidationHappyOnVcf, ValidationParseHappy],
+    'large_cohort': [LoadVqsr, Frequencies, AncestryPlots, GvcfMultiQC, CramMultiQC],
+    'gatk_sv_singlesample': [CreateSampleBatches],
+    'gatk_sv_multisample_1': [FilterBatch, GenotypeBatch],
+    'gatk_sv_sandwich': [MergeBatchSites],  # stage to run between FilterBatch & GenotypeBatch
+    'gatk_sv_multisample_2': [MtToEsSv],
+    'rare_disease_rnaseq': [Outrider, Fraser],
+    'gcnv': [AnnotateCohortgCNV, AnnotateDatasetCNV, MtToEsCNV],
+}
\ No newline at end of file