generate_workflow_snakefile.py

#!/usr/bin/python3
# This script will take an workflow and generate snakefile
# Accepts 1 parameter:
# Usage: ./generate_workflow_snakefile.py workflow_name


from tools import read_yaml
import sys
import tempfile
import shutil
import os

dirpath = tempfile.mkdtemp()

DEFAULT_PATH_OUTPUT = "./"
DEFAULT_PATH_INPUT = "./workflows/"

# Return True when {sample} is in inputs files
def wildcard_in_inputs(rule_name, tool_yaml):
    found = False
    for command in tool_yaml["commands"]:
        if command["name"] == rule_name:
            found = True
            for inputt in command["inputs"]:
                if "file" in inputt:
                    if "{sample}" in inputt["file"]:
                        return True
                if inputt["type"] == "reads":
                    return True
    if found: # command found but no wildcard
        return False
    else:
        exit("Command "+rule_name+" not found")

# Return True when {sample} is in outputs files
def wildcard_in_outputs(rule_name, tool_yaml):
    found = False
    for command in tool_yaml["commands"]:
        if command["name"] == rule_name:
            found = True
            for output in command["outputs"]:
                if "file" in output:
                    if "{sample}" in output["file"]:
                        return True
    if found: # command found but no wildcard
        return False
    else:
        exit("Command "+rule_name+" not found")

def generate(name, path_yaml = "", path_input = DEFAULT_PATH_INPUT, path_output = DEFAULT_PATH_OUTPUT):

    result = "# File generated with generate_workflow_snakefile.py\n\n"

    yaml = None

    if(path_yaml == ""):
        yaml = read_yaml(path_input+name+"/"+name+".yaml")
    else:
        #print(path_yaml)
        yaml = read_yaml(path_yaml)

    # imports

    result += "{import global_imports}\n\n"

    # intputs

    result += "##########\n"
    result += "# Inputs #\n"
    result += "##########\n"

    result += "\n"

    if "input" in yaml:
        for raw_input in yaml["input"]:
                raw_inputs_yaml = read_yaml(path_input+"/raw_inputs/"+raw_input+".yaml")
                result += "# raw_inputs function call\n"
                result += raw_input+" = "+raw_inputs_yaml["function_call"]+"\n"
                result += "config.update("+raw_input+")\n"
                # !!! possibilité de définir samples plusieurs fois si il y a raw_reads et raw_vcf ... !!!
                # il faut vérifier que les samples sont les mêmes ?
                if  not raw_input.endswith("_index") and not raw_input.endswith("_popmap") and not raw_input.endswith("_vcfFile") and not raw_input.endswith("_gff3"): 
                     result += "SAMPLES = "+raw_input+"['samples']\n"

    result += "\n"

    result += "# Tools inputs functions\n"

    if "steps_in" in yaml:

        for step in yaml["steps_in"]:

            TOOL_YAML = read_yaml(path_input+"/tools/"+step["tool_name"]+"/"+step["tool_name"]+".yaml")

            if "rule_name" in step:
                # get rule yaml
                RULE_YAML = dict()
                for command in TOOL_YAML["commands"]:
                    if command["name"] == step["rule_name"]:
                        RULE_YAML = command

                # wildcards in inputs and outputs
                wildcard_in = wildcard_in_inputs(step["rule_name"],TOOL_YAML)
                wildcard_out = wildcard_in_outputs(step["rule_name"],TOOL_YAML)

                # wildcard in inputs but not in outputs => expand on inputs
                if wildcard_in and not wildcard_out:
                    expand_begin = "expand("
                    expand_end = ",sample=SAMPLES)"
                else:
                    expand_begin = ""
                    expand_end = ""

                # Liste des inputs dans la commande du yaml de l'outil
                #inputs_list_from_command = [inputt["name"] for inputt in RULE_YAML["inputs"]]
                # Liste des inputs dans le steps_in du yaml du workflow
                inputs_list_from_yaml = list()

                result += "def " + step["step_name"] + "__" + step["rule_name"] + "_inputs():\n"
                result += "\tinputs = dict()\n"

                # Parcourir le yaml du tool, les commandes et leurs inputs
                # Traiter les inputs de type liste
                # Gérer les "from params" (ex: fasta genomes)

                if "params" in step:
                    raw_inputs = list()
                    for param in step["params"]:

                        # get input yaml
                        INPUT_YAML = dict()
                        for inputt in RULE_YAML["inputs"]:
                            if inputt["name"] == param["input_name"]:
                                INPUT_YAML = inputt

                        # Inputs de type expand
                        if ("expand" in INPUT_YAML and INPUT_YAML["expand"]):
                            if param["input_name"] not in inputs_list_from_yaml:
                                result += "\tinputs[\"" + param["input_name"] + "\"] = list()\n"
                            if "raw_" in param["origin_command"]:
                                #if not param["origin_command"] in raw_inputs:
                                #   result += "\tinputs[\"" + param["input_name"] + "\"].append(" + expand_begin + param["origin_command"] + "[\"" + param["origin_name"] + "\"]" + expand_end + ")\n"
                                #   raw_inputs.append(param["origin_command"])
                                result += "\tinputs[\"" + param["input_name"] + "\"] = " + expand_begin + param["origin_command"] + "[\"" + param["origin_name"] + "\"] " + expand_end + "\n"
                            else:
                                result += "\tinputs[\"" + param["input_name"] + "\"] = expand(" + expand_begin + "rules." + param["origin_step"] + "__" + param["origin_command"] + ".output." + param["origin_name"] + expand_end + ", sample=SAMPLES)\n"
                        # Inputs de type liste
                        elif ("list" in INPUT_YAML and INPUT_YAML["list"]):
                            if param["input_name"] not in inputs_list_from_yaml:
                                result += "\tinputs[\"" + param["input_name"] + "\"] = list()\n"
                            if "raw_" in param["origin_command"]:
                                if not param["origin_command"] in raw_inputs:
                                   result += "\tinputs[\"" + param["input_name"] + "\"].append(" + expand_begin + param["origin_command"] + "[\"" + param["origin_name"] + "\"]" + expand_end + ")\n"
                                   raw_inputs.append(param["origin_command"])
                            else:
                                result += "\tinputs[\"" + param["input_name"] + "\"].append(" + expand_begin + "rules." + param["origin_step"] + "__" + param["origin_command"] + ".output." + param["origin_name"] + expand_end + ")\n"
                        # inputs classiques
                        else:
                            if "raw_" in param["origin_command"] :
                                result += "\tinputs[\"" + param["input_name"] + "\"] = " + expand_begin + param["origin_command"] + "[\"" + param["origin_name"] + "\"] " + expand_end + "\n"
                            else:
                                result += "\tinputs[\"" + param["input_name"] + "\"] = " + expand_begin + "rules." + param["origin_step"] + "__" + param["origin_command"] + ".output." + param["origin_name"] + expand_end + "\n"
                        
                        inputs_list_from_yaml.append(param["input_name"]) # input traités (ou vu une fois si de type liste)
                    
                    # Pour tous les inputs d'une commande qui ne sont pas liés dans le steps_in
                    for inputt in RULE_YAML["inputs"]:
                        if inputt["name"] not in inputs_list_from_yaml and "from" in inputt and inputt["from"] == "parameter":
                            result += "\tinputs[\"" + inputt["name"] + "\"] = config[\""+ step["step_name"] + "__" + step["rule_name"] + "_" + inputt["name"] +"\"]\n"

                result += "\treturn inputs\n"
                result += "\n"

    result += "\n"

    result += "{import global_functions}\n"

    result += "\n"

    # outputs

    result += "###########\n"
    result += "# Outputs #\n"
    result += "###########\n"

    result += "\n"

    result += "def step_outputs(step):\n"
    result += "\toutputs = list()\n"

    if "steps_in" in yaml:

        for step in yaml["steps_in"]:

            if ("step_name" in step) and ("rule_name" in step):
                # Check in tool yaml if {sample} in output files to add expand
                tool_yaml = read_yaml(path_input+"/tools/"+step["tool_name"]+"/"+step["tool_name"]+".yaml")
                need_expand = wildcard_in_outputs(step["rule_name"],tool_yaml)
                
                if need_expand:
                    expand_start = "expand("
                    expand_end = ", sample=SAMPLES)"
                else:
                    expand_start = ""
                    expand_end = ""

                add = ""
                if ("_SE" in step["rule_name"]):
                    add = "and config['SeOrPe'] == 'SE'"
                if ("_PE" in step["rule_name"]):
                    add = "and config['SeOrPe'] == 'PE'"
                result += "\tif (step == \"" + step["step_name"] + "\" " + add + " ):\n"
                result += "\t\toutputs = " + expand_start + "rules." + step["step_name"] + "__" + step["rule_name"] + ".output " + expand_end + "\n"
                result += "\t\t\n"


    result += "\tif (step == \"all\"):\n"
    result += "\t\toutputs = list(rules.multiqc.output)\n"

    result += "\n"

    result += "\treturn outputs\n"

    result += "\n"

    result += "# get outputs for each choosen tools\n"
    result += "def workflow_outputs(step):\n"
    result += "\toutputs = list()\n"
    result += "\toutputs.extend(step_outputs(step))\n"
    result += "\treturn outputs\n"

    result += "\n"

    # rules

    result += "#########\n"
    result += "# Rules #\n"
    result += "#########\n"

    result += "\n"

    result += "{import rules}\n"

    result += "{import global_rules}\n"

    #print(result)

    path_snakefile = path_output + name + "/" + name + ".snakefile"
    #print(path_snakefile)
    f = open(path_snakefile, "w")
    f.write(result)
    f.close()


def create_directory(output_path):

    access_rights = 0o755

    try:
        os.mkdir(output_path, access_rights)
    except OSError:
        print ("Creation of the directory %s failed" % output_path)



def main():

    if len(sys.argv) == 2:
        generate(sys.argv[1])

    elif len(sys.argv) > 2:

        name = sys.argv[1]
        output_path = sys.argv[2]
        path_yaml = output_path + name + "/" + name + ".yaml"
        input_path = sys.argv[3]
        
        generate(name, path_yaml, input_path, output_path)

    else:
        exit("""Needs 1 argument minimum: workflow_name
Usage: ./generate_workflow_snakefile.py workflow_name """)


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