runsnake.py

#!/bin/env python

"""
Run Snakemake targets directly using the same environment SMRTSV would set up.
"""

import os
import re
import sys

from smrtsvlib import smrtsvrunner


class _EmtpyArguments:
    """
    An object class for storing attributes. Takes the place of the object generated by parsing with `argparse`.
    """

    def __init__(self):
        self.distribute = False
        self.dryrun = False
        self.drmaalib = None
        self.verbose = False
        self.cluster_config = None


if __name__ == '__main__':

    # Get SMRTSV base directory
    smrtsv_dir = os.path.dirname(os.path.abspath(__file__))

    # Modify environment and get a copy
    process_env = smrtsvrunner.get_env(smrtsv_dir)

    # Setup arguments
    args = _EmtpyArguments()

    # Check arguments
    if len(sys.argv) < 2:
        raise RuntimeError('Need at least one argument specifying the snakefile file')

    # Get target snakefile
    snakefile = sys.argv[1].strip()

    snake_command = sys.argv[2:]  # Strip snakefile name from command

    if snakefile == 'Snakefile':
        pass  # Do no further interpretation

    elif re.match('[A-Za-z_]+$', snakefile):
        snakefile_rules = os.path.join(smrtsv_dir, 'rules', snakefile + '.snakefile')

        if os.path.isfile(snakefile_rules):
            snakefile = snakefile_rules

    elif not os.path.isfile(snakefile):
        raise RuntimeError('Cannot locate snakefile: ' + snakefile)

    # Get additional arguments to this command (remaining arguments are for snakemake)
    while len(snake_command) > 0:

        next_arg = snake_command[0]

        if next_arg == '--distribute':
            args.distribute = True
            snake_command = snake_command[1:]

        elif next_arg == '--drmaalib':
            if len(snake_command) < 2:
                raise RuntimeError('Option --drmaalib is missing a value')

            args.drmaalib = snake_command[1]
            snake_command = snake_command[2:]

        elif next_arg == '--verbose':
            args.verbose = True
            snake_command = snake_command[1:]

        elif next_arg == '--cluster-config':
            args.cluster_config = snake_command[1]
            snake_command = snake_command[2:]

        elif next_arg == '--':
            snake_command = snake_command[1:]
            break

        else:
            break

    # Set DRMAA library path
    if args.drmaalib is not None:
        process_env['DRMAA_LIBRARY_PATH'] = args.drmaalib

    elif args.distribute and 'DRMAA_LIBRARY_PATH' not in process_env:
        sys.stderr.write(
            'WARNING: --distribute is set, but DRMAA_LIBRARY_PATH is not set in the environment or via the '
            '--drmaalib option: Searching only in Python\'s library path for libdrmaa.so\n'
        )

    # Output verbose information
    if args.verbose:

        # Print python version
        print('Python version: {0}'.format(re.sub('\s*\n\s*', ' - ', sys.version)))

        # Print environment
        print('PATH:')
        for PATH_ELEMENT in process_env['PATH'].split(':'):
            print('\t* {}'.format(PATH_ELEMENT))

        print('LD_LIBRARY_PATH:')
        for PATH_ELEMENT in process_env['LD_LIBRARY_PATH'].split(':'):
            print('\t* {}'.format(PATH_ELEMENT))

        if 'DRMAA_LIBRARY_PATH' in process_env:
            print('DRMAA_LIBRARY_PATH: {}'.format(process_env['DRMAA_LIBRARY_PATH']))
        else:
            print('DRMAA_LIBRARY_PATH: <NOT_SET>\n\t* Not required unless --distribute is set')

        # Print snake command
        print('Snake command:')
        for command_element in snake_command:
            if command_element is not None:
                print('\t* {}'.format(command_element))
            else:
                print('\t- None')

        # Flush output
        sys.stdout.flush()

    # Make log directory for distributed jobs
    if args.distribute:
        os.makedirs('log', exist_ok=True)

    # Run snake target
    smrtsvrunner.run_snake_target(
        snakefile, args, process_env, smrtsv_dir, snake_command
    )