generate_workflow.py

#!/usr/bin/python3
# This script will take a workflow yaml file and generate Snakefile, params file, docker recipe and copy all needed files (scripts)
# Accepts four parameters: workflow_yaml, output_dir, waw_dir, (local_config)
# Usage: ./generate_workflow.py ./workflows/my_workflow/my_workflow.yaml ./output_dir ./ ./local_configs/cluster_mbb_isem.yml
import os
import collections
import shutil
import sys
import re
import subprocess

from tools import *

WORKFLOW_YAML = collections.OrderedDict()
TOOLS_YAML = collections.OrderedDict()
BASE_CITATIONS_YAML = collections.OrderedDict()
BASE_VERSIONS_YAML = collections.OrderedDict()

STEPS = collections.OrderedDict()
TOOLS = list()

# Import all yaml files (workflow and tools)
def import_yaml_files(workflow_yaml, waw_dir):
    # Import workflow yaml
    global WORKFLOW_YAML
    WORKFLOW_YAML = read_yaml(workflow_yaml)
    
    # Import steps
    global STEPS
    for step in WORKFLOW_YAML["steps"]:
        STEPS[step["name"]] = dict()
        STEPS[step["name"]]["name"] = step["name"]
        STEPS[step["name"]]["title"] = step["title"]
        STEPS[step["name"]]["tools"] = step["tools"]
        STEPS[step["name"]]["default"] = step["default"]

        # Import every tools yaml
        global TOOLS
        global TOOLS_YAML
        for tool in step["tools"]:
            TOOLS.append(tool)
            TOOLS_YAML[tool] = read_yaml(os.path.join(waw_dir + "/tools/" + tool + "/" + tool + ".yaml"))
    
    # Import base yaml (versions and citations for base container)
    global BASE_CITATIONS_YAML
    BASE_CITATIONS_YAML = read_yaml(os.path.join(waw_dir+"/Docker_base/citations.yaml"))
    global BASE_VERSIONS_YAML
    BASE_VERSIONS_YAML = read_yaml(os.path.join(waw_dir+"/Docker_base/versions.yaml"))

def get_params_to_remove():
    # Get params to remove
    to_remove = list()
    if "params_equals" in WORKFLOW_YAML.keys():
        params_equals = WORKFLOW_YAML["params_equals"]
        for line in params_equals:
            if (len(line) == 1):
                to_remove.append(line["remove"])
            if (len(line) == 2):
                to_remove.append(line["param_B"])
    return to_remove

# Generate tool parameters with default values
def generate_tool_params(to_remove, waw_dir = "./"):
    TOOL_PARAMS = collections.OrderedDict()

    # Global workflow parameters
    for option in WORKFLOW_YAML["options"]:
        if (option["type"] == "input_file"):
            TOOL_PARAMS[option["name"]+"_select"] = "server"
        TOOL_PARAMS[option["name"]] = option["value"] if ("value" in option.keys()) else ""

    # input parameters
    if "input" in WORKFLOW_YAML:
        for raw_input in WORKFLOW_YAML["input"]:
               raw_inputs_yaml = read_yaml(os.path.join(waw_dir+"/raw_inputs/" + raw_input + ".yaml"))
               for option in raw_inputs_yaml["options"]:
                      TOOL_PARAMS[option["name"]] = option["value"] if ("value" in option.keys()) else ""

    # For each step
    for step_name, step_yaml in STEPS.items():
        TOOL_PARAMS[step_name] = step_yaml["default"]
        # For each tool that can be selected in a step
        for tool in step_yaml["tools"]:
            if "script" in TOOLS_YAML[tool]:
                TOOL_PARAMS[step_name + "__" + TOOLS_YAML[tool]["id"] + "_script"] = "scripts/" + TOOLS_YAML[tool]["script"]
            # For each command (rule snakemake) in the tool
            for command in TOOLS_YAML[tool]["commands"]:
                TOOL_PARAMS[step_name + "__" + command["name"] + "_output_dir"] = step_name + "/" + command["output_dir"]
                TOOL_PARAMS[step_name + "__" + command["name"] + "_command"] = command["command"]
                # For each option in the command
                for option in command["options"]:
                    if step_name + "__" + option["name"] not in to_remove:
                        # Add parameter (and 'select parameter' in case of input_file)
                        if (option["type"] == "input_file"):
                            TOOL_PARAMS[step_name + "__" + option["name"]+"_select"] = "server"
                        TOOL_PARAMS[step_name + "__" + option["name"]] = option["value"] if ("value" in option.keys()) else ""
    return TOOL_PARAMS

# Generate parameters info (usefull for shiny app)
def generate_params_info(to_remove, waw_dir = "./"):
    PARAMS_INFO = collections.OrderedDict()

    for option in WORKFLOW_YAML["options"]:
        if (option["type"] == "input_file"):
            PARAMS_INFO[option["name"]+"_select"] = collections.OrderedDict()
            PARAMS_INFO[option["name"]+"_select"]["type"] = "select"
        PARAMS_INFO[option["name"]] = collections.OrderedDict()
        PARAMS_INFO[option["name"]]["type"] = option["type"]

    # input parameters
    if "input" in WORKFLOW_YAML:
        for raw_input in WORKFLOW_YAML["input"]:
            raw_inputs_yaml = read_yaml(os.path.join(waw_dir+"/raw_inputs/" + raw_input + ".yaml"))
            for option in raw_inputs_yaml["options"]:
                PARAMS_INFO[option["name"]] = collections.OrderedDict()
                PARAMS_INFO[option["name"]]["type"] = option["type"]

    # For each step
    for step_name, step_yaml in STEPS.items():
        # For each tool that can be selected in a step
        for tool in step_yaml["tools"]:
            # For each command (rule snakemake) in the tool
            for command in TOOLS_YAML[tool]["commands"]:
                # For each option in the command
                for option in command["options"]:
                    if step_name + "__" + option["name"] not in to_remove:
                        # Add parameter info (and 'select parameter' in case of input_file)
                        if (option["type"] == "input_file"):
                            PARAMS_INFO[step_name + "__" + option["name"] + "_select"] = collections.OrderedDict()
                            PARAMS_INFO[step_name + "__" + option["name"] + "_select"]["tool"] = tool
                            PARAMS_INFO[step_name + "__" + option["name"] + "_select"]["rule"] = step_name + "_" + command["name"]
                            PARAMS_INFO[step_name + "__" + option["name"] + "_select"]["type"] = "select"
                        PARAMS_INFO[step_name + "__" + option["name"]] = collections.OrderedDict()
                        PARAMS_INFO[step_name + "__" + option["name"]]["tool"] = tool
                        PARAMS_INFO[step_name + "__" + option["name"]]["rule"] = step_name + "_" + command["name"]
                        PARAMS_INFO[step_name + "__" + option["name"]]["type"] = option["type"]
                        PARAMS_INFO[step_name + "__" + option["name"]]["label"] = option["label"]

    return PARAMS_INFO

def generate_stop_cases():
    STOP_CASES = collections.OrderedDict()

    if "stop_cases" in WORKFLOW_YAML.keys():
        STOP_CASES = WORKFLOW_YAML["stop_cases"]

    return STOP_CASES

# Generate tools outputs
def generate_outputs():
    OUTPUTS = collections.OrderedDict()
    # For each step
    for step_name, step_yaml in STEPS.items():
        # For each tool that can be selected in a step
        for tool in step_yaml["tools"]:
            # For each command (rule snakemake) in the tool
            OUTPUTS[step_name + "__" + tool] = dict()
            for command in TOOLS_YAML[tool]["commands"]:
                OUTPUTS[step_name + "__" + tool][step_name + "__" + command["name"]] = command["outputs"]
    return OUTPUTS

# Generate multiqc (for multiqc config file generation)
def generate_multiqc():
    MULTIQC = collections.OrderedDict()

    for tool_name, tool_yaml in TOOLS_YAML.items():
        MULTIQC[tool_name] = tool_yaml["multiqc"]

    return MULTIQC

# Generate citations for all tools and dependencies
def generate_citations(output_dir):
    CITATIONS = collections.OrderedDict()

    CITATIONS["base_tools"] = BASE_CITATIONS_YAML["citations"]
    for tool_name, tool_yaml in TOOLS_YAML.items():
        if "citations" in tool_yaml:
            CITATIONS[tool_name] = tool_yaml["citations"]
    
    return CITATIONS

# Generate versions of all tools and dependencies
def generate_versions(output_dir):
    VERSIONS = collections.OrderedDict()

    VERSIONS["base_tools"] = BASE_VERSIONS_YAML["versions"]
    for tool_name, tool_yaml in TOOLS_YAML.items():
        if "version" in tool_yaml:
            VERSIONS[tool_name] = tool_yaml["version"]

    return VERSIONS

# Generate list of prepare report scripts
def generate_prepare_report_scripts():
    PREPARE_REPORT_SCRIPTS = dict()
    scripts_list = list()
    # For each step
    for step_name, step_yaml in STEPS.items():
        # For each tool that can be selected in a step
        for tool in step_yaml["tools"]:
            if "prepare_report_script" in TOOLS_YAML[tool]:
                res = dict()
                res["tool"] = tool
                res["step"] = step_name
                res["script"] = TOOLS_YAML[tool]["prepare_report_script"]
                PREPARE_REPORT_SCRIPTS[step_name + "__" + tool] = res
                scripts_list.append(step_name + "__" + TOOLS_YAML[tool]["prepare_report_script"])
    PREPARE_REPORT_SCRIPTS["SCRIPTS"] = scripts_list
    return PREPARE_REPORT_SCRIPTS

# Generate list of prepare report scripts
def generate_tools_scripts():
    TOOLS_SCRIPTS = dict()
    for tool_yaml in TOOLS_YAML.values():
        if "script" in tool_yaml.keys():
            TOOLS_SCRIPTS[tool_yaml["id"]] =tool_yaml["script"]

    return TOOLS_SCRIPTS

# Generate list of prepare report outputs
def generate_prepare_report_outputs():
    PREPARE_REPORT_OUTPUTS = collections.OrderedDict()
    # For each step
    for step_name, step_yaml in STEPS.items():
        # For each tool that can be selected in a step
        for tool in step_yaml["tools"]:
            if "prepare_report_outputs" in TOOLS_YAML[tool]:
                PREPARE_REPORT_OUTPUTS[step_name + "__" + tool] = TOOLS_YAML[tool]["prepare_report_outputs"]

    return PREPARE_REPORT_OUTPUTS

def generate_snake_rule(step_name,tool_name, waw_dir = "./"):
    RULE = "\n"
    
    with open(os.path.join(waw_dir + "/tools/" + tool_name + "/" + tool_name + ".rule.snakefile"), "r") as rule:
        RULE += rule.read()

    RULE += "\n"

    RULE = RULE.replace("<step_name>",step_name)

    return RULE

# Generate Snakefile
def generate_snakefile(workflow_yaml, waw_dir = "./"):
    SNAKEFILE = ""
    RULES = ""

    # import tools
    SNAKEFILE += "from tools import *\n"

    # import raw_input function
    if "input" in WORKFLOW_YAML:
       for raw_input in WORKFLOW_YAML["input"]:
             SNAKEFILE += "from "+ raw_input +" import "+ raw_input+"\n"

    # configure workdir (snakemake -d)
    SNAKEFILE += "workdir: config['params']['results_dir']\n"

    # Open template

    with open(workflow_yaml.replace(".yaml",".snakefile"), "r") as template:
        SNAKEFILE += template.read()

    # import wildcards and imports
    with open(os.path.join(waw_dir + "/workflows/global_imports.py"), "r") as imports:
        GLOBAL_IMPORTS = imports.read()
        SNAKEFILE = SNAKEFILE.replace("{import global_imports}",GLOBAL_IMPORTS)

    # Global functions
    with open(os.path.join(waw_dir + "/workflows/global_functions.py")) as global_functions:
        GLOBAL_FUNCTIONS = global_functions.read()
        SNAKEFILE = SNAKEFILE.replace("{import global_functions}",GLOBAL_FUNCTIONS)

    # Global rules
    with open(os.path.join(waw_dir + "/workflows/global_rules.snakefile")) as global_rules:
        GLOBAL_RULES = global_rules.read()
        SNAKEFILE = SNAKEFILE.replace("{import global_rules}",GLOBAL_RULES)
    
    for step_name, step_yaml in STEPS.items():
        for tool in step_yaml["tools"]:
            RULES += generate_snake_rule(step_name,tool,waw_dir)

    SNAKEFILE = SNAKEFILE.replace("{import rules}", RULES)
    
    #####
    # Replace params
    if "params_equals" in WORKFLOW_YAML.keys():
        for line in WORKFLOW_YAML["params_equals"]:
            if (len(line) == 2):
                SNAKEFILE = SNAKEFILE.replace("config[\""+line["param_B"]+"\"]","config[\""+line["param_A"]+"\"]")
    #####

    return SNAKEFILE

# Generate Dockerfile
def generate_dockerfile(output_dir, local_config, waw_dir = "./"):
    workflow_name = WORKFLOW_YAML["name"]
    DOCKERFILE = ""
    # Open template
    with open(waw_dir+"/Dockerfile.template", "r") as template:
        DOCKERFILE = template.read()
    DOCKERFILE += "\n"

    # COPY files if data in workflow yaml
    files = ""
    if "data" in WORKFLOW_YAML.keys():
        for data in WORKFLOW_YAML["data"]:
            file = os.path.normpath(waw_dir + "/workflows/" + workflow_name + "/data/" + data["name"])
            files += "COPY ./files/data/" + data["name"] + " /" + data["name"] + "\n"
            # Copy files
            copy_dir(file,output_dir+'/files/data/'+data["name"])

    # COPY files if data in tools yaml
    # And import install commands
    tools_installs = collections.OrderedDict()
    for tool_name, tool_yaml in TOOLS_YAML.items():
        tools_installs.update(tool_yaml["install"])
        if "data" in tool_yaml.keys():
            for data in tool_yaml["data"]:
                file = os.path.normpath(waw_dir + "/tools/" + tool_name + "/data/" + data["name"])
                files += "COPY ./files/data/" + data["name"] + " /" + data["name"] + "\n"
                # Copy files
                shutil.copytree(file, output_dir+workflow_name + '/files/data/' + data["name"])

    files += "COPY files /workflow\n"
    files += "COPY sagApp /sagApp\n\n"

    for value in tools_installs.values():
        tool_install="RUN "
        tool_path=""
        for line in value:
            if "ENV" in line:
                tool_path+= line + "\n"
            else:
                tool_install += line + " \\\n && "
        tool_install = tool_install[:-7] # remove trailing '\n && '
        tool_install += "\n\n"

        if (tool_path!=""):
            DOCKERFILE += tool_path    
        DOCKERFILE += tool_install

    # Local Config
    if local_config != "default":
        local_config_data = read_yaml(local_config)
        if "docker" in local_config_data.keys():
            DOCKERFILE += '\n'.join(local_config_data["docker"]["env"])
            DOCKERFILE += '\n\n'
            DOCKERFILE += '\n'.join(local_config_data["docker"]["run"])
    DOCKERFILE += "\n\n"
    DOCKERFILE += "EXPOSE 3838\n"
    DOCKERFILE += "CMD [\"Rscript\", \"-e\", \"setwd('/sagApp/'); shiny::runApp('/sagApp/app.R',port=3838 , host='0.0.0.0')\"]\n"

    # Ajout des copy après les installs
    DOCKERFILE += "\n\n"
    DOCKERFILE += "FROM alltools\n\n"
    DOCKERFILE += files
    
    return DOCKERFILE

# Generate params yaml with all parameters and informations
def generate_params_yaml(workflow_name, to_remove, waw_dir = "./"):
    PARAMS = collections.OrderedDict()

    PARAMS["pipeline"] = workflow_name
    PARAMS["params"] = generate_tool_params(to_remove, waw_dir)
    PARAMS["steps"] = list(STEPS.values())
    PARAMS["params_info"] = generate_params_info(to_remove, waw_dir)
    PARAMS["prepare_report_scripts"] = list(generate_prepare_report_scripts()["SCRIPTS"])
    PARAMS["prepare_report_outputs"] = generate_prepare_report_outputs()
    PARAMS["outputs"] = generate_outputs()
    PARAMS["multiqc"] = generate_multiqc()
    PARAMS["stop_cases"] = generate_stop_cases()

    return PARAMS

def generate_pipeline_files(workflow_yaml, output_dir, waw_dir, local_config="default"):

    import_yaml_files(workflow_yaml, waw_dir)

    workflow_name = WORKFLOW_YAML["name"]

    # Create output directory if needed
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)
    if not os.path.isdir(output_dir+"/files"):
        os.mkdir(output_dir+"/files")

    ### Generate all output files

    write_yaml(output_dir + "/files/citations.yaml", generate_citations(output_dir))

    write_yaml(output_dir + "/files/versions.yaml", generate_versions(output_dir))
    
    with open(output_dir + "/files/Snakefile", "w") as out:
        out.write(generate_snakefile(workflow_yaml, waw_dir))

    with open(output_dir + "/Dockerfile", "w") as out:
        out.write(generate_dockerfile(output_dir, local_config, waw_dir))

    write_yaml(output_dir + "/files/params.total.yml", generate_params_yaml(workflow_name, get_params_to_remove(), waw_dir))
    ###1

    ### Copy scripts and other files

    NB_SCRIPTS = len(list(generate_prepare_report_scripts()["SCRIPTS"])) + len(list(generate_tools_scripts()))

    if not os.path.isdir(output_dir + "/files/scripts"):
        os.mkdir(output_dir + "/files/scripts")
    if (NB_SCRIPTS==0):
        with open(output_dir + "/files/scripts/.gitignore",'w') as gitignore:
            to_write ="""# gitignore to force creation of scripts dir
!.gitignore
"""
            gitignore.write(to_write)
    else:

        for name, script_dict in generate_prepare_report_scripts().items():
            if name!= "SCRIPTS":
                tool = script_dict["tool"]
                script = script_dict["script"]
                step = script_dict["step"]
                returned_value = subprocess.call("sed 's/<step_name>/"+step+"/g' " + os.path.join(waw_dir + "/tools/" + tool + "/" + script) + " > " + output_dir + "/files/scripts/" + step + "__" + script, shell=True)

        for tool, script in generate_tools_scripts().items():
            shutil.copy(os.path.join(waw_dir+"/tools/" + tool + "/" + script), output_dir+ "/files/scripts")

    shutil.copy(os.path.join(waw_dir+"/generate_multiqc_config.py"), output_dir+ "/files")
    shutil.copy(os.path.join(waw_dir+"/tools.py"), output_dir + "/files")
    shutil.copy(os.path.join(waw_dir+"/mbb.png"), output_dir + "/files")

    if "input" in WORKFLOW_YAML:
        for raw_input in WORKFLOW_YAML["input"]:
            shutil.copy(os.path.join(waw_dir+"/raw_inputs/"+raw_input+".py"), output_dir+"/files")

    if (local_config != "default"):
        local_config_data = read_yaml(local_config)
        if ("qsub" in local_config_data.keys()):
            with open(output_dir+"/waw_workflow.qsub", 'w') as qsub_file:
                qsub_file.writelines([ a.replace("{workflow_name}",WORKFLOW_YAML["docker_name"]) +'\n' for a in local_config_data["qsub"]] )

    # deploys scripts
    for script in os.listdir(os.path.join(waw_dir+"/deploys")):
        with open(os.path.join(waw_dir+"/deploys/"+script),"r") as infile:
            lines = infile.readlines()
            with open(output_dir+"/"+script,"w") as outfile:
                outfile.writelines([ a.replace("{workflow_name}",WORKFLOW_YAML["docker_name"]) for a in lines] )


    ###

def main():
    #print(len(sys.argv))
    if len(sys.argv) == 5:
        generate_pipeline_files(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
    elif len(sys.argv) == 4:
        generate_pipeline_files(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        exit("""Need at least three parameters: workflow_yaml, output_dir, waw_dir, (local_config)
Usage: ./generate_workflow.py workflow_yaml output_dir waw_dir (local_config)""")

if __name__ == "__main__":
    # execute only if run as a script
    main()