# -*- coding: utf-8 -*-
"""
The skeleton of loon package.
Set command in [options.entry_points] section of setup.cfg file
Install with `python setup.py install`
"""
import os
import sys
import argparse
import logging
if __package__ == '' or __package__ is None: # Use for test
from __init__ import __version__, __author__, __license__
from classes import Host, PBS
from tool import batch
else:
from loon import __version__, __author__, __license__
from loon.classes import Host, PBS
from loon.tool import batch
_logger = logging.getLogger(__name__)
[docs]def parse_args(args):
"""Parse command line parameters
Args:
args ([str]): command line parameters as list of strings
Returns:
:obj:`argparse.Namespace`: command line parameters namespace
"""
parser = argparse.ArgumentParser(description="Be an efficient loon.")
# Show version info
parser.add_argument(
"-V",
"--version",
action="version",
version="loon {ver} released under {license} license.".format(
ver=__version__, license=__license__.upper()))
parser.add_argument(
"--author",
action="version",
help="show info of program's author",
version="Author: 王诗翔 Email: w_shixiang@163.com GitHub: @{author}".
format(author=__author__))
# Common arguments for host
host_parent_parser = argparse.ArgumentParser(add_help=False)
host_parent_parser.add_argument(
'-N',
'--name',
dest="name",
help='Host alias, default is value from -U',
type=str)
host_parent_parser.add_argument('-U',
'--username',
dest='username',
help='Username for remote host',
type=str)
host_parent_parser.add_argument(
'-H',
'--host',
dest='host',
help='IP address for remote host (e.g. 192.168.0.1)',
type=str)
host_parent_parser.add_argument('-P',
'--port',
dest="port",
help='Port for remote host, default is 22',
type=int,
default=22)
# Common arguments for all commands
verbose_parser = argparse.ArgumentParser(add_help=False)
verbose_parser.add_argument("-v",
"--verbose",
dest="loglevel",
help="set loglevel to INFO",
action="store_const",
const=logging.INFO)
verbose_parser.add_argument('--dry',
help="Dry run the commands",
action='store_true')
# Subcommands
subparsers = parser.add_subparsers(
title='subcommands',
#description='valid subcommands',
help="description",
dest="subparsers_name")
# Create the parser for the "add" command
parser_add = subparsers.add_parser(
'add',
help="Add a remote host",
parents=[host_parent_parser, verbose_parser])
parser_add.add_argument('-A',
'--active',
dest='switch_active',
help='Set new host as active host',
action='store_true')
# Create the parser for the "add" command
parser_del = subparsers.add_parser(
'delete',
help="Delete a remote host",
parents=[host_parent_parser, verbose_parser])
# Create the parser for the "switch" command
parser_switch = subparsers.add_parser(
'switch',
help="Switch active remote host",
parents=[host_parent_parser, verbose_parser])
# Create the parser for the "list" command
parser_list = subparsers.add_parser('list',
help="List all remote hosts",
parents=[verbose_parser])
# Create the parser for the "rename" command
parser_rename = subparsers.add_parser('rename',
help="Rename host alias",
parents=[verbose_parser])
parser_rename.add_argument('old', help="Old host alias", type=str)
parser_rename.add_argument('new', help="New host alias", type=str)
# Create the parser for the "run" command
parser_run = subparsers.add_parser(
'run',
help='Run commands or scripts on remote',
parents=[verbose_parser])
parser_run.add_argument(
nargs='+',
dest='commands',
help=
"Commands/scripts to run, special symbol or option should be quoted, e.g. 'ls -l ~', 'ls -l'"
)
parser_run.add_argument('-f',
'--file',
dest='run_file',
help='Run scripts instead of commands',
action='store_true')
parser_run.add_argument(
'--data',
help=('Include a data directory when run local scripts'),
type=str,
required=False)
parser_run.add_argument(
'--remote',
dest='remote_file',
help='Scripts are directly from the active remote host',
action='store_true')
parser_run.add_argument(
'--dir',
help=
'Remote directory for storing local scripts. Only used when flag --file sets and --remote does not set. Default is /tmp',
default='/tmp')
parser_run.add_argument(
'--prog',
help=
'Specified program to run scripts, if not set, scripts will be executed directly assuming shbang exist',
required=False)
# Create the parser for the "upload" command
parser_upload = subparsers.add_parser(
'upload',
help='Upload files to active remote host',
parents=[verbose_parser])
parser_upload.add_argument('source',
nargs='+',
help='Source files to upload')
parser_upload.add_argument('destination',
help="Remote destination directory",
type=str)
parser_upload.add_argument('--rsync',
help="Use rsync instead of scp",
action='store_true')
# Create the parser for the "download" command
parser_download = subparsers.add_parser(
'download',
help='Download files from active remote host',
parents=[verbose_parser])
parser_download.add_argument('source',
nargs='+',
help='Source files to download')
parser_download.add_argument(
'destination',
help=
"Local destination directory, note '~' should be quoted in some cases",
type=str)
parser_download.add_argument('--rsync',
help="Use rsync instead of scp",
action='store_true')
# Create the parser for the "gen" command
parser_gen = subparsers.add_parser(
'gen',
help='Generate a batch of (script) files',
parents=[verbose_parser])
parser_gen.add_argument('-t',
'--template',
help="A template file containing placeholders")
parser_gen.add_argument(
'-s',
'--samplefile',
help=
"A csv file containing unique filenames (the first column) and replacing labels"
)
parser_gen.add_argument(
'-m',
'--mapfile',
help=
"A csv file containing placeholders and column index (0-based) indicating replacing labels in samplefile"
)
parser_gen.add_argument('-o', '--output', help="Output directory")
# Create the parser for the "batch" command
parser_batch = subparsers.add_parser(
'batch',
help="Batch process commands with placeholders",
parents=[verbose_parser])
parser_batch.add_argument(
'-f',
'--file',
help=
r'A structed file/stdin like CSV, TSV etc. Each column is placeholder target, i.e. {0} targets the first column, {1} targets the second column, etc.',
type=argparse.FileType('r', encoding='UTF-8'),
default=sys.stdin,
nargs='?')
# parser_batch.add_argument(
# '-o',
# '--output',
# help=
# r'default is stdout, can be a file',
# type=argparse.FileType('w'),
# default=sys.stdout,
# nargs='?')
parser_batch.add_argument(
'-s',
'--sep',
help=r"File separator, ',' for CSV (default) and '\t' for TSV",
default=',',
required=False)
parser_batch.add_argument('-T',
'--thread',
help="Thread number, default is 1",
required=False,
default=1,
type=int)
parser_batch.add_argument('--header',
help="Set it if input file contains header",
action='store_true')
parser_batch.add_argument('cmds',
type=str,
help="A sample command with placeholders")
# Create the parser for the "pbstemp" command
parser_pbstemp = subparsers.add_parser('pbstemp',
help='Generate a PBS template file',
parents=[verbose_parser])
parser_pbstemp.add_argument(
'-i',
'--input',
help='A template file, if not set, a default template is used',
type=str,
required=False)
parser_pbstemp.add_argument('-o',
'--output',
help="Output file, default is work.pbs",
type=str,
required=False)
# Create the parser for the "pbsgen" command
parser_pbsgen = subparsers.add_parser(
'pbsgen',
help='Generate a batch of PBS files (with .pbs extension)',
parents=[verbose_parser])
parser_pbsgen.add_argument(
'-t', '--template', help="A PBS template file containing placeholders")
parser_pbsgen.add_argument(
'-s',
'--samplefile',
help=
"A csv file containing unique filenames (the first column) and replacing labels"
)
parser_pbsgen.add_argument(
'-m',
'--mapfile',
help=
"A csv file containing placeholders and column index (0-based) indicating replacing labels in samplefile"
)
parser_pbsgen.add_argument('-o', '--output', help="Output directory")
# Create the parser for the "pbsgen_example" command
parser_genexample = subparsers.add_parser(
'pbsgen_example',
help='Generate example files for pbsgen command',
parents=[verbose_parser])
parser_genexample.add_argument('output', help='Output directory')
# Create the parser for the "pbssub" command
parser_pbssub = subparsers.add_parser('pbssub',
help='Submit PBS tasks',
parents=[verbose_parser])
parser_pbssub.add_argument(
'--remote',
dest='remote_file',
help='PBS task files are located at the active remote host',
action='store_true')
parser_pbssub.add_argument(
'--workdir',
help=
'Working directory, default is /tmp for remote host and otherwise the command executed path',
required=False)
parser_pbssub.add_argument(
nargs='+',
dest='tasks',
help="Tasks to submit, can be a directory containing only PBS files")
# Create the parser for the "pbsdeploy" command
parser_deploy = subparsers.add_parser(
'pbsdeploy',
help='Deploy target destination to remote host',
parents=[verbose_parser])
parser_deploy.add_argument(
'target', help='Target directory containing PBS files and more')
parser_deploy.add_argument(
'destination',
help=
"Local destination directory, note '~' should be quoted in some cases",
nargs='?')
parser_deploy.add_argument('--rsync',
help="Use rsync instead of scp",
action='store_true')
# Create the parser for the "pbscheck" command
parser_pbscheck = subparsers.add_parser(
'pbscheck',
help='Check status of PBS job on remote host',
parents=[verbose_parser])
parser_pbscheck.add_argument(
'job_id',
help="ID of job, if not set, all running jobs will be returned",
type=str,
nargs='?')
return parser.parse_args(args), parser
[docs]def setup_logging(loglevel):
"""Setup basic logging
Args:
loglevel (int): minimum loglevel for emitting messages
"""
logformat = "[%(asctime)s] %(levelname)s:%(name)s:%(message)s"
logging.basicConfig(level=loglevel,
stream=sys.stdout,
format=logformat,
datefmt="%Y-%m-%d %H:%M:%S")
[docs]def main(args):
"""Main entry point allowing external calls
Args:
args ([str]): command line parameter list
"""
args_bk = args
args, parser = parse_args(args)
if len(args_bk) == 0:
parser.print_help()
parser.exit()
setup_logging(args.loglevel)
_logger.info("Starting loon...")
host = Host()
pbs = PBS()
if hasattr(args, 'rsync') and args.rsync:
use_rsync = True
else:
use_rsync = False
# Deparse arguments
if args.subparsers_name == 'add':
_logger.info("Add command is detected.")
if args.username is None or args.host is None:
print("Error: username and host are both required in add command.")
sys.exit(1)
if args.name is None:
args.name = args.username
host.add(name=args.name,
username=args.username,
host=args.host,
port=args.port,
dry_run=args.dry)
if args.switch_active:
host.switch(name=args.name,
username=args.username,
host=args.host,
port=args.port,
dry_run=args.dry)
elif args.subparsers_name == 'delete':
_logger.info("Delete command is detected.")
if args.username is None or args.host is None:
if args.name is None:
print("Error: either specify name or both username and host")
sys.exit(1)
host.delete(name=args.name,
username=args.username,
host=args.host,
port=args.port,
dry_run=args.dry)
elif args.subparsers_name == 'switch':
_logger.info("Switch command is detected.")
if args.username is None or args.host is None:
if args.name is None:
print("Error: either specify name or both username and host")
sys.exit(1)
host.switch(name=args.name,
username=args.username,
host=args.host,
port=args.port,
dry_run=args.dry)
elif args.subparsers_name == 'list':
_logger.info("List command is detected.")
host.list()
elif args.subparsers_name == 'rename':
_logger.info("Rename command is detected.")
host.rename(args.old, args.new, dry_run=args.dry)
elif args.subparsers_name == 'run':
_logger.info("Run command is detected.")
if args.run_file:
commands = args.commands
else:
commands = " ".join(args.commands)
host.cmd(commands,
_logger=_logger,
run_file=args.run_file,
data_dir=args.data,
remote_file=args.remote_file,
dir=args.dir,
prog=args.prog,
dry_run=args.dry)
elif args.subparsers_name == 'upload':
_logger.info("Upload command is detected.")
#host.connect(open_channel=False)
host.upload(args.source,
args.destination,
_logger=_logger,
use_rsync=use_rsync,
dry_run=args.dry)
elif args.subparsers_name == 'download':
_logger.info("Download command is detected.")
#host.connect(open_channel=False)
host.download(args.source,
args.destination,
_logger=_logger,
use_rsync=use_rsync,
dry_run=args.dry)
elif args.subparsers_name == 'batch':
_logger.info("Batch command is detected.")
batch(args.file,
args.cmds,
sep=args.sep,
thread=args.thread,
header=args.header,
dry_run=args.dry,
_logger=_logger)
elif args.subparsers_name == 'pbstemp':
_logger.info("pbstemp command is detected.")
pbs.gen_template(args.input, args.output, dry_run=args.dry)
elif args.subparsers_name == 'gen':
_logger.info("gen command is detected.")
pbs.gen_pbs(args.template,
args.samplefile,
args.mapfile,
args.output,
_logger=_logger,
pbs_mode=False,
dry_run=args.dry)
elif args.subparsers_name == 'pbsgen':
_logger.info("pbsgen command is detected.")
pbs.gen_pbs(args.template,
args.samplefile,
args.mapfile,
args.output,
_logger=_logger,
dry_run=args.dry)
elif args.subparsers_name == 'pbsgen_example':
pbs.gen_pbs_example(args.output, _logger=_logger, dry_run=args.dry)
elif args.subparsers_name == 'pbssub':
_logger.info("pbssub command is detected.")
pbs.sub(host,
args.tasks,
args.remote_file,
args.workdir,
_logger=_logger,
dry_run=args.dry)
elif args.subparsers_name == 'pbsdeploy':
_logger.info("pbsdeploy command is detected.")
pbs.deploy(host,
args.target,
args.destination,
_logger=_logger,
use_rsync=use_rsync,
dry_run=args.dry)
elif args.subparsers_name == 'pbscheck':
_logger.info("pbscheck command is detected.")
pbs.check(host, args.job_id, dry_run=args.dry)
_logger.info("loon ends here")
[docs]def run():
"""Entry point for console_scripts"""
main(sys.argv[1:])
if __name__ == "__main__":
run()