generate_classes_edit.py

"""This script generates Slicer Interfaces based on the CLI modules XML. CLI
modules are selected from the hardcoded list below and generated code is placed
in the cli_modules.py file (and imported in __init__.py). For this to work
correctly you must have your CLI executabes in $PATH"""

from __future__ import print_function

import xml.dom.minidom
import subprocess
import os, glob
from shutil import rmtree

import keyword
python_keywords = keyword.kwlist  # If c++ SEM module uses one of these key words as a command line parameter, we need to modify variable

#from ...external.six import string_types
from six import string_types


def force_to_valid_python_variable_name(old_name):
    """  Valid c++ names are not always valid in python, so
    provide alternate naming

    >>> force_to_valid_python_variable_name('lambda')
    'opt_lambda'
    >>> force_to_valid_python_variable_name('inputVolume')
    'inputVolume'
    """
    new_name = old_name
    new_name = new_name.lstrip().rstrip()
    if old_name in python_keywords:
        new_name = 'opt_' + old_name
    return new_name


def add_class_to_package(class_codes, class_names, module_name, package_dir):
    module_python_filename = os.path.join(package_dir, "%s.py" % module_name)
    f_m = open(module_python_filename, 'w')
    f_i = open(os.path.join(package_dir, "__init__.py"), 'a+')
    f_m.write("""# -*- coding: utf8 -*-
\"\"\"Autogenerated file - DO NOT EDIT
If you spot a bug, please report it on the mailing list and/or change the generator.\"\"\"\n\n""")
    imports = """from nipype.interfaces.base import CommandLine, CommandLineInputSpec, SEMLikeCommandLine, TraitedSpec, File, Directory, traits, isdefined, InputMultiPath, OutputMultiPath
import os\n\n\n"""
    f_m.write(imports)
    f_m.write("\n\n".join(class_codes).encode('utf8'))
    f_i.write("from %s import %s\n" % (module_name, ", ".join(class_names)))
    f_m.close()
    f_i.close()


def crawl_code_struct(code_struct, package_dir):
    subpackages = []
    for k, v in code_struct.items():
        if isinstance(v, str) or isinstance(v, string_types):
            module_name = k.lower()
            class_name = k
            class_code = v
            add_class_to_package(
                [class_code], [class_name], module_name, package_dir)
        else:
            l1 = {}
            l2 = {}
            for key in list(v.keys()):
                if (isinstance(v[key], str) or isinstance(v[key], string_types)):
                    l1[key] = v[key]
                else:
                    l2[key] = v[key]
            if l2:
                v = l2
                subpackages.append(k.lower())
                f_i = open(os.path.join(package_dir, "__init__.py"), 'a+')
                f_i.write("from %s import *\n" % k.lower())
                f_i.close()
                new_pkg_dir = os.path.join(package_dir, k.lower())
                if os.path.exists(new_pkg_dir):
                    rmtree(new_pkg_dir)
                os.mkdir(new_pkg_dir)
                crawl_code_struct(v, new_pkg_dir)
                if l1:
                    for ik, iv in l1.items():
                        crawl_code_struct({ik: {ik: iv}}, new_pkg_dir)
            elif l1:
                v = l1
                module_name = k.lower()
                add_class_to_package(
                    list(v.values()), list(v.keys()), module_name, package_dir)
        if subpackages:
            f = open(os.path.join(package_dir, "setup.py"), 'w')
            f.write("""# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
def configuration(parent_package='',top_path=None):
    from numpy.distutils.misc_util import Configuration

    config = Configuration('{pkg_name}', parent_package, top_path)

    {sub_pks}

    return config

if __name__ == '__main__':
    from numpy.distutils.core import setup
    setup(**configuration(top_path='').todict())
""".format(pkg_name=package_dir.split("/")[-1], sub_pks="\n    ".join(["config.add_data_dir('%s')" % sub_pkg for sub_pkg in subpackages])))
            f.close()


def generate_all_classes(modules_list=[], launcher=[], redirect_x=False, mipav_hacks=False):
    """ modules_list contains all the SEM compliant tools that should have wrappers created for them.
        launcher containtains the command line prefix wrapper arugments needed to prepare
        a proper environment for each of the modules.
    """
    all_code = {}
    for module in modules_list:
        print("=" * 80)
        print("Generating Definition for module {0}".format(module))
        print("^" * 80)
        package, code, module = generate_class(module, launcher, redirect_x=redirect_x, mipav_hacks=mipav_hacks)
        cur_package = all_code
        module_name = package.strip().split(" ")[0].split(".")[-1]
        for package in package.strip().split(" ")[0].split(".")[:-1]:
            if package not in cur_package:
                cur_package[package] = {}
            cur_package = cur_package[package]
        if module_name not in cur_package:
            cur_package[module_name] = {}
        cur_package[module_name][module] = code
    if os.path.exists("__init__.py"):
        os.unlink("__init__.py")
    crawl_code_struct(all_code, os.getcwd())


def generate_class(module, launcher, strip_module_name_prefix=True, redirect_x=False, mipav_hacks=False):

    print(glob.glob('*'+module.lower()+'*'))

    dom = xml.dom.minidom.parseString(open(module.lower()+'.xml','rU').read().strip())
    if strip_module_name_prefix:
        module_name = module.split(".")[-1]
    else:
        module_name = module
    inputTraits = []
    outputTraits = []
    outputs_filenames = {}

    # self._outputs_nodes = []

    class_string = "\"\"\""

    for desc_str in ['title', 'category', 'description', 'version',
                     'documentation-url', 'license', 'contributor',
                     'acknowledgements']:
        el = dom.getElementsByTagName(desc_str)
        if el and el[0].firstChild and el[0].firstChild.nodeValue.strip():
            class_string += desc_str + ": " + el[0].firstChild.nodeValue.strip() + "\n\n"
        if desc_str == 'category':
            category = el[0].firstChild.nodeValue.strip()
    class_string += "\"\"\""

    for paramGroup in dom.getElementsByTagName("parameters"):
        indices = paramGroup.getElementsByTagName('index')
        max_index = 0
        for index in indices:
            if int(index.firstChild.nodeValue) > max_index:
                max_index = int(index.firstChild.nodeValue)
        for param in paramGroup.childNodes:
            if param.nodeName in ['label', 'description', '#text', '#comment']:
                continue
            traitsParams = {}

            longFlagNode = param.getElementsByTagName('longflag')
            if longFlagNode:
                # Prefer to use longFlag as name if it is given, rather than the parameter name
                longFlagName = longFlagNode[0].firstChild.nodeValue
                # SEM automatically strips prefixed "--" or "-" from from xml before processing
                # we need to replicate that behavior here The following
                # two nodes in xml have the same behavior in the program
                # <longflag>--test</longflag>
                # <longflag>test</longflag>
                longFlagName = longFlagName.lstrip(" -").rstrip(" ")
                name = longFlagName
                name = force_to_valid_python_variable_name(name)
                traitsParams["argstr"] = "--" + longFlagName + " "
            else:
                name = param.getElementsByTagName(
                    'name')[0].firstChild.nodeValue
                name = force_to_valid_python_variable_name(name)
                if param.getElementsByTagName('index'):
                    traitsParams["argstr"] = ""
                else:
                    traitsParams["argstr"] = "--" + name + " "

            if param.getElementsByTagName('description') and param.getElementsByTagName('description')[0].firstChild:
                traitsParams["desc"] = param.getElementsByTagName('description')[0].firstChild.nodeValue.replace('"', "\\\"").replace("\n", ", ")

            argsDict = {'directory': '%s', 'file': '%s', 'integer': "%d",
                        'double': "%f", 'float': "%f", 'image': "%s",
                        'transform': "%s", 'boolean': '',
                        'string-enumeration': '%s', 'string': "%s",
                        'integer-enumeration': '%s',
                        'table': '%s', 'point': '%s', 'region': '%s', 'geometry': '%s'}

            if param.nodeName.endswith('-vector'):
                traitsParams["argstr"] += "%s"
            else:
                traitsParams["argstr"] += argsDict[param.nodeName]

            index = param.getElementsByTagName('index')
            if index:
                traitsParams["position"] = int(
                    index[0].firstChild.nodeValue) - (max_index + 1)

            desc = param.getElementsByTagName('description')
            if index:
                traitsParams["desc"] = desc[0].firstChild.nodeValue

            typesDict = {'integer': "traits.Int", 'double': "traits.Float",
                         'float': "traits.Float", 'image': "File",
                         'transform': "File", 'boolean': "traits.Bool",
                         'string': "traits.Str", 'file': "File", 'geometry': "File",
                         'directory': "Directory", 'table': "File",
                         'point': "traits.List", 'region': "traits.List"}

            if param.nodeName.endswith('-enumeration'):
                type = "traits.Enum"
                values = ['"%s"' % str(el.firstChild.nodeValue).replace('"', '') for el in param.getElementsByTagName('element')]
            elif param.nodeName.endswith('-vector'):
                type = "InputMultiPath"
                if param.nodeName in ['file', 'directory', 'image', 'geometry', 'transform', 'table']:
                    values = ["%s(exists=True)" % typesDict[
                              param.nodeName.replace('-vector', '')]]
                else:
                    values = [typesDict[param.nodeName.replace('-vector', '')]]
                if mipav_hacks is True:
                    traitsParams["sep"] = ";"
                else:
                    traitsParams["sep"] = ','
            elif param.getAttribute('multiple') == "true":
                type = "InputMultiPath"
                if param.nodeName in ['file', 'directory', 'image', 'geometry', 'transform', 'table']:
                    values = ["%s(exists=True)" % typesDict[param.nodeName]]
                elif param.nodeName in ['point', 'region']:
                    values = ["%s(traits.Float(), minlen=3, maxlen=3)" %
                              typesDict[param.nodeName]]
                else:
                    values = [typesDict[param.nodeName]]
                traitsParams["argstr"] += "..."
            else:
                values = []
                type = typesDict[param.nodeName]

            if param.nodeName in ['file', 'directory', 'image', 'geometry', 'transform', 'table']:
                if not param.getElementsByTagName('channel'):
                    raise RuntimeError("Insufficient XML specification: each element of type 'file', 'directory', 'image', 'geometry', 'transform',  or 'table' requires 'channel' field.\n{0}".format(traitsParams))
                elif param.getElementsByTagName('channel')[0].firstChild.nodeValue == 'output':
                    traitsParams["hash_files"] = False
                    inputTraits.append(
                        "%s = traits.Either(traits.Bool, %s(%s), %s)" % (name,
                                                                         type,
                                                                         parse_values(
                                                                             values).replace("exists=True", ""),
                                                                         parse_params(traitsParams)))
                    traitsParams["exists"] = True
                    traitsParams.pop("argstr")
                    traitsParams.pop("hash_files")
                    outputTraits.append("%s = %s(%s%s)" % (name, type.replace("Input", "Output"), parse_values(values), parse_params(traitsParams)))

                    outputs_filenames[
                        name] = gen_filename_from_param(param, name)
                elif param.getElementsByTagName('channel')[0].firstChild.nodeValue == 'input':
                    if param.nodeName in ['file', 'directory', 'image', 'geometry', 'transform', 'table'] and type not in ["InputMultiPath", "traits.List"]:
                        traitsParams["exists"] = True
                    inputTraits.append("%s = %s(%s%s)" % (name, type, parse_values(values), parse_params(traitsParams)))
                else:
                    raise RuntimeError("Insufficient XML specification: each element of type 'file', 'directory', 'image', 'geometry', 'transform',  or 'table' requires 'channel' field to be in ['input','output'].\n{0}".format(traitsParams))
            else:  # For all other parameter types, they are implicitly only input types
                inputTraits.append("%s = %s(%s%s)" % (name, type, parse_values(
                    values), parse_params(traitsParams)))

    if mipav_hacks:
        blacklisted_inputs = ["maxMemoryUsage"]
        inputTraits = [trait for trait in inputTraits if trait.split()[0] not in blacklisted_inputs]

        compulsory_inputs = ['xDefaultMem = traits.Int(desc="Set default maximum heap size", argstr="-xDefaultMem %d")',
                             'xMaxProcess = traits.Int(1, desc="Set default maximum number of processes.", argstr="-xMaxProcess %d", usedefault=True)']
        inputTraits += compulsory_inputs

    input_spec_code = "class " + module_name + "InputSpec(CommandLineInputSpec):\n"
    for trait in inputTraits:
        input_spec_code += "    " + trait + "\n"

    output_spec_code = "class " + module_name + "OutputSpec(TraitedSpec):\n"
    if not outputTraits:
        output_spec_code += "    pass\n"
    else:
        for trait in outputTraits:
            output_spec_code += "    " + trait + "\n"

    output_filenames_code = "_outputs_filenames = {"
    output_filenames_code += ",".join(["'%s':'%s'" % (
        key, value) for key, value in outputs_filenames.items()])
    output_filenames_code += "}"

    input_spec_code += "\n\n"
    output_spec_code += "\n\n"

    template = """class %module_name%(SEMLikeCommandLine):
    %class_str%

    input_spec = %module_name%InputSpec
    output_spec = %module_name%OutputSpec
    _cmd = "%launcher% %name% "
    %output_filenames_code%\n"""
    template += "    _redirect_x = {0}\n".format(str(redirect_x))

    main_class = template.replace('%class_str%', class_string).replace("%module_name%", module_name).replace("%name%", module).replace("%output_filenames_code%", output_filenames_code).replace("%launcher%", " ".join(launcher))

    return category, input_spec_code + output_spec_code + main_class, module_name


def grab_xml(module, launcher, mipav_hacks=False):
    #        cmd = CommandLine(command = "Slicer3", args="--launch %s --xml"%module)
    #        ret = cmd.run()
    command_list = launcher[:]  # force copy to preserve original
    command_list.extend([module, "--xml"])
    final_command = " ".join(command_list)
    xmlReturnValue = subprocess.Popen(
        final_command, stdout=subprocess.PIPE, shell=True).communicate()[0]
    if mipav_hacks:
        # workaround for a jist bug https://www.nitrc.org/tracker/index.php?func=detail&aid=7234&group_id=228&atid=942
        new_xml = ""
        replace_closing_tag = False
        for line in xmlReturnValue.splitlines():
            if line.strip() == "<file collection: semi-colon delimited list>":
                new_xml += "<file-vector>\n"
                replace_closing_tag = True
            elif replace_closing_tag and line.strip() == "</file>":
                new_xml += "</file-vector>\n"
                replace_closing_tag = False
            else:
                new_xml += line + "\n"

        xmlReturnValue = new_xml

        # workaround for a JIST bug https://www.nitrc.org/tracker/index.php?func=detail&aid=7233&group_id=228&atid=942
        if xmlReturnValue.strip().endswith("XML"):
            xmlReturnValue = xmlReturnValue.strip()[:-3]
        if xmlReturnValue.strip().startswith("Error: Unable to set default atlas"):
            xmlReturnValue = xmlReturnValue.strip()[len("Error: Unable to set default atlas"):]
    try:
        dom = xml.dom.minidom.parseString(xmlReturnValue.strip())
    except Exception as e:
        print(xmlReturnValue.strip())
        raise e
    return dom
#        if ret.runtime.returncode == 0:
#            return xml.dom.minidom.parseString(ret.runtime.stdout)
#        else:
#            raise Exception(cmd.cmdline + " failed:\n%s"%ret.runtime.stderr)


def parse_params(params):
    list = []
    for key, value in params.items():
        if isinstance(value, string_types):
            list.append('%s="%s"' % (key, value.replace('"', "'")))
        else:
            list.append('%s=%s' % (key, value))

    return ", ".join(list)


def parse_values(values):
    values = ['%s' % value for value in values]
    if len(values) > 0:
        retstr = ", ".join(values) + ", "
    else:
        retstr = ""
    return retstr


def gen_filename_from_param(param, base):
    fileExtensions = param.getAttribute("fileExtensions")
    if fileExtensions:
        # It is possible that multiple file extensions can be specified in a
        # comma separated list,  This will extract just the first extension
        firstFileExtension = fileExtensions.split(',')[0]
        ext = firstFileExtension
    else:
        ext = {'image': '.nii', 'transform': '.mat', 'file': '',
               'directory': '', 'geometry': '.vtk'}[param.nodeName]
    return base + ext

if __name__ == "__main__":
    # NOTE:  For now either the launcher needs to be found on the default path, or
    # every tool in the modules list must be found on the default path
    # AND calling the module with --xml must be supported and compliant.
    modules_list = [#'MedianImageFilter',
                    #'CheckerBoardFilter',
                    #'EMSegmentCommandLine',
                    #'GrayscaleFillHoleImageFilter',
                    # 'CreateDICOMSeries', #missing channel
                    #'TractographyLabelMapSeeding',
                    #'IntensityDifferenceMetric',
                    #'DWIToDTIEstimation',
                    #'MaskScalarVolume',
                    #'ImageLabelCombine',
                    #'DTIimport',
                    #'OtsuThresholdImageFilter',
                    #'ExpertAutomatedRegistration',
                    #'ThresholdScalarVolume',
                    #'DWIUnbiasedNonLocalMeansFilter',
                    #'BRAINSFit',
                    #'MergeModels',
                    #'ResampleDTIVolume',
                    #'MultiplyScalarVolumes',
                    #'LabelMapSmoothing',
                    #'RigidRegistration',
                    #'VotingBinaryHoleFillingImageFilter',
                    #'BRAINSROIAuto',
                    #'RobustStatisticsSegmenter',
                    #'GradientAnisotropicDiffusion',
                    #'ProbeVolumeWithModel',
                    #'ModelMaker',
                    #'ExtractSkeleton',
                    #'GrayscaleGrindPeakImageFilter',
                    #'N4ITKBiasFieldCorrection',
                    #'BRAINSResample',
                    #'DTIexport',
                    #'VBRAINSDemonWarp',
                    #'ResampleScalarVectorDWIVolume',
                    #'ResampleScalarVolume',
                    #'OtsuThresholdSegmentation',
                    #'ExecutionModelTour',
                    #'HistogramMatching',
                    #'BRAINSDemonWarp',
                    #'ModelToLabelMap',
                    #'GaussianBlurImageFilter',
                    #'DiffusionWeightedVolumeMasking',
                    #'GrayscaleModelMaker',
                    #'CastScalarVolume',
                    #'DicomToNrrdConverter',
                    #'AffineRegistration',
                    #'AddScalarVolumes',
                    #'LinearRegistration',
                    #'SimpleRegionGrowingSegmentation',
                    #'DWIJointRicianLMMSEFilter',
                    #'MultiResolutionAffineRegistration',
                    #'SubtractScalarVolumes',
                    #'DWIRicianLMMSEFilter',
                    #'OrientScalarVolume',
                    #'FiducialRegistration',
                    #'BSplineDeformableRegistration',
                    #'CurvatureAnisotropicDiffusion',
                    #'PETStandardUptakeValueComputation',
                    #'DiffusionTensorScalarMeasurements',
                    #'ACPCTransform',
                    #'EMSegmentTransformToNewFormat',
                    'DTIPrep',
                    'DWIConvert']

    # SlicerExecutionModel compliant tools that are usually statically built, and don't need the Slicer3 --launcher
    generate_all_classes(modules_list=modules_list, launcher=[])
    # Tools compliant with SlicerExecutionModel called from the Slicer environment (for shared lib compatibility)
    # launcher = ['/home/raid3/gorgolewski/software/slicer/Slicer', '--launch']
    # generate_all_classes(modules_list=modules_list, launcher=launcher)
    # generate_all_classes(modules_list=['BRAINSABC'], launcher=[] )