<commit_before># coding: utf-8

""" Make an SCF run given initial conditions, timestep, and potential parameters """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys
import shutil

# Third-party
from astropy.constants import G
from astropy import log as logger
import astropy.units as u

# Project
from scf import project_path
import scf.potential as sp

def main(name, pos, vel, scfpars, potentials, run_path, overwrite=False, submit=False):
    # run_path = os.path.abspath(os.path.join(project_path, "simulations", "runs"))
    template_path = os.path.abspath(os.path.join(project_path, "templates"))
    src_path = os.path.abspath(os.path.join(project_path, "src"))
    logger.debug("Run path: {}".format(run_path))

    path = os.path.join(run_path, name)
    if os.path.exists(path) and overwrite:
        logger.info("nuking directory...")
        shutil.rmtree(path)
    elif os.path.exists(path) and not overwrite:
        logger.warning("Path '{}' already exists -- do you want to overwrite?"
                       .format(path))
        yn = raw_input("[y/N]: ")
        if yn.lower().strip() == 'y':
            logger.info("nuking directory...")
            shutil.rmtree(path)
        else:
            logger.info("aborting...")
            sys.exit(0)

        raise IOError("Path '{}' already exists!".format(path))
        sys.exit(0)

    if not os.path.exists(path):
        os.makedirs(path)

    # parse pos and vel
    x_vals,x_unit = pos.split()
    v_vals,v_unit = vel.split()

    pos = map(float, x_vals.split(',')) * u.Unit(x_unit)
    pos = pos.to(u.kpc).value

    vel = map(float, v_vals.split(',')) * u.Unit(v_unit)
    vel = vel.to(u.km/u.s).value

    # combine all of the different blocks
    allblocks = dict()
    for potential_name in potentials:
        blocks = getattr(sp,potential_name)()
        for k,v in blocks.items():
            if k not in allblocks:
                allblocks[k] = ""

            allblocks[k] += v

            if k != "SCFPOT":
                allblocks[k] += "\n"

    # ------------------------------------------------------------------------
    # SCFPAR
    with open(os.path.join(template_path, "SCFPAR.tpl"), 'r') as f:
        base_SCFPAR = f.read()

    with open(os.path.join(path, "SCFPAR"), 'w') as f:
        f.write(base_SCFPAR.format(x=pos, v=vel, **scfpars))

    # ------------------------------------------------------------------------
    # Makefile
    with open(os.path.join(template_path, "Makefile.tpl"), 'r') as f:
        base_Makefile = f.read()

    with open(os.path.join(path, "Makefile"), 'w') as f:
        f.write(base_Makefile)

    # cmd = "unexpand {fn} > {fn}2".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)
    # cmd = "mv {fn}2 {fn}".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)

    # ------------------------------------------------------------------------
    # potential.h
    with open(os.path.join(template_path, "potential.h.tpl"), 'r') as f:
        base_potentialh = f.read()

    with open(os.path.join(path, "potential.h"), 'w') as f:
        f.write(base_potentialh.format(hblock=allblocks['hblock']))

    # ------------------------------------------------------------------------
    # potential.f
    with open(os.path.join(template_path, "potential.f.tpl"), 'r') as f:
        base_potentialf = f.read()

    with open(os.path.join(path, "potential.f"), 'w') as f:
        f.write(base_potentialf.format(**allblocks))

    # ------------------------------------------------------------------------
    # SCFPOT
    with open(os.path.join(path, "SCFPOT"), 'w') as f:
        f.write(allblocks['SCFPOT'])

    if submit:
        with open(os.path.join(template_path, "submit.tpl"), 'r') as f:
            base_submit = f.read()

        with open(os.path.join(path, "submit.sh"), 'w') as f:
            f.write(base_submit.format(name=name, path=path))

    # Copy scf.f and scf.h
    shutil.copyfile(os.path.join(src_path, 'scf.f'),
                    os.path.join(path, 'scf.f'))
    shutil.copyfile(os.path.join(src_path, 'scf.h'),
                    os.path.join(path, 'scf.h'))

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY.")

    parser.add_argument("--path", "-p", dest="path", type=str, help="Path.", required=True)
    parser.add_argument("--name", dest="name", type=str, help="Name.", required=True)
    parser.add_argument("--potentials", dest="potentials", type=str,
                        help="Potential names. One of: {0}".format(",".join(sp.__all__)),
                        nargs='+', required=True)
    parser.add_argument("-s", "--submit", action="store_true", dest="submit",
                        default=True, help="Generate a submit.sh file as well.")
    parser.add_argument("--pos", dest="x", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --x='15.324,123.314,51.134 kpc'")
    parser.add_argument("--vel", dest="v", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --v='15.324,123.314,51.134 km/s'")

    # Parameters for SCF
    parser.add_argument("--dt", dest="dt", default=1., type=float,
                        help="Time step in Myr.")
    parser.add_argument("--nsteps", dest="nsteps", type=int,
                        help="Number of steps.", required=True)
    parser.add_argument("--ncen", dest="ncen", type=int, default=10,
                        help="Output to SCFCEN every (this number) steps.")
    parser.add_argument("--nsnap", dest="nsnap", type=int, default=1000,
                        help="Output to a SNAP file every (this number) steps.")
    parser.add_argument("--ntide", dest="ntide", default=500, type=int,
                        help="Number of steps over which to turn on the tidal field.")
    parser.add_argument("--mass", dest="mass", type=float, required=True,
                        help="Mass in solar masses within 35 scale radii of the "
                             "cluster/satellite.")
    parser.add_argument("--rscale", dest="rscale", default=None, type=float,
                        help="Scale radius of the cluster/satellite.")
    parser.add_argument("--SCFBI", dest="scfbi_path", default=None, type=str,
                        help="Path to an SCFBI file.")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    # Contains user-specified parameters for SCF
    scfpars = dict()

    if args.rscale is None:
        ru = 0.43089*(args.mass/2.5e9)**(1/3.)
    else:
        ru = args.rscale
    scfpars['rscale'] = ru
    scfpars['mass'] = args.mass

    _G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
    X = (_G / ru**3 * args.mass)**-0.5

    length_unit = u.Unit("{0} kpc".format(ru))
    mass_unit = u.Unit("{0} M_sun".format(args.mass))
    time_unit = u.Unit("{:08f} Myr".format(X))

    scfpars['dt'] = args.dt / (1*time_unit).to(u.Myr).value
    scfpars['nsteps'] = args.nsteps
    scfpars['ncen'] = args.ncen
    scfpars['nsnap'] = args.nsnap
    scfpars['ntide'] = args.ntide
    if args.scfbi_path is None:
        scfpars['SCFBIpath'] = os.path.abspath(os.path.join(project_path, 'src', 'SCFBI'))

    main(name=args.name, pos=args.x, vel=args.v, scfpars=scfpars,
         potentials=args.potentials, overwrite=args.overwrite,
         submit=args.submit, path=args.path)
<commit_msg>use abspath<commit_after># coding: utf-8

""" Make an SCF run given initial conditions, timestep, and potential parameters """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys
import shutil

# Third-party
from astropy.constants import G
from astropy import log as logger
import astropy.units as u

# Project
from scf import project_path
import scf.potential as sp

def main(name, pos, vel, scfpars, potentials, run_path, overwrite=False, submit=False):
    run_path = os.path.abspath(run_path)
    # run_path = os.path.abspath(os.path.join(project_path, "simulations", "runs"))
    template_path = os.path.abspath(os.path.join(project_path, "templates"))
    src_path = os.path.abspath(os.path.join(project_path, "src"))
    logger.debug("Run path: {}".format(run_path))

    path = os.path.join(run_path, name)
    if os.path.exists(path) and overwrite:
        logger.info("nuking directory...")
        shutil.rmtree(path)
    elif os.path.exists(path) and not overwrite:
        logger.warning("Path '{}' already exists -- do you want to overwrite?"
                       .format(path))
        yn = raw_input("[y/N]: ")
        if yn.lower().strip() == 'y':
            logger.info("nuking directory...")
            shutil.rmtree(path)
        else:
            logger.info("aborting...")
            sys.exit(0)

        raise IOError("Path '{}' already exists!".format(path))
        sys.exit(0)

    if not os.path.exists(path):
        os.makedirs(path)

    # parse pos and vel
    x_vals,x_unit = pos.split()
    v_vals,v_unit = vel.split()

    pos = map(float, x_vals.split(',')) * u.Unit(x_unit)
    pos = pos.to(u.kpc).value

    vel = map(float, v_vals.split(',')) * u.Unit(v_unit)
    vel = vel.to(u.km/u.s).value

    # combine all of the different blocks
    allblocks = dict()
    for potential_name in potentials:
        blocks = getattr(sp,potential_name)()
        for k,v in blocks.items():
            if k not in allblocks:
                allblocks[k] = ""

            allblocks[k] += v

            if k != "SCFPOT":
                allblocks[k] += "\n"

    # ------------------------------------------------------------------------
    # SCFPAR
    with open(os.path.join(template_path, "SCFPAR.tpl"), 'r') as f:
        base_SCFPAR = f.read()

    with open(os.path.join(path, "SCFPAR"), 'w') as f:
        f.write(base_SCFPAR.format(x=pos, v=vel, **scfpars))

    # ------------------------------------------------------------------------
    # Makefile
    with open(os.path.join(template_path, "Makefile.tpl"), 'r') as f:
        base_Makefile = f.read()

    with open(os.path.join(path, "Makefile"), 'w') as f:
        f.write(base_Makefile)

    # cmd = "unexpand {fn} > {fn}2".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)
    # cmd = "mv {fn}2 {fn}".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)

    # ------------------------------------------------------------------------
    # potential.h
    with open(os.path.join(template_path, "potential.h.tpl"), 'r') as f:
        base_potentialh = f.read()

    with open(os.path.join(path, "potential.h"), 'w') as f:
        f.write(base_potentialh.format(hblock=allblocks['hblock']))

    # ------------------------------------------------------------------------
    # potential.f
    with open(os.path.join(template_path, "potential.f.tpl"), 'r') as f:
        base_potentialf = f.read()

    with open(os.path.join(path, "potential.f"), 'w') as f:
        f.write(base_potentialf.format(**allblocks))

    # ------------------------------------------------------------------------
    # SCFPOT
    with open(os.path.join(path, "SCFPOT"), 'w') as f:
        f.write(allblocks['SCFPOT'])

    if submit:
        with open(os.path.join(template_path, "submit.tpl"), 'r') as f:
            base_submit = f.read()

        with open(os.path.join(path, "submit.sh"), 'w') as f:
            f.write(base_submit.format(name=name, path=path))

    # Copy scf.f and scf.h
    shutil.copyfile(os.path.join(src_path, 'scf.f'),
                    os.path.join(path, 'scf.f'))
    shutil.copyfile(os.path.join(src_path, 'scf.h'),
                    os.path.join(path, 'scf.h'))

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY.")

    parser.add_argument("--path", "-p", dest="path", type=str, help="Path.", required=True)
    parser.add_argument("--name", dest="name", type=str, help="Name.", required=True)
    parser.add_argument("--potentials", dest="potentials", type=str,
                        help="Potential names. One of: {0}".format(",".join(sp.__all__)),
                        nargs='+', required=True)
    parser.add_argument("-s", "--submit", action="store_true", dest="submit",
                        default=True, help="Generate a submit.sh file as well.")
    parser.add_argument("--pos", dest="x", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --x='15.324,123.314,51.134 kpc'")
    parser.add_argument("--vel", dest="v", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --v='15.324,123.314,51.134 km/s'")

    # Parameters for SCF
    parser.add_argument("--dt", dest="dt", default=1., type=float,
                        help="Time step in Myr.")
    parser.add_argument("--nsteps", dest="nsteps", type=int,
                        help="Number of steps.", required=True)
    parser.add_argument("--ncen", dest="ncen", type=int, default=10,
                        help="Output to SCFCEN every (this number) steps.")
    parser.add_argument("--nsnap", dest="nsnap", type=int, default=1000,
                        help="Output to a SNAP file every (this number) steps.")
    parser.add_argument("--ntide", dest="ntide", default=500, type=int,
                        help="Number of steps over which to turn on the tidal field.")
    parser.add_argument("--mass", dest="mass", type=float, required=True,
                        help="Mass in solar masses within 35 scale radii of the "
                             "cluster/satellite.")
    parser.add_argument("--rscale", dest="rscale", default=None, type=float,
                        help="Scale radius of the cluster/satellite.")
    parser.add_argument("--SCFBI", dest="scfbi_path", default=None, type=str,
                        help="Path to an SCFBI file.")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    # Contains user-specified parameters for SCF
    scfpars = dict()

    if args.rscale is None:
        ru = 0.43089*(args.mass/2.5e9)**(1/3.)
    else:
        ru = args.rscale
    scfpars['rscale'] = ru
    scfpars['mass'] = args.mass

    _G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
    X = (_G / ru**3 * args.mass)**-0.5

    length_unit = u.Unit("{0} kpc".format(ru))
    mass_unit = u.Unit("{0} M_sun".format(args.mass))
    time_unit = u.Unit("{:08f} Myr".format(X))

    scfpars['dt'] = args.dt / (1*time_unit).to(u.Myr).value
    scfpars['nsteps'] = args.nsteps
    scfpars['ncen'] = args.ncen
    scfpars['nsnap'] = args.nsnap
    scfpars['ntide'] = args.ntide
    if args.scfbi_path is None:
        scfpars['SCFBIpath'] = os.path.abspath(os.path.join(project_path, 'src', 'SCFBI'))

    main(name=args.name, pos=args.x, vel=args.v, scfpars=scfpars,
         potentials=args.potentials, overwrite=args.overwrite,
         submit=args.submit, run_path=args.path)
<|endoftext|>


<commit_before>import fnmatch
import hexchat

__module_name__ = 'BanSearch'
__module_author__ = 'TingPing'
__module_version__ = '1'
__module_description__ = 'Search for bans/quiets matching a user'

banhook = 0
quiethook = 0
endbanhook = 0
endquiethook = 0

banlist = []
quietlist = []

def print_result(mask, matchlist, matchnum, _type):
	if matchnum:
		print('\00318{}\017 had \00320{}\017 {} matches:'.format(mask, matchnum, _type))
		for match in matchlist:
			print('\t\t\t{}'.format(match))
	else:
		print('No {} matches for {} were found.'.format(_type, mask))

def match_mask(mask, searchmask):
	# The fnmatch function works for * and ?, to be replaced by a more irc friendly version...
	return fnmatch.fnmatch(mask, searchmask)

def match_extban(mask, host, account, realname, usermask):
	extban, banmask = mask.split(':')

	if '~' in extban:
		invert = True
	else:
		invert = False

	# From http://freenode.net/using_the_network.shtml
	if 'a' in extban:
		ret = banmask == account
	elif 'r' in extban:
		ret = banmask == realname
	elif 'x' in extban:
		ret = match_mask ('{}#{}'.format(host, realname), banmask)
	else:
		return False

	if invert:
		return not ret
	else:
		return ret

def get_user_info(nick):
	invalid_chars = ['*', '?', '$', '@', '!']
	if any(char in nick for char in invalid_chars):
		return (None, None, None) # It's a mask not a nick.

	for user in hexchat.get_list('users'):
		if user.nick == nick:
			host = user.host
			account = user.account
			realname = user.realname
			return (host, account, realname)

def banlist_cb(word, word_eol, userdata):
	global banlist
	banlist.append(word[4])
	return hexchat.EAT_HEXCHAT

def endbanlist_cb(word, word_eol, usermask):
	global banhook
	global endbanhook
	global banlist
	matchlist = []
	matchnum = 0

	# Cleanup hooks
	hexchat.unhook(banhook)
	banhook = 0
	hexchat.unhook(endbanhook)
	endbanhook = 0

	host, account, realname = get_user_info (usermask)

	if banlist:
		for mask in banlist:
			if mask[0] == '$' and host: # If extban we need user info
				if match_extban (mask, host, account, realname, usermask):
					matchlist.append(mask)
					matchnum = matchnum + 1
			elif match_mask (mask, usermask):
				matchlist.append(mask)
				matchnum = matchnum + 1

		banlist = []

	print_result (usermask, matchlist, matchnum, 'Ban')

	return hexchat.EAT_HEXCHAT

def quietlist_cb(word, word_eol, userdata):
	global quietlist
	quietlist.append(word[5])
	return hexchat.EAT_HEXCHAT

def endquietlist_cb(word, word_eol, usermask):
	global quiethook
	global endquiethook
	global quietlist
	matchlist = []
	matchnum = 0

	# Cleanup hooks
	hexchat.unhook(quiethook)
	quiethook = 0
	hexchat.unhook(endquiethook)
	endquiethook = 0

	host, account, realname = get_user_info (usermask)

	if quietlist:
		for mask in quietlist:
			if mask[0] == '$' and host: # If extban we need user info
				if match_extban (mask, host, account, realname, usermask):
					matchlist.append(mask)
					matchnum = matchnum + 1
			elif match_mask (mask, usermask):
				matchlist.append(mask)
				matchnum = matchnum + 1

		quietlist = []

	print_result (usermask, matchlist, matchnum, 'Quiet')

	return hexchat.EAT_HEXCHAT

def search_cb(word, word_eol, userdata):
	global banhook
	global quiethook
	global endbanhook
	global endquiethook

	if len(word) == 2:
		if not banhook and not quiethook:
			banhook = hexchat.hook_server ('367', banlist_cb)
			quiethook = hexchat.hook_server ('728', quietlist_cb)
			endbanhook = hexchat.hook_server ('368', endbanlist_cb, word[1])
			endquiethook = hexchat.hook_server ('729', endquietlist_cb, word[1])

			hexchat.command('ban')
			hexchat.command('quiet')
		else:
			print('A ban search is already in progress.')

	else:
		hexchat.command('help bansearch')

	return hexchat.EAT_ALL

hexchat.hook_command('bansearch', search_cb, help='BANSEARCH <mask|nick>')
hexchat.prnt(__module_name__ + ' version ' + __module_version__ + ' loaded.')
<commit_msg>Update bansearch.py<commit_after>import fnmatch
import hexchat

__module_name__ = 'BanSearch'
__module_author__ = 'TingPing'
__module_version__ = '1'
__module_description__ = 'Search for bans/quiets matching a user'

banhook = 0
quiethook = 0
endbanhook = 0
endquiethook = 0

banlist = []
quietlist = []

def print_result(mask, matchlist, matchnum, _type):
	if matchnum:
		print('\00318{}\017 had \00320{}\017 {} matches:'.format(mask, matchnum, _type))
		for match in matchlist:
			print('\t\t\t{}'.format(match))
	else:
		print('No {} matches for \00318{}\017 were found.'.format(_type, mask))

def match_mask(mask, searchmask):
	# The fnmatch function works for * and ?, to be replaced by a more irc friendly version...
	return fnmatch.fnmatch(mask, searchmask)

def match_extban(mask, host, account, realname, usermask):
	extban, banmask = mask.split(':')

	if '~' in extban:
		invert = True
	else:
		invert = False

	# From http://freenode.net/using_the_network.shtml
	if 'a' in extban:
		ret = banmask == account
	elif 'r' in extban:
		ret = banmask == realname
	elif 'x' in extban:
		ret = match_mask ('{}#{}'.format(host, realname), banmask)
	else:
		return False

	if invert:
		return not ret
	else:
		return ret

def get_user_info(nick):
	invalid_chars = ['*', '?', '$', '@', '!']
	if any(char in nick for char in invalid_chars):
		return (None, None, None) # It's a mask not a nick.

	for user in hexchat.get_list('users'):
		if user.nick == nick:
			host = user.host
			account = user.account
			realname = user.realname
			return (host, account, realname)

def banlist_cb(word, word_eol, userdata):
	global banlist
	banlist.append(word[4])
	return hexchat.EAT_HEXCHAT

def endbanlist_cb(word, word_eol, usermask):
	global banhook
	global endbanhook
	global banlist
	matchlist = []
	matchnum = 0

	# Cleanup hooks
	hexchat.unhook(banhook)
	banhook = 0
	hexchat.unhook(endbanhook)
	endbanhook = 0

	host, account, realname = get_user_info (usermask)

	if banlist:
		for mask in banlist:
			if mask[0] == '$' and host: # If extban we need user info
				if match_extban (mask, host, account, realname, usermask):
					matchlist.append(mask)
					matchnum = matchnum + 1
			elif match_mask (mask, usermask):
				matchlist.append(mask)
				matchnum = matchnum + 1

		banlist = []

	print_result (usermask, matchlist, matchnum, 'Ban')

	return hexchat.EAT_HEXCHAT

def quietlist_cb(word, word_eol, userdata):
	global quietlist
	quietlist.append(word[5])
	return hexchat.EAT_HEXCHAT

def endquietlist_cb(word, word_eol, usermask):
	global quiethook
	global endquiethook
	global quietlist
	matchlist = []
	matchnum = 0

	# Cleanup hooks
	hexchat.unhook(quiethook)
	quiethook = 0
	hexchat.unhook(endquiethook)
	endquiethook = 0

	host, account, realname = get_user_info (usermask)

	if quietlist:
		for mask in quietlist:
			if mask[0] == '$' and host: # If extban we need user info
				if match_extban (mask, host, account, realname, usermask):
					matchlist.append(mask)
					matchnum = matchnum + 1
			elif match_mask (mask, usermask):
				matchlist.append(mask)
				matchnum = matchnum + 1

		quietlist = []

	print_result (usermask, matchlist, matchnum, 'Quiet')

	return hexchat.EAT_HEXCHAT

def search_cb(word, word_eol, userdata):
	global banhook
	global quiethook
	global endbanhook
	global endquiethook

	if len(word) == 2:
		if not banhook and not quiethook:
			banhook = hexchat.hook_server ('367', banlist_cb)
			quiethook = hexchat.hook_server ('728', quietlist_cb)
			endbanhook = hexchat.hook_server ('368', endbanlist_cb, word[1])
			endquiethook = hexchat.hook_server ('729', endquietlist_cb, word[1])

			hexchat.command('ban')
			hexchat.command('quiet')
		else:
			print('A ban search is already in progress.')

	else:
		hexchat.command('help bansearch')

	return hexchat.EAT_ALL

hexchat.hook_command('bansearch', search_cb, help='BANSEARCH <mask|nick>')
hexchat.prnt(__module_name__ + ' version ' + __module_version__ + ' loaded.')
<|endoftext|>


<commit_before>"""
SkCode medias tag definitions code.
"""

from urllib.parse import urlsplit, parse_qs, quote_plus
from html import escape as escape_html

from .base import TagOptions
from ..tools import escape_attrvalue, sanitize_url

class ImageTagOptions(TagOptions):
    """ Image tag options container class. """

    # Alternate text attribute name
    alt_attr_name = 'alt'

    # Image width attribute name
    width_attr_name = 'width'

    # Image height attribute name
    height_attr_name = 'height'

    def get_image_src_link(self, tree_node):
        """
        Get the image source link URL.
        :param tree_node: The current tree node instance.
        :return The image source link URL (not sanitized).
        """
        return tree_node.get_raw_content()

    def get_alt_text(self, tree_node):
        """
        Get the image alt text.
        :param tree_node: The current tree node instance.
        :return The image alternative text as string, or an empty string.
        """
        return tree_node.attrs.get(self.alt_attr_name, '')

    def get_img_width(self, tree_node):
        """
        Get the image width, or zero.
        :param tree_node: The current tree node instance.
        :return The image width, or zero.
        """
        try:
            width = int(tree_node.attrs.get(self.width_attr_name, 0))
            if width <= 0:
                return 0
            return width
        except ValueError:
            return 0

    def get_img_height(self, tree_node):
        """
        Get the image height, or zero.
        :param tree_node: The current tree node instance.
        :return The image height, or zero.
        """
        try:
            height = int(tree_node.attrs.get(self.height_attr_name, 0))
            if height <= 0:
                return 0
            return height
        except ValueError:
            return 0

    def render_html(self, tree_node, inner_html, force_rel_nofollow=True):
        """
        Callback function for rendering HTML.
        :param force_rel_nofollow: If set, all links in the rendered HTML will have "rel=nofollow" (default to True).
        :param tree_node: Current tree node to be rendered.
        :param inner_html: Inner HTML of this tree node.
        :return Rendered HTML of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)
        src_link = sanitize_url(src_link, allowed_schemes=('http', 'https'))

        # Shortcut if no source link
        if not src_link:
            return ''

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            extra_attrs = ' alt="%s"' % escape_html(alt_text)
        else:
            extra_attrs = ''

        # Get the image width
        img_width = self.get_img_width(tree_node)
        if img_width:
            extra_attrs += ' width="%d"' % img_width

        # Get the image height
        img_height = self.get_img_height(tree_node)
        if img_width:
            extra_attrs += ' height="%d"' % img_height

        # Render the image
        return '' % (src_link, extra_attrs)

    def render_text(self, tree_node, inner_text):
        """
        Callback function for rendering text.
        :param tree_node: Current tree node to be rendered.
        :param inner_text: Inner text of this tree node.
        :return Rendered text of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)
        src_link = sanitize_url(src_link, allowed_schemes=('http', 'https'))

        # Shortcut if no source link
        if not src_link:
            return inner_text

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            return '%s (%s)' % (src_link, alt_text)
        else:
            return src_link

    def render_skcode(self, tree_node, inner_skcode):
        """
        Callback function for rendering SkCode.
        :param tree_node: Current tree node to be rendered.
        :param inner_skcode: Inner SkCode of this tree node.
        :return Rendered SkCode of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)
        src_link = sanitize_url(src_link, allowed_schemes=('http', 'https'))

        # Shortcut if no source link
        if not src_link:
            return inner_skcode

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            extra_attrs = ' %s=%s' % (self.alt_attr_name, escape_attrvalue(alt_text))
        else:
            extra_attrs = ''

        # Get the image width
        img_width = self.get_img_width(tree_node)
        if img_width:
            extra_attrs += ' %s="%d"' % (self.width_attr_name, img_width)

        # Get the image height
        img_height = self.get_img_height(tree_node)
        if img_width:
            extra_attrs += ' %s="%d"' % (self.height_attr_name, img_height)

        # Render the image
        node_name = tree_node.name
        return '[%s%s]%s[/%s]' % (node_name, extra_attrs, src_link, node_name)

class YoutubeTagOptions(TagOptions):
    """ Youtube video integration tag options container class. """

    # Default iframe width
    default_iframe_width = 560

    # Default iframe height
    default_iframe_height = 315

    def get_youtube_video_id(self, tree_node):
        """
        Get the Youtube video ID.
        :param tree_node: Current tree node to be rendered.
        :return The video ID, or an empty string.
        """

        # Get the URL
        url = tree_node.get_raw_content()

        # Split the URL
        try:
            scheme, netloc, path, query, fragment = urlsplit(url)
        except ValueError:

            # Handle malformed URL
            return ''

        # Check for bad netloc
        if netloc != 'www.youtube.com' and netloc != 'youtube.com':
            return ''

        # Check for bad url query
        if not query:
            return ''

        # Get the video ID
        query_args = parse_qs(query)
        video_id = query_args.get('v', None)

        # Query args return a list of values
        if not video_id:
            return ''
        else:
            return video_id[0].strip()

    def render_html(self, tree_node, inner_html, force_rel_nofollow=True):
        """
        Callback function for rendering HTML.
        :param force_rel_nofollow: If set, all links in the rendered HTML will have "rel=nofollow" (default to True).
        :param tree_node: Current tree node to be rendered.
        :param inner_html: Inner HTML of this tree node.
        :return Rendered HTML of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the iframe
        if video_id:
            return '' % (self.default_iframe_width,
                                                                  self.default_iframe_height,
                                                                  quote_plus(video_id))
        else:
            return inner_html

    def render_text(self, tree_node, inner_text):
        """
        Callback function for rendering text.
        :param tree_node: Current tree node to be rendered.
        :param inner_text: Inner text of this tree node.
        :return Rendered text of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the video link
        if video_id:
            return 'https://www.youtube.com/watch?v=%s' % quote_plus(video_id)
        else:
            return inner_text

    def render_skcode(self, tree_node, inner_skcode):
        """
        Callback function for rendering SkCode.
        :param tree_node: Current tree node to be rendered.
        :param inner_skcode: Inner SkCode of this tree node.
        :return Rendered SkCode of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the video tag
        if video_id:
            node_name = tree_node.name
            return '[%s]https://www.youtube.com/watch?v=%s[/%s]' % (node_name, quote_plus(video_id), node_name)
        else:
            return inner_skcode
<commit_msg>Rework image tag.<commit_after>"""
SkCode medias tag definitions code.
"""

from urllib.parse import (urlsplit,
                          parse_qs,
                          quote_plus)

from html import escape as escape_html
from html import unescape as unescape_html_entities

from .base import TagOptions
from ..tools import (escape_attrvalue,
                     sanitize_url)

class ImageTagOptions(TagOptions):
    """ Image tag options container class. """

    # Alternate text attribute name
    alt_attr_name = 'alt'

    # Image width attribute name
    width_attr_name = 'width'

    # Image height attribute name
    height_attr_name = 'height'

    # Allowed schemes for URL
    allowed_schemes = ('http', 'https')

    def get_image_src_link(self, tree_node):
        """
        Get the image source link URL.
        :param tree_node: The current tree node instance.
        :return The image source link URL (not sanitized).
        """
        src_link = tree_node.get_raw_content()
        return sanitize_url(src_link, allowed_schemes=self.allowed_schemes)

    def get_alt_text(self, tree_node):
        """
        Get the image alt text.
        :param tree_node: The current tree node instance.
        :return The image alternative text as string, or an empty string.
        """
        alt_text = tree_node.attrs.get(self.alt_attr_name, '')
        return unescape_html_entities(alt_text)

    def get_img_width(self, tree_node):
        """
        Get the image width, or zero.
        :param tree_node: The current tree node instance.
        :return The image width, or zero.
        """
        try:
            width = int(tree_node.attrs.get(self.width_attr_name, 0))
            if width <= 0:
                return 0
            return width
        except ValueError:
            return 0

    def get_img_height(self, tree_node):
        """
        Get the image height, or zero.
        :param tree_node: The current tree node instance.
        :return The image height, or zero.
        """
        try:
            height = int(tree_node.attrs.get(self.height_attr_name, 0))
            if height <= 0:
                return 0
            return height
        except ValueError:
            return 0

    def render_html(self, tree_node, inner_html, force_rel_nofollow=True):
        """
        Callback function for rendering HTML.
        :param force_rel_nofollow: If set, all links in the rendered HTML will have "rel=nofollow" (default to True).
        :param tree_node: Current tree node to be rendered.
        :param inner_html: Inner HTML of this tree node.
        :return Rendered HTML of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)

        # Shortcut if no source link
        if not src_link:
            return ''

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            extra_attrs = ' alt="%s"' % escape_html(alt_text)
        else:
            extra_attrs = ''

        # Get the image width
        img_width = self.get_img_width(tree_node)
        if img_width:
            extra_attrs += ' width="%d"' % img_width

        # Get the image height
        img_height = self.get_img_height(tree_node)
        if img_width:
            extra_attrs += ' height="%d"' % img_height

        # Render the image
        return '' % (src_link, extra_attrs)

    def render_text(self, tree_node, inner_text):
        """
        Callback function for rendering text.
        :param tree_node: Current tree node to be rendered.
        :param inner_text: Inner text of this tree node.
        :return Rendered text of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)

        # Shortcut if no source link
        if not src_link:
            return inner_text

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            return '%s (%s)' % (src_link, alt_text)
        else:
            return src_link

    def render_skcode(self, tree_node, inner_skcode):
        """
        Callback function for rendering SkCode.
        :param tree_node: Current tree node to be rendered.
        :param inner_skcode: Inner SkCode of this tree node.
        :return Rendered SkCode of this node.
        """

        # Get the image source link
        src_link = self.get_image_src_link(tree_node)

        # Shortcut if no source link
        if not src_link:
            return inner_skcode

        # Get the alternative text
        alt_text = self.get_alt_text(tree_node)
        if alt_text:
            extra_attrs = ' %s=%s' % (self.alt_attr_name,
                                      escape_attrvalue(alt_text))
        else:
            extra_attrs = ''

        # Get the image width
        img_width = self.get_img_width(tree_node)
        if img_width:
            extra_attrs += ' %s="%d"' % (self.width_attr_name, img_width)

        # Get the image height
        img_height = self.get_img_height(tree_node)
        if img_width:
            extra_attrs += ' %s="%d"' % (self.height_attr_name, img_height)

        # Render the image
        node_name = tree_node.name
        return '[%s%s]%s[/%s]' % (node_name, extra_attrs, src_link, node_name)

class YoutubeTagOptions(TagOptions):
    """ Youtube video integration tag options container class. """

    # Default iframe width
    default_iframe_width = 560

    # Default iframe height
    default_iframe_height = 315

    def get_youtube_video_id(self, tree_node):
        """
        Get the Youtube video ID.
        :param tree_node: Current tree node to be rendered.
        :return The video ID, or an empty string.
        """

        # Get the URL
        url = tree_node.get_raw_content()

        # Split the URL
        try:
            scheme, netloc, path, query, fragment = urlsplit(url)
        except ValueError:

            # Handle malformed URL
            return ''

        # Check for bad netloc
        if netloc != 'www.youtube.com' and netloc != 'youtube.com':
            return ''

        # Check for bad url query
        if not query:
            return ''

        # Get the video ID
        query_args = parse_qs(query)
        video_id = query_args.get('v', None)

        # Query args return a list of values
        if not video_id:
            return ''
        else:
            return video_id[0].strip()

    def render_html(self, tree_node, inner_html, force_rel_nofollow=True):
        """
        Callback function for rendering HTML.
        :param force_rel_nofollow: If set, all links in the rendered HTML will have "rel=nofollow" (default to True).
        :param tree_node: Current tree node to be rendered.
        :param inner_html: Inner HTML of this tree node.
        :return Rendered HTML of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the iframe
        if video_id:
            return '' % (self.default_iframe_width,
                                                                  self.default_iframe_height,
                                                                  quote_plus(video_id))
        else:
            return inner_html

    def render_text(self, tree_node, inner_text):
        """
        Callback function for rendering text.
        :param tree_node: Current tree node to be rendered.
        :param inner_text: Inner text of this tree node.
        :return Rendered text of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the video link
        if video_id:
            return 'https://www.youtube.com/watch?v=%s' % quote_plus(video_id)
        else:
            return inner_text

    def render_skcode(self, tree_node, inner_skcode):
        """
        Callback function for rendering SkCode.
        :param tree_node: Current tree node to be rendered.
        :param inner_skcode: Inner SkCode of this tree node.
        :return Rendered SkCode of this node.
        """

        # Get the video ID
        video_id = self.get_youtube_video_id(tree_node)

        # Render the video tag
        if video_id:
            node_name = tree_node.name
            return '[%s]https://www.youtube.com/watch?v=%s[/%s]' % (node_name, quote_plus(video_id), node_name)
        else:
            return inner_skcode
<|endoftext|>


<commit_before>import datetime
import os

import dotenv
root_path = os.path.dirname(os.path.realpath(__file__)) + '/../'
dotenv.read_dotenv(os.path.join(root_path, '.env'))
from getenv import env
import pytz

def prune_note_files(note_files):
    files = [note_file for note_file in note_files if '~' not in note_file]
    return files

def get_note_files():
    current_directory = os.path.dirname(os.path.realpath(__file__))
    notes_directory = os.path.join(current_directory, 'notes')
    files = os.listdir(notes_directory)
    files.sort(reverse=True)
    files = prune_note_files(files)
    files = [os.path.join(notes_directory, note_file) for note_file in files]
    return files

def get_notes():
    note_files = get_note_files()
    timezone = pytz.timezone(env('DISPLAY_TIMEZONE'))
    notes = []
    for note_file in note_files:
        with open(note_file) as note_handle:
            note = note_handle.readlines()
        note = [line.strip() for line in note]
        if len(note) < 3 or not note[1].isdigit():
            continue
        timestamp = int(note[1])
        note_parsed = {}
        note_parsed['title'] = note[0]
        note_parsed['time'] = datetime.datetime.fromtimestamp(
            timestamp, timezone)
        note_parsed['note'] = "\n".join(note[2:])
        notes.append(note_parsed)
    return notes
<commit_msg>Make notes support markdown<commit_after>import datetime
import os

import dotenv
root_path = os.path.dirname(os.path.realpath(__file__)) + '/../'
dotenv.read_dotenv(os.path.join(root_path, '.env'))
from getenv import env
import markdown2
import pytz

def prune_note_files(note_files):
    files = [note_file for note_file in note_files if '~' not in note_file]
    return files

def get_note_files():
    current_directory = os.path.dirname(os.path.realpath(__file__))
    notes_directory = os.path.join(current_directory, 'notes')
    files = os.listdir(notes_directory)
    files.sort(reverse=True)
    files = prune_note_files(files)
    files = [os.path.join(notes_directory, note_file) for note_file in files]
    return files

def get_notes():
    note_files = get_note_files()
    timezone = pytz.timezone(env('DISPLAY_TIMEZONE'))
    notes = []
    for note_file in note_files:
        with open(note_file) as note_handle:
            note = note_handle.readlines()
        note = [line.strip() for line in note]
        if len(note) < 3 or not note[1].isdigit():
            continue
        timestamp = int(note[1])
        note_parsed = {}
        note_parsed['title'] = note[0]
        note_parsed['time'] = datetime.datetime.fromtimestamp(
            timestamp, timezone)
        note_parsed['note'] = markdown2.markdown("\n".join(note[2:]))
        notes.append(note_parsed)
    return notes
<|endoftext|>


<commit_before>''' Test receiving, primarely RawConnection and LircdConnnection. '''

import asyncio
import os
import os.path
import subprocess
import sys
import time
import unittest

testdir = os.path.abspath(os.path.dirname(__file__))
os.chdir(testdir)

sys.path.insert(0, os.path.abspath(os.path.join(testdir, '..')))

from lirc import RawConnection, LircdConnection, CommandConnection
from lirc import AsyncConnection
import lirc

import signal
from contextlib import contextmanager, suppress
from concurrent.futures import TimeoutError

_PACKET_ONE = '0123456789abcdef 00 KEY_1 mceusb'
_LINE_0 = '0123456789abcdef 00 KEY_1 mceusb'
_SOCKET = 'lircd.socket'
_SOCAT = subprocess.check_output('which socat', shell=True) \
    .decode('ascii').strip()
_EXPECT = subprocess.check_output('which expect', shell=True) \
    .decode('ascii').strip()

class TimeoutException(Exception):
    pass

@contextmanager
def event_loop(suppress=[]):

    if isinstance(suppress, type) and issubclass(suppress, Exception):
        suppress = [suppress]
    elif isinstance(suppress, list):
        for ex_type in suppress:
            if not isinstance(ex_type, type) or not issubclass(ex_type, Exception):
                raise ValueError('suppress is not an array of exception types')
    else:
        raise ValueError('suppress is not an exception type')

    ex = []
    def exception_handler(loop, context):
        nonlocal ex
        nonlocal suppress
        for ex_type in suppress:
            if isinstance(context['exception'], ex_type):
                return
        ex.append(context['exception'])

    loop = asyncio.get_event_loop()
    if loop.is_closed():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    loop.set_exception_handler(exception_handler)

    try:
        yield loop
    finally:
        loop.close()
        if len(ex):
            raise Exception('Unhandled exceptions in async code') from ex[0]

def _wait_for_socket():
    ''' Wait until the ncat process has setup the lircd.socket dummy. '''
    i = 0
    while not os.path.exists(_SOCKET):
        time.sleep(0.01)
        i += 1
        if i > 100:
            raise OSError('Cannot find socket file')

class ReceiveTests(unittest.TestCase):
    ''' Test various Connections. '''

    @contextmanager
    def assertCompletedBeforeTimeout(self, timeout):

        triggered = False

        def handle_timeout(signum, frame):
            nonlocal triggered
            triggered = True
            for task in asyncio.Task.all_tasks():
                task.cancel()
            raise TimeoutException()

        try:
            signal.signal(signal.SIGALRM, handle_timeout)
            signal.alarm(timeout)
            with suppress(TimeoutException, asyncio.CancelledError):
                yield
        finally:
            signal.signal(signal.SIGALRM, signal.SIG_DFL)
            signal.alarm(0)
            if triggered:
                raise self.failureException('Code block did not complete before the {} seconds timeout'.format(timeout)) from None

    def testReceiveOneRawLine(self):
        ''' Receive a single, raw line. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT,  'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"echo %s"' % _PACKET_ONE]
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with RawConnection(socket_path=_SOCKET) as conn:
                line = conn.readline()
            self.assertEqual(line, _PACKET_ONE)

    def testReceive10000RawLines(self):
        ''' Receive 10000 raw lines. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
                'EXEC:"%s ./dummy-server 0"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            lines = []
            with RawConnection(socket_path=_SOCKET) as conn:
                for i in range(0, 10000):
                    lines.append(conn.readline())
            self.assertEqual(lines[0], _LINE_0)
            self.assertEqual(lines[9999], _LINE_0.replace(" 00 ", " 09 "))

    def testReceiveOneLine(self):
        ''' Receive a single, translated line OK. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT,  'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"echo %s"' % _PACKET_ONE]
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                line = conn.readline()
        self.assertEqual(line, 'foo-cmd')

    def testReceive1AsyncLines(self):
        ''' Receive 1000 lines using the async interface. '''

        async def get_lines(raw_conn, count, loop):

            nonlocal lines
            async with AsyncConnection(raw_conn, loop) as conn:
                async for keypress in conn:
                    lines.append(keypress)
                    if len(lines) >= count:
                        return lines

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 0"' % _EXPECT]
        lines = []
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                with event_loop() as loop:
                    loop.run_until_complete(get_lines(conn, 1000, loop))

        self.assertEqual(len(lines), 1000)
        self.assertEqual(lines[0], 'foo-cmd')
        self.assertEqual(lines[999], 'foo-cmd')

    def testReceiveTimeout(self):
        ''' Generate a TimeoutException if there is no data '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET, 'EXEC:"sleep 1"']
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                self.assertRaises(lirc.TimeoutException, conn.readline, 0.1)

    def testReceiveAsyncDisconnectDontBlock(self):
        ''' Do not block the loop if connection is lost '''

        async def readline(raw_conn):
            async with AsyncConnection(raw_conn, loop) as conn:
                return await conn.readline()

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET, 'EXEC:"sleep 1"']
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                with event_loop(suppress=[ConnectionResetError, TimeoutException]) as loop:
                    with self.assertCompletedBeforeTimeout(3):
                        with suppress(TimeoutError):
                            loop.run_until_complete(asyncio.wait_for(readline(conn), 2))

class CommandTests(unittest.TestCase):
    ''' Test Command, Reply, ReplyParser and some Commands samples. '''

    def testRemotesCommmand(self):
        ''' Do LIST without arguments . '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 100"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with CommandConnection(socket_path=_SOCKET) as conn:
                reply = lirc.ListRemotesCommand(conn).run()
            self.assertEqual(len(reply.data), 2)
            self.assertEqual(reply.success, True)
            self.assertEqual(reply.data[0], 'mceusb1')
            self.assertEqual(reply.data[1], 'mceusb2')
            self.assertEqual(reply.sighup, False)

    def testSighupReply(self):
        ''' Handle an unexpected SIGHUP in SEND_STOP reply. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 100"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with CommandConnection(socket_path=_SOCKET) as conn:
                reply = lirc.StopRepeatCommand(conn, 'mceusb', 'KEY_1').run()
            self.assertEqual(len(reply.data), 0)
            self.assertEqual(reply.success, True)
            self.assertEqual(reply.sighup, True)

if __name__ == '__main__':
    unittest.main()
<commit_msg>python-pkg: test: if the connection is lost while using sync code, a ConnectionResetError should be raised<commit_after>''' Test receiving, primarely RawConnection and LircdConnnection. '''

import asyncio
import os
import os.path
import subprocess
import sys
import time
import unittest

testdir = os.path.abspath(os.path.dirname(__file__))
os.chdir(testdir)

sys.path.insert(0, os.path.abspath(os.path.join(testdir, '..')))

from lirc import RawConnection, LircdConnection, CommandConnection
from lirc import AsyncConnection
import lirc

import signal
from contextlib import contextmanager, suppress
from concurrent.futures import TimeoutError

_PACKET_ONE = '0123456789abcdef 00 KEY_1 mceusb'
_LINE_0 = '0123456789abcdef 00 KEY_1 mceusb'
_SOCKET = 'lircd.socket'
_SOCAT = subprocess.check_output('which socat', shell=True) \
    .decode('ascii').strip()
_EXPECT = subprocess.check_output('which expect', shell=True) \
    .decode('ascii').strip()

class TimeoutException(Exception):
    pass

@contextmanager
def event_loop(suppress=[]):

    if isinstance(suppress, type) and issubclass(suppress, Exception):
        suppress = [suppress]
    elif isinstance(suppress, list):
        for ex_type in suppress:
            if not isinstance(ex_type, type) or not issubclass(ex_type, Exception):
                raise ValueError('suppress is not an array of exception types')
    else:
        raise ValueError('suppress is not an exception type')

    ex = []
    def exception_handler(loop, context):
        nonlocal ex
        nonlocal suppress
        for ex_type in suppress:
            if isinstance(context['exception'], ex_type):
                return
        ex.append(context['exception'])

    loop = asyncio.get_event_loop()
    if loop.is_closed():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    loop.set_exception_handler(exception_handler)

    try:
        yield loop
    finally:
        loop.close()
        if len(ex):
            raise Exception('Unhandled exceptions in async code') from ex[0]

def _wait_for_socket():
    ''' Wait until the ncat process has setup the lircd.socket dummy. '''
    i = 0
    while not os.path.exists(_SOCKET):
        time.sleep(0.01)
        i += 1
        if i > 100:
            raise OSError('Cannot find socket file')

class ReceiveTests(unittest.TestCase):
    ''' Test various Connections. '''

    @contextmanager
    def assertCompletedBeforeTimeout(self, timeout):

        triggered = False

        def handle_timeout(signum, frame):
            nonlocal triggered
            triggered = True
            for task in asyncio.Task.all_tasks():
                task.cancel()
            raise TimeoutException()

        try:
            signal.signal(signal.SIGALRM, handle_timeout)
            signal.alarm(timeout)
            with suppress(TimeoutException, asyncio.CancelledError):
                yield
        finally:
            signal.signal(signal.SIGALRM, signal.SIG_DFL)
            signal.alarm(0)
            if triggered:
                raise self.failureException('Code block did not complete before the {} seconds timeout'.format(timeout)) from None

    def testReceiveOneRawLine(self):
        ''' Receive a single, raw line. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT,  'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"echo %s"' % _PACKET_ONE]
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with RawConnection(socket_path=_SOCKET) as conn:
                line = conn.readline()
            self.assertEqual(line, _PACKET_ONE)

    def testReceive10000RawLines(self):
        ''' Receive 10000 raw lines. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
                'EXEC:"%s ./dummy-server 0"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            lines = []
            with RawConnection(socket_path=_SOCKET) as conn:
                for i in range(0, 10000):
                    lines.append(conn.readline())
            self.assertEqual(lines[0], _LINE_0)
            self.assertEqual(lines[9999], _LINE_0.replace(" 00 ", " 09 "))

    def testReceiveOneLine(self):
        ''' Receive a single, translated line OK. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT,  'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"echo %s"' % _PACKET_ONE]
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                line = conn.readline()
        self.assertEqual(line, 'foo-cmd')

    def testReceive1AsyncLines(self):
        ''' Receive 1000 lines using the async interface. '''

        async def get_lines(raw_conn, count, loop):

            nonlocal lines
            async with AsyncConnection(raw_conn, loop) as conn:
                async for keypress in conn:
                    lines.append(keypress)
                    if len(lines) >= count:
                        return lines

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 0"' % _EXPECT]
        lines = []
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                with event_loop() as loop:
                    loop.run_until_complete(get_lines(conn, 1000, loop))

        self.assertEqual(len(lines), 1000)
        self.assertEqual(lines[0], 'foo-cmd')
        self.assertEqual(lines[999], 'foo-cmd')

    def testReceiveTimeout(self):
        ''' Generate a TimeoutException if there is no data '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET, 'EXEC:"sleep 1"']
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                self.assertRaises(lirc.TimeoutException, conn.readline, 0.1)

    def testReceiveDisconnect(self):
        ''' Generate a ConnectionResetError if connection is lost '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET, 'EXEC:"sleep 1"']
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                with self.assertRaises(ConnectionResetError):
                    with self.assertCompletedBeforeTimeout(3):
                        conn.readline(2)

    def testReceiveAsyncDisconnectDontBlock(self):
        ''' Do not block the loop if connection is lost '''

        async def readline(raw_conn):
            async with AsyncConnection(raw_conn, loop) as conn:
                return await conn.readline()

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET, 'EXEC:"sleep 1"']
        with subprocess.Popen(cmd) as child:
            _wait_for_socket()
            with LircdConnection('foo',
                                 socket_path=_SOCKET,
                                 lircrc_path='lircrc.conf') as conn:
                with event_loop(suppress=[ConnectionResetError, TimeoutException]) as loop:
                    with self.assertCompletedBeforeTimeout(3):
                        with suppress(TimeoutError):
                            loop.run_until_complete(asyncio.wait_for(readline(conn), 2))

class CommandTests(unittest.TestCase):
    ''' Test Command, Reply, ReplyParser and some Commands samples. '''

    def testRemotesCommmand(self):
        ''' Do LIST without arguments . '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 100"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with CommandConnection(socket_path=_SOCKET) as conn:
                reply = lirc.ListRemotesCommand(conn).run()
            self.assertEqual(len(reply.data), 2)
            self.assertEqual(reply.success, True)
            self.assertEqual(reply.data[0], 'mceusb1')
            self.assertEqual(reply.data[1], 'mceusb2')
            self.assertEqual(reply.sighup, False)

    def testSighupReply(self):
        ''' Handle an unexpected SIGHUP in SEND_STOP reply. '''

        if os.path.exists(_SOCKET):
            os.unlink(_SOCKET)
        cmd = [_SOCAT, 'UNIX-LISTEN:' + _SOCKET,
               'EXEC:"%s ./dummy-server 100"' % _EXPECT]
        with subprocess.Popen(cmd,
                              stdout = subprocess.PIPE,
                              stderr = subprocess.STDOUT) as child:
            _wait_for_socket()
            with CommandConnection(socket_path=_SOCKET) as conn:
                reply = lirc.StopRepeatCommand(conn, 'mceusb', 'KEY_1').run()
            self.assertEqual(len(reply.data), 0)
            self.assertEqual(reply.success, True)
            self.assertEqual(reply.sighup, True)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before># coding: utf-8

""" Make an SCF run given initial conditions, timestep, and potential parameters """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys
import shutil

# Third-party
from astropy.constants import G
from astropy import log as logger
import astropy.units as u

# Project
from scf import project_path
import scf.potential as sp

def main(name, pos, vel, scfpars, potentials, run_path, overwrite=False, submit=False):
    run_path = os.path.abspath(run_path)
    # run_path = os.path.abspath(os.path.join(project_path, "simulations", "runs"))
    template_path = os.path.abspath(os.path.join(project_path, "templates"))
    src_path = os.path.abspath(os.path.join(project_path, "src"))
    logger.debug("Run path: {}".format(run_path))

    path = os.path.join(run_path, name)
    if os.path.exists(path) and overwrite:
        logger.info("nuking directory...")
        shutil.rmtree(path)
    elif os.path.exists(path) and not overwrite:
        logger.warning("Path '{}' already exists -- do you want to overwrite?"
                       .format(path))
        yn = raw_input("[y/N]: ")
        if yn.lower().strip() == 'y':
            logger.info("nuking directory...")
            shutil.rmtree(path)
        else:
            logger.info("aborting...")
            sys.exit(0)

        raise IOError("Path '{}' already exists!".format(path))
        sys.exit(0)

    if not os.path.exists(path):
        os.makedirs(path)

    # parse pos and vel
    x_vals,x_unit = pos.split()
    v_vals,v_unit = vel.split()

    pos = map(float, x_vals.split(',')) * u.Unit(x_unit)
    pos = pos.to(u.kpc).value

    vel = map(float, v_vals.split(',')) * u.Unit(v_unit)
    vel = vel.to(u.km/u.s).value

    # combine all of the different blocks
    allblocks = dict()
    for potential_name in potentials:
        blocks = getattr(sp,potential_name)()
        for k,v in blocks.items():
            if k not in allblocks:
                allblocks[k] = ""

            allblocks[k] += v

            if k != "SCFPOT":
                allblocks[k] += "\n"

    # ------------------------------------------------------------------------
    # SCFPAR
    with open(os.path.join(template_path, "SCFPAR.tpl"), 'r') as f:
        base_SCFPAR = f.read()

    with open(os.path.join(path, "SCFPAR"), 'w') as f:
        f.write(base_SCFPAR.format(x=pos, v=vel, **scfpars))

    # ------------------------------------------------------------------------
    # Makefile
    with open(os.path.join(template_path, "Makefile.tpl"), 'r') as f:
        base_Makefile = f.read()

    with open(os.path.join(path, "Makefile"), 'w') as f:
        f.write(base_Makefile)

    # cmd = "unexpand {fn} > {fn}2".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)
    # cmd = "mv {fn}2 {fn}".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)

    # ------------------------------------------------------------------------
    # potential.h
    with open(os.path.join(template_path, "potential.h.tpl"), 'r') as f:
        base_potentialh = f.read()

    with open(os.path.join(path, "potential.h"), 'w') as f:
        f.write(base_potentialh.format(hblock=allblocks['hblock']))

    # ------------------------------------------------------------------------
    # potential.f
    with open(os.path.join(template_path, "potential.f.tpl"), 'r') as f:
        base_potentialf = f.read()

    with open(os.path.join(path, "potential.f"), 'w') as f:
        f.write(base_potentialf.format(**allblocks))

    # ------------------------------------------------------------------------
    # SCFPOT
    with open(os.path.join(path, "SCFPOT"), 'w') as f:
        f.write(allblocks['SCFPOT'])

    if submit:
        with open(os.path.join(template_path, "submit.tpl"), 'r') as f:
            base_submit = f.read()

        with open(os.path.join(path, "submit.sh"), 'w') as f:
            f.write(base_submit.format(path=path))

    # Copy scf.f and scf.h
    shutil.copyfile(os.path.join(src_path, 'scf.f'),
                    os.path.join(path, 'scf.f'))
    shutil.copyfile(os.path.join(src_path, 'scf.h'),
                    os.path.join(path, 'scf.h'))

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY.")

    parser.add_argument("--path", "-p", dest="path", type=str, help="Path.", required=True)
    parser.add_argument("--name", dest="name", type=str, help="Name.", required=True)
    parser.add_argument("--potentials", dest="potentials", type=str,
                        help="Potential names. One of: {0}".format(",".join(sp.__all__)),
                        nargs='+', required=True)
    parser.add_argument("-s", "--submit", action="store_true", dest="submit",
                        default=True, help="Generate a submit.sh file as well.")
    parser.add_argument("--pos", dest="x", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --x='15.324,123.314,51.134 kpc'")
    parser.add_argument("--vel", dest="v", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --v='15.324,123.314,51.134 km/s'")

    # Parameters for SCF
    parser.add_argument("--dt", dest="dt", default=1., type=float,
                        help="Time step in Myr.")
    parser.add_argument("--nsteps", dest="nsteps", type=int,
                        help="Number of steps.", required=True)
    parser.add_argument("--ncen", dest="ncen", type=int, default=10,
                        help="Output to SCFCEN every (this number) steps.")
    parser.add_argument("--nsnap", dest="nsnap", type=int, default=1000,
                        help="Output to a SNAP file every (this number) steps.")
    parser.add_argument("--ntide", dest="ntide", default=500, type=int,
                        help="Number of steps over which to turn on the tidal field.")
    parser.add_argument("--mass", dest="mass", type=float, required=True,
                        help="Mass in solar masses within 35 scale radii of the "
                             "cluster/satellite.")
    parser.add_argument("--rscale", dest="rscale", default=None, type=float,
                        help="Scale radius of the cluster/satellite.")
    parser.add_argument("--SCFBI", dest="scfbi_path", default=None, type=str,
                        help="Path to an SCFBI file.")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    # Contains user-specified parameters for SCF
    scfpars = dict()

    if args.rscale is None:
        ru = 0.43089*(args.mass/2.5e9)**(1/3.)
    else:
        ru = args.rscale
    scfpars['rscale'] = ru
    scfpars['mass'] = args.mass

    _G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
    X = (_G / ru**3 * args.mass)**-0.5

    length_unit = u.Unit("{0} kpc".format(ru))
    mass_unit = u.Unit("{0} M_sun".format(args.mass))
    time_unit = u.Unit("{:08f} Myr".format(X))

    scfpars['dt'] = args.dt / (1*time_unit).to(u.Myr).value
    scfpars['nsteps'] = args.nsteps
    scfpars['ncen'] = args.ncen
    scfpars['nsnap'] = args.nsnap
    scfpars['ntide'] = args.ntide
    if args.scfbi_path is None:
        scfpars['SCFBIpath'] = os.path.abspath(os.path.join(project_path, 'src', 'SCFBI'))

    main(name=args.name, pos=args.x, vel=args.v, scfpars=scfpars,
         potentials=args.potentials, overwrite=args.overwrite,
         submit=args.submit, run_path=args.path)
<commit_msg>name too<commit_after># coding: utf-8

""" Make an SCF run given initial conditions, timestep, and potential parameters """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys
import shutil

# Third-party
from astropy.constants import G
from astropy import log as logger
import astropy.units as u

# Project
from scf import project_path
import scf.potential as sp

def main(name, pos, vel, scfpars, potentials, run_path, overwrite=False, submit=False):
    run_path = os.path.abspath(run_path)
    # run_path = os.path.abspath(os.path.join(project_path, "simulations", "runs"))
    template_path = os.path.abspath(os.path.join(project_path, "templates"))
    src_path = os.path.abspath(os.path.join(project_path, "src"))
    logger.debug("Run path: {}".format(run_path))

    path = os.path.join(run_path, name)
    if os.path.exists(path) and overwrite:
        logger.info("nuking directory...")
        shutil.rmtree(path)
    elif os.path.exists(path) and not overwrite:
        logger.warning("Path '{}' already exists -- do you want to overwrite?"
                       .format(path))
        yn = raw_input("[y/N]: ")
        if yn.lower().strip() == 'y':
            logger.info("nuking directory...")
            shutil.rmtree(path)
        else:
            logger.info("aborting...")
            sys.exit(0)

        raise IOError("Path '{}' already exists!".format(path))
        sys.exit(0)

    if not os.path.exists(path):
        os.makedirs(path)

    # parse pos and vel
    x_vals,x_unit = pos.split()
    v_vals,v_unit = vel.split()

    pos = map(float, x_vals.split(',')) * u.Unit(x_unit)
    pos = pos.to(u.kpc).value

    vel = map(float, v_vals.split(',')) * u.Unit(v_unit)
    vel = vel.to(u.km/u.s).value

    # combine all of the different blocks
    allblocks = dict()
    for potential_name in potentials:
        blocks = getattr(sp,potential_name)()
        for k,v in blocks.items():
            if k not in allblocks:
                allblocks[k] = ""

            allblocks[k] += v

            if k != "SCFPOT":
                allblocks[k] += "\n"

    # ------------------------------------------------------------------------
    # SCFPAR
    with open(os.path.join(template_path, "SCFPAR.tpl"), 'r') as f:
        base_SCFPAR = f.read()

    with open(os.path.join(path, "SCFPAR"), 'w') as f:
        f.write(base_SCFPAR.format(x=pos, v=vel, **scfpars))

    # ------------------------------------------------------------------------
    # Makefile
    with open(os.path.join(template_path, "Makefile.tpl"), 'r') as f:
        base_Makefile = f.read()

    with open(os.path.join(path, "Makefile"), 'w') as f:
        f.write(base_Makefile)

    # cmd = "unexpand {fn} > {fn}2".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)
    # cmd = "mv {fn}2 {fn}".format(fn=os.path.join(path, "Makefile"))
    # os.system(cmd)

    # ------------------------------------------------------------------------
    # potential.h
    with open(os.path.join(template_path, "potential.h.tpl"), 'r') as f:
        base_potentialh = f.read()

    with open(os.path.join(path, "potential.h"), 'w') as f:
        f.write(base_potentialh.format(hblock=allblocks['hblock']))

    # ------------------------------------------------------------------------
    # potential.f
    with open(os.path.join(template_path, "potential.f.tpl"), 'r') as f:
        base_potentialf = f.read()

    with open(os.path.join(path, "potential.f"), 'w') as f:
        f.write(base_potentialf.format(**allblocks))

    # ------------------------------------------------------------------------
    # SCFPOT
    with open(os.path.join(path, "SCFPOT"), 'w') as f:
        f.write(allblocks['SCFPOT'])

    if submit:
        with open(os.path.join(template_path, "submit.tpl"), 'r') as f:
            base_submit = f.read()

        with open(os.path.join(path, "submit.sh"), 'w') as f:
            f.write(base_submit.format(path=path, name=name))

    # Copy scf.f and scf.h
    shutil.copyfile(os.path.join(src_path, 'scf.f'),
                    os.path.join(path, 'scf.f'))
    shutil.copyfile(os.path.join(src_path, 'scf.h'),
                    os.path.join(path, 'scf.h'))

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY.")

    parser.add_argument("--path", "-p", dest="path", type=str, help="Path.", required=True)
    parser.add_argument("--name", dest="name", type=str, help="Name.", required=True)
    parser.add_argument("--potentials", dest="potentials", type=str,
                        help="Potential names. One of: {0}".format(",".join(sp.__all__)),
                        nargs='+', required=True)
    parser.add_argument("-s", "--submit", action="store_true", dest="submit",
                        default=True, help="Generate a submit.sh file as well.")
    parser.add_argument("--pos", dest="x", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --x='15.324,123.314,51.134 kpc'")
    parser.add_argument("--vel", dest="v", required=True,
                        type=str, help="Initial position in 3 space. Must specify a "
                                       "unit, e.g., --v='15.324,123.314,51.134 km/s'")

    # Parameters for SCF
    parser.add_argument("--dt", dest="dt", default=1., type=float,
                        help="Time step in Myr.")
    parser.add_argument("--nsteps", dest="nsteps", type=int,
                        help="Number of steps.", required=True)
    parser.add_argument("--ncen", dest="ncen", type=int, default=10,
                        help="Output to SCFCEN every (this number) steps.")
    parser.add_argument("--nsnap", dest="nsnap", type=int, default=1000,
                        help="Output to a SNAP file every (this number) steps.")
    parser.add_argument("--ntide", dest="ntide", default=500, type=int,
                        help="Number of steps over which to turn on the tidal field.")
    parser.add_argument("--mass", dest="mass", type=float, required=True,
                        help="Mass in solar masses within 35 scale radii of the "
                             "cluster/satellite.")
    parser.add_argument("--rscale", dest="rscale", default=None, type=float,
                        help="Scale radius of the cluster/satellite.")
    parser.add_argument("--SCFBI", dest="scfbi_path", default=None, type=str,
                        help="Path to an SCFBI file.")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    # Contains user-specified parameters for SCF
    scfpars = dict()

    if args.rscale is None:
        ru = 0.43089*(args.mass/2.5e9)**(1/3.)
    else:
        ru = args.rscale
    scfpars['rscale'] = ru
    scfpars['mass'] = args.mass

    _G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
    X = (_G / ru**3 * args.mass)**-0.5

    length_unit = u.Unit("{0} kpc".format(ru))
    mass_unit = u.Unit("{0} M_sun".format(args.mass))
    time_unit = u.Unit("{:08f} Myr".format(X))

    scfpars['dt'] = args.dt / (1*time_unit).to(u.Myr).value
    scfpars['nsteps'] = args.nsteps
    scfpars['ncen'] = args.ncen
    scfpars['nsnap'] = args.nsnap
    scfpars['ntide'] = args.ntide
    if args.scfbi_path is None:
        scfpars['SCFBIpath'] = os.path.abspath(os.path.join(project_path, 'src', 'SCFBI'))

    main(name=args.name, pos=args.x, vel=args.v, scfpars=scfpars,
         potentials=args.potentials, overwrite=args.overwrite,
         submit=args.submit, run_path=args.path)
<|endoftext|>


<commit_before>import logging
from pprint import pprint, pformat  # noqa

from aleph.core import es
from aleph.index.util import authz_query, field_filter_query
from aleph.search.result import SearchQueryResult
from aleph.search.parser import SearchQueryParser

log = logging.getLogger(__name__)

def convert_filters(filters):
    ret = []
    for field, values in filters.items():
        ret.append(field_filter_query(field, values))
    return ret

class Query(object):
    RESULT_CLASS = SearchQueryResult
    INCLUDE_FIELDS = None
    EXCLUDE_FIELDS = None
    TEXT_FIELDS = ['text']
    PREFIX_FIELD = 'name'
    SKIP_FILTERS = []
    SORT_FIELDS = {
        'label': 'label.kw',
        'name': 'name.kw',
        'score': '_score',
    }
    SORT_DEFAULT = ['_score']

    def __init__(self, parser):
        self.parser = parser

    def get_text_query(self):
        query = []
        if self.parser.text:
            query.append({
                "query_string": {
                    "query": self.parser.text,
                    "fields": self.TEXT_FIELDS,
                    "analyzer": "icu_latin",
                    "default_operator": "AND",
                    "minimum_should_match": "3<80%",
                    "lenient": True
                }
            })
        if self.parser.prefix:
            query.append({
                "match_phrase_prefix": {
                    self.PREFIX_FIELD: self.parser.prefix
                }
            })
        if not len(query):
            query.append({'match_all': {}})
        return query

    def get_filters(self):
        """Apply query filters from the user interface."""
        filters = []
        # This enforces the authorization (access control) rules on
        # a particular query by comparing the collections a user is
        # authorized for with the one on the document.
        if self.parser.authz and not self.parser.authz.is_admin:
            filters.append(authz_query(self.parser.authz))

        for field, values in self.parser.filters.items():
            if field in self.SKIP_FILTERS:
                continue
            if field not in self.parser.facet_names:
                filters.append(field_filter_query(field, values))
        return filters

    def get_negative_filters(self):
        """Apply negative filters."""
        filters = []
        for field, values in self.parser.empties.items():
            filters.append({'exists': {'field': field}})

        if len(self.parser.exclude):
            filters.append({'ids': {'values': self.parser.exclude}})
        return filters

    def get_post_filters(self, exclude=None):
        """Apply post-aggregation query filters."""
        filters = []
        for field, values in self.parser.filters.items():
            if field in self.SKIP_FILTERS or field == exclude:
                continue
            if field in self.parser.facet_filters:
                filters.append(field_filter_query(field, values))
        return {'bool': {'filter': filters}}

    def get_query(self):
        return {
            'bool': {
                'should': self.get_text_query(),
                'must': [],
                'must_not': self.get_negative_filters(),
                'filter': self.get_filters(),
                'minimum_should_match': 1
            }
        }

    def get_aggregations(self):
        """Aggregate the query in order to generate faceted results."""
        aggregations = {}
        for facet_name in self.parser.facet_names:
            facet_aggregations = {}
            if self.parser.get_facet_values(facet_name):
                facet_aggregations[facet_name] = {
                    'terms': {
                        'field': facet_name,
                        'size': self.parser.get_facet_size(facet_name)
                    }
                }

            if self.parser.get_facet_total(facet_name):
                # Option to return total distinct value counts for
                # a given facet, instead of the top buckets.
                agg_name = '%s.cardinality' % facet_name
                facet_aggregations[agg_name] = {
                    'cardinality': {
                        'field': facet_name
                    }
                }

            if not len(facet_aggregations):
                break

            # See here for an explanation of the whole post_filters and
            # aggregation filters thing:
            # https://www.elastic.co/guide/en/elasticsearch/reference/6.2/search-request-post-filter.html  # noqa
            other_filters = self.get_post_filters(exclude=facet_name)
            if len(other_filters['bool']['filter']):
                agg_name = '%s.filtered' % facet_name
                aggregations[agg_name] = {
                    'filter': other_filters,
                    'aggregations': facet_aggregations
                }
            else:
                aggregations.update(facet_aggregations)

        return aggregations

    def get_sort(self):
        """Pick one of a set of named result orderings."""
        if not len(self.parser.sorts):
            return self.SORT_DEFAULT

        sort_fields = ['_score']
        for (field, direction) in self.parser.sorts:
            field = self.SORT_FIELDS.get(field, field)
            sort_fields.append({field: direction})
        return list(reversed(sort_fields))

    def get_highlight(self):
        if not self.parser.highlight:
            return {}
        return {
            'fields': {
                'text': {
                    'type': 'fvh',
                    'number_of_fragments': self.parser.highlight_count,
                    'fragment_size': self.parser.highlight_length
                }
            }
        }

    def get_source(self):
        source = {}
        if self.INCLUDE_FIELDS:
            source['includes'] = self.INCLUDE_FIELDS
        elif self.EXCLUDE_FIELDS:
            source['excludes'] = self.EXCLUDE_FIELDS
        return source

    def get_body(self):
        body = {
            'query': self.get_query(),
            'post_filter': self.get_post_filters(),
            'from': self.parser.offset,
            'size': self.parser.limit,
            'aggregations': self.get_aggregations(),
            'sort': self.get_sort(),
            'highlight': self.get_highlight(),
            '_source': self.get_source()
        }
        # log.info("Query: %s", pformat(body))
        return body

    def search(self):
        """Execute the query as assmbled."""
        # log.info("Search index: %s", self.get_index())
        result = es.search(index=self.get_index(),
                           body=self.get_body())
        log.info("Took: %sms", result.get('took'))
        # log.info("%s", pformat(result.get('profile')))
        return result

    @classmethod
    def handle(cls, request, parser=None, **kwargs):
        if parser is None:
            parser = SearchQueryParser(request.args, request.authz)
        result = cls(parser, **kwargs).search()
        return cls.RESULT_CLASS(request, parser, result)
<commit_msg>Cleaning<commit_after>import logging
from pprint import pprint, pformat  # noqa

from aleph.core import es
from aleph.index.util import authz_query, field_filter_query
from aleph.search.result import SearchQueryResult
from aleph.search.parser import SearchQueryParser

log = logging.getLogger(__name__)

def convert_filters(filters):
    ret = []
    for field, values in filters.items():
        ret.append(field_filter_query(field, values))
    return ret

class Query(object):
    RESULT_CLASS = SearchQueryResult
    INCLUDE_FIELDS = None
    EXCLUDE_FIELDS = None
    TEXT_FIELDS = ['text']
    PREFIX_FIELD = 'name'
    SKIP_FILTERS = []
    SORT_FIELDS = {
        'label': 'label.kw',
        'name': 'name.kw',
        'score': '_score',
    }
    SORT_DEFAULT = ['_score']

    def __init__(self, parser):
        self.parser = parser

    def get_text_query(self):
        query = []
        if self.parser.text:
            query.append({
                "query_string": {
                    "query": self.parser.text,
                    "fields": self.TEXT_FIELDS,
                    "analyzer": "icu_latin",
                    "default_operator": "AND",
                    "minimum_should_match": "3<80%",
                    "lenient": True
                }
            })
        if self.parser.prefix:
            query.append({
                "match_phrase_prefix": {
                    self.PREFIX_FIELD: self.parser.prefix
                }
            })
        if not len(query):
            query.append({'match_all': {}})
        return query

    def get_filters(self):
        """Apply query filters from the user interface."""
        filters = []
        # This enforces the authorization (access control) rules on
        # a particular query by comparing the collections a user is
        # authorized for with the one on the document.
        if self.parser.authz and not self.parser.authz.is_admin:
            filters.append(authz_query(self.parser.authz))

        for field, values in self.parser.filters.items():
            if field in self.SKIP_FILTERS:
                continue
            if field not in self.parser.facet_names:
                filters.append(field_filter_query(field, values))
        return filters

    def get_negative_filters(self):
        """Apply negative filters."""
        filters = []
        for field, _ in self.parser.empties.items():
            filters.append({'exists': {'field': field}})

        if len(self.parser.exclude):
            filters.append({'ids': {'values': self.parser.exclude}})
        return filters

    def get_post_filters(self, exclude=None):
        """Apply post-aggregation query filters."""
        filters = []
        for field, values in self.parser.filters.items():
            if field in self.SKIP_FILTERS or field == exclude:
                continue
            if field in self.parser.facet_filters:
                filters.append(field_filter_query(field, values))
        return {'bool': {'filter': filters}}

    def get_query(self):
        return {
            'bool': {
                'should': self.get_text_query(),
                'must': [],
                'must_not': self.get_negative_filters(),
                'filter': self.get_filters(),
                'minimum_should_match': 1
            }
        }

    def get_aggregations(self):
        """Aggregate the query in order to generate faceted results."""
        aggregations = {}
        for facet_name in self.parser.facet_names:
            facet_aggregations = {}
            if self.parser.get_facet_values(facet_name):
                facet_aggregations[facet_name] = {
                    'terms': {
                        'field': facet_name,
                        'size': self.parser.get_facet_size(facet_name)
                    }
                }

            if self.parser.get_facet_total(facet_name):
                # Option to return total distinct value counts for
                # a given facet, instead of the top buckets.
                agg_name = '%s.cardinality' % facet_name
                facet_aggregations[agg_name] = {
                    'cardinality': {
                        'field': facet_name
                    }
                }

            if not len(facet_aggregations):
                break

            # See here for an explanation of the whole post_filters and
            # aggregation filters thing:
            # https://www.elastic.co/guide/en/elasticsearch/reference/6.2/search-request-post-filter.html  # noqa
            other_filters = self.get_post_filters(exclude=facet_name)
            if len(other_filters['bool']['filter']):
                agg_name = '%s.filtered' % facet_name
                aggregations[agg_name] = {
                    'filter': other_filters,
                    'aggregations': facet_aggregations
                }
            else:
                aggregations.update(facet_aggregations)

        return aggregations

    def get_sort(self):
        """Pick one of a set of named result orderings."""
        if not len(self.parser.sorts):
            return self.SORT_DEFAULT

        sort_fields = ['_score']
        for (field, direction) in self.parser.sorts:
            field = self.SORT_FIELDS.get(field, field)
            sort_fields.append({field: direction})
        return list(reversed(sort_fields))

    def get_highlight(self):
        if not self.parser.highlight:
            return {}
        return {
            'fields': {
                'text': {
                    'type': 'fvh',
                    'number_of_fragments': self.parser.highlight_count,
                    'fragment_size': self.parser.highlight_length
                }
            }
        }

    def get_source(self):
        source = {}
        if self.INCLUDE_FIELDS:
            source['includes'] = self.INCLUDE_FIELDS
        elif self.EXCLUDE_FIELDS:
            source['excludes'] = self.EXCLUDE_FIELDS
        return source

    def get_body(self):
        body = {
            'query': self.get_query(),
            'post_filter': self.get_post_filters(),
            'from': self.parser.offset,
            'size': self.parser.limit,
            'aggregations': self.get_aggregations(),
            'sort': self.get_sort(),
            'highlight': self.get_highlight(),
            '_source': self.get_source()
        }
        # log.info("Query: %s", pformat(body))
        return body

    def get_index(self):
        raise NotImplementedError

    def search(self):
        """Execute the query as assmbled."""
        # log.info("Search index: %s", self.get_index())
        result = es.search(index=self.get_index(),
                           body=self.get_body())
        log.info("Took: %sms", result.get('took'))
        # log.info("%s", pformat(result.get('profile')))
        return result

    @classmethod
    def handle(cls, request, parser=None, **kwargs):
        if parser is None:
            parser = SearchQueryParser(request.args, request.authz)
        result = cls(parser, **kwargs).search()
        return cls.RESULT_CLASS(request, parser, result)
<|endoftext|>