<commit_before># Standard library
from os.path import abspath, expanduser, join
import os
import time

# Third-party
import astropy.units as u
import h5py
import numpy as np
from schwimmbad import choose_pool
from sqlalchemy import func
from thejoker.log import log as joker_logger
from thejoker.sampler import JokerParams, TheJoker, save_prior_samples
from thejoker.utils import quantity_to_hdf5
import yaml

# Project
from twoface.data import APOGEERVData
from twoface.log import log as logger
from twoface.db import Session, db_connect
from twoface.db import JokerRun, AllStar, AllVisit, StarResult, Status, AllVisitToAllStar
from twoface.config import TWOFACE_CACHE_PATH

def main(config_file, pool, seed, overwrite=False, _continue=False):
    config_file = abspath(expanduser(config_file))

    # parse config file
    with open(config_file, 'r') as f:
        config = yaml.load(f.read())

    # filename of sqlite database
    if 'database_file' not in config:
        database_file = None

    else:
        database_file = config['database_file']

    db_path = join(TWOFACE_CACHE_PATH, database_file)

    if not os.path.exists(db_path):
        raise IOError("sqlite database not found at '{}'\n Did you run scripts/initdb.py yet?"
                      .format(db_path))

    logger.debug("Connecting to sqlite database at '{}'".format(db_path))
    engine = db_connect(database_path=db_path, ensure_db_exists=False)
    session = Session()
    logger.debug("...connected!")

    # HACK: for testing
    session.query(StarResult).delete()
    session.query(JokerRun).delete()
    session.commit()

    # see if this run (by name) is in db already, if so, just grab
    _run = session.query(JokerRun).filter(JokerRun.name == config['name']).all()
    if len(_run) == 0:
        logger.info("JokerRun '{}' not found in database - creating new entry..."
                    .format(config['name']))

        # create a JokerRun for this run
        run = JokerRun()
        run.config_file = config_file
        run.name = config['name']
        run.P_min = u.Quantity(*config['hyperparams']['P_min'].split())
        run.P_max = u.Quantity(*config['hyperparams']['P_max'].split())
        run.requested_samples_per_star = int(config['hyperparams']['requested_samples_per_star'])
        run.max_prior_samples = int(config['prior']['max_samples'])
        run.prior_samples_file = join(TWOFACE_CACHE_PATH, config['prior']['samples_file'])

        if 'jitter' in config['hyperparams']:
            run.jitter = u.Quantity(*config['hyperparams']['jitter'].split())

        elif 'jitter_prior_mean' in config['hyperparams']:
            raise NotImplementedError("NOT YET SORRY")

        else:
            # no jitter specified, set to 0
            run.jitter = 0 * u.m/u.s

        # TODO: need a way to pass in velocity trend info

        # TODO: get all stars we're going to process - some filter on number of
        #   observations based on the number of linear parameters
        # TODO: need some way of specifying a constraint on the stars that are to be processed
        stars = session.query(AllStar).join(AllVisitToAllStar, AllVisit)\
                                      .group_by(AllStar.apogee_id)\
                                      .having(func.count(AllVisit.id) >= 3)\
                                      .all()

        run.stars = stars
        session.add(run)
        session.commit()

    elif len(_run) == 1:
        logger.info("JokerRun '{}' already found in database - retrieving".format(config['name']))
        run = _run[0]

    else:
        raise ValueError("Multiple JokerRun rows found for name '{}'".format(config['name']))

    # create TheJoker sampler instance
    rnd = np.random.RandomState(seed=seed)
    logger.debug("Creating TheJoker instance with {}, {}".format(rnd, pool))
    params = JokerParams(P_min=run.P_min, P_max=run.P_max, jitter=run.jitter,
                         anomaly_tol=1E-11)
    joker = TheJoker(params, random_state=rnd, pool=pool)

    # create prior samples cache, store to file and store filename in DB
    if not os.path.exists(run.prior_samples_file):
        logger.debug("Prior samples file not found - generating now...")

        prior_samples = joker.sample_prior(config['prior']['num_cache'])
        prior_units = save_prior_samples(run.prior_samples_file, prior_samples, u.km/u.s) # data in km/s
        del prior_samples
        # TODO: for large prior samples libraries, may want a way to write
        #   without storing all samples in memory

        logger.debug("...done")

    else:
        with h5py.File(run.prior_samples_file, 'r') as f:
            prior_units = [u.Unit(uu) for uu in f.attrs['units']]

    # build a query to get all stars associated with this JokerRun that need processing
    star_query = session.query(AllStar).join(StarResult, JokerRun, Status)\
                                       .filter(JokerRun.name == run.name)\
                                       .filter(Status.id == 0)

    # create a file to cache the results
    n_stars = star_query.count()
    logger.info("{} stars left to process for run '{}'".format(n_stars, run.name))

    results_filename = join(TWOFACE_CACHE_PATH, "{}.hdf5".format(run.name))

    # TODO: what should structure be? currently thinking /APOGEE_ID/key, e.g.,
    #       /2M00000222+5625359/P for period, etc.
    # TODO: deal with existing file
    with h5py.File(results_filename, 'w') as f:
        pass

    # TODO: grab a batch of targets to process
    # star = star_query.limit(1).one() # get a single untouched star
    for star in star_query.all():
        logger.debug("Starting star {}".format(star.apogee_id))
        t0 = time.time()

        data = APOGEERVData.from_visits(star.visits)

        # adaptive scheme for the rejection sampling:
        n_process = 256 * run.requested_samples_per_star
        n_samples = 0 # running total of good samples returned
        start_idx = 0
        all_samples = None
        for n in range(128): # magic number we should never hit
            logger.debug("Iteration {}, using {} prior samples".format(n, n_process))

            # process n_process samples from prior cache
            samples = joker._rejection_sample_from_cache(data, n_process,
                                                         run.prior_samples_file, start_idx)

            n_survive = samples.shape[0]
            if n_survive > 1:
                n_samples += n_survive

                if all_samples is None:
                    all_samples = samples
                else:
                    all_samples = np.vstack((all_samples, samples))

            if n_samples >= run.requested_samples_per_star:
                break

            start_idx += n_process
            n_process *= 2

            if n_process > run.max_prior_samples:
                if n_process == run.max_prior_samples:
                    break

                else:
                    n_process = run.max_prior_samples

        else: # hit maxiter
            # TODO: error, should never get here
            pass

        # for now, it's sufficient to write the run results to an HDF5 file
        n = run.requested_samples_per_star
        samples_dict = joker.unpack_full_samples(all_samples[:n], data.t_offset, prior_units)

        with h5py.File(results_filename, 'r+') as f:
            g = f.create_group(star.apogee_id)
            for key in samples_dict.keys():
                quantity_to_hdf5(g, key, samples_dict[key])

        logger.debug("...done with star {} - {:.2f} seconds".format(star.apogee_id,
                                                                    time.time()-t0))

    pool.close()
    session.close()

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

    # Define parser object
    parser = ArgumentParser(description="")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')

    oc_group = parser.add_mutually_exclusive_group()
    oc_group.add_argument("--overwrite", dest="overwrite", default=False,
                          action="store_true", help="Overwrite any existing results for "
                                                    "this JokerRun.")
    oc_group.add_argument("--continue", dest="_continue", default=False,
                          action="store_true", help="Continue the the JokerRun.")

    parser.add_argument("-s", "--seed", dest="seed", default=None, type=int,
                        help="Random number seed")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--procs", dest="n_procs", default=1,
                       type=int, help="Number of processes.")
    group.add_argument("--mpi", dest="mpi", default=False,
                       action="store_true", help="Run with MPI.")

    #
    parser.add_argument("-c", "--config", dest="config_file", required=True,
                        type=str, help="Path to config file that specifies the parameters for "
                                       "this JokerRun.")

    args = parser.parse_args()

    loggers = [joker_logger, logger]

    for _logger in loggers:
        # Set logger level based on verbose flags
        if args.verbosity != 0:
            if args.verbosity == 1:
                _logger.setLevel(logging.DEBUG)
            else: # anything >= 2
                _logger.setLevel(1)

        elif args.quietness != 0:
            if args.quietness == 1:
                _logger.setLevel(logging.WARNING)
            else: # anything >= 2
                _logger.setLevel(logging.ERROR)

        else: # default
            _logger.setLevel(logging.INFO)

    pool_kwargs = dict(mpi=args.mpi, processes=args.n_procs)
    pool = choose_pool(**pool_kwargs)

    main(config_file=args.config_file, pool=pool, seed=args.seed,
         overwrite=args.overwrite, _continue=args._continue)
<commit_msg>store result for done star<commit_after># Standard library
from os.path import abspath, expanduser, join
import os
import time

# Third-party
import astropy.units as u
import h5py
import numpy as np
from schwimmbad import choose_pool
from sqlalchemy import func
from thejoker.log import log as joker_logger
from thejoker.sampler import JokerParams, TheJoker, save_prior_samples
from thejoker.utils import quantity_to_hdf5
import yaml

# Project
from twoface.data import APOGEERVData
from twoface.log import log as logger
from twoface.db import Session, db_connect
from twoface.db import JokerRun, AllStar, AllVisit, StarResult, Status, AllVisitToAllStar
from twoface.config import TWOFACE_CACHE_PATH

def main(config_file, pool, seed, overwrite=False, _continue=False):
    config_file = abspath(expanduser(config_file))

    # parse config file
    with open(config_file, 'r') as f:
        config = yaml.load(f.read())

    # filename of sqlite database
    if 'database_file' not in config:
        database_file = None

    else:
        database_file = config['database_file']

    db_path = join(TWOFACE_CACHE_PATH, database_file)

    if not os.path.exists(db_path):
        raise IOError("sqlite database not found at '{}'\n Did you run scripts/initdb.py yet?"
                      .format(db_path))

    logger.debug("Connecting to sqlite database at '{}'".format(db_path))
    engine = db_connect(database_path=db_path, ensure_db_exists=False)
    session = Session()
    logger.debug("...connected!")

    # HACK: for testing
    session.query(StarResult).delete()
    session.query(JokerRun).delete()
    session.commit()

    # see if this run (by name) is in db already, if so, just grab
    _run = session.query(JokerRun).filter(JokerRun.name == config['name']).all()
    if len(_run) == 0:
        logger.info("JokerRun '{}' not found in database - creating new entry..."
                    .format(config['name']))

        # create a JokerRun for this run
        run = JokerRun()
        run.config_file = config_file
        run.name = config['name']
        run.P_min = u.Quantity(*config['hyperparams']['P_min'].split())
        run.P_max = u.Quantity(*config['hyperparams']['P_max'].split())
        run.requested_samples_per_star = int(config['hyperparams']['requested_samples_per_star'])
        run.max_prior_samples = int(config['prior']['max_samples'])
        run.prior_samples_file = join(TWOFACE_CACHE_PATH, config['prior']['samples_file'])

        if 'jitter' in config['hyperparams']:
            run.jitter = u.Quantity(*config['hyperparams']['jitter'].split())

        elif 'jitter_prior_mean' in config['hyperparams']:
            raise NotImplementedError("NOT YET SORRY")

        else:
            # no jitter specified, set to 0
            run.jitter = 0 * u.m/u.s

        # TODO: need a way to pass in velocity trend info

        # TODO: get all stars we're going to process - some filter on number of
        #   observations based on the number of linear parameters
        # TODO: need some way of specifying a constraint on the stars that are to be processed
        stars = session.query(AllStar).join(AllVisitToAllStar, AllVisit)\
                                      .group_by(AllStar.apogee_id)\
                                      .having(func.count(AllVisit.id) >= 3)\
                                      .all()

        run.stars = stars
        session.add(run)
        session.commit()

    elif len(_run) == 1:
        logger.info("JokerRun '{}' already found in database - retrieving".format(config['name']))
        run = _run[0]

    else:
        raise ValueError("Multiple JokerRun rows found for name '{}'".format(config['name']))

    # create TheJoker sampler instance
    rnd = np.random.RandomState(seed=seed)
    logger.debug("Creating TheJoker instance with {}, {}".format(rnd, pool))
    params = JokerParams(P_min=run.P_min, P_max=run.P_max, jitter=run.jitter,
                         anomaly_tol=1E-11)
    joker = TheJoker(params, random_state=rnd, pool=pool)

    # create prior samples cache, store to file and store filename in DB
    if not os.path.exists(run.prior_samples_file):
        logger.debug("Prior samples file not found - generating now...")

        prior_samples = joker.sample_prior(config['prior']['num_cache'])
        prior_units = save_prior_samples(run.prior_samples_file, prior_samples, u.km/u.s) # data in km/s
        del prior_samples
        # TODO: for large prior samples libraries, may want a way to write
        #   without storing all samples in memory

        logger.debug("...done")

    else:
        with h5py.File(run.prior_samples_file, 'r') as f:
            prior_units = [u.Unit(uu) for uu in f.attrs['units']]

    # build a query to get all stars associated with this JokerRun that need processing
    star_query = session.query(AllStar).join(StarResult, JokerRun, Status)\
                                       .filter(JokerRun.name == run.name)\
                                       .filter(Status.id == 0)

    # create a file to cache the results
    n_stars = star_query.count()
    logger.info("{} stars left to process for run '{}'".format(n_stars, run.name))

    results_filename = join(TWOFACE_CACHE_PATH, "{}.hdf5".format(run.name))

    # TODO: what should structure be? currently thinking /APOGEE_ID/key, e.g.,
    #       /2M00000222+5625359/P for period, etc.
    if not os.path.exists(results_filename):
        with h5py.File(results_filename, 'w') as f:
            pass

    else:
        with h5py.File(results_filename, 'r') as f:
            already_done = list(f.keys())

    # TODO: grab a batch of targets to process
    for star in star_query.limit(1000).all():
        if star.apogee_id in already_done:
            logger.debug("Star {} already done".format(star.apogee_id))
            continue

        logger.debug("Starting star {}".format(star.apogee_id))
        t0 = time.time()

        data = APOGEERVData.from_visits(star.visits)

        # adaptive scheme for the rejection sampling:
        n_process = 256 * run.requested_samples_per_star
        n_samples = 0 # running total of good samples returned
        start_idx = 0
        all_samples = None
        for n in range(128): # magic number we should never hit
            logger.debug("Iteration {}, using {} prior samples".format(n, n_process))

            # process n_process samples from prior cache
            samples = joker._rejection_sample_from_cache(data, n_process,
                                                         run.prior_samples_file, start_idx)

            n_survive = samples.shape[0]
            if n_survive > 1:
                n_samples += n_survive

                if all_samples is None:
                    all_samples = samples
                else:
                    all_samples = np.vstack((all_samples, samples))

            if n_samples >= run.requested_samples_per_star:
                break

            start_idx += n_process
            n_process *= 2

            if n_process > run.max_prior_samples:
                if n_process == run.max_prior_samples:
                    break

                else:
                    n_process = run.max_prior_samples

        else: # hit maxiter
            # TODO: error, should never get here
            pass

        # for now, it's sufficient to write the run results to an HDF5 file
        n = run.requested_samples_per_star
        samples_dict = joker.unpack_full_samples(all_samples[:n], data.t_offset, prior_units)

        with h5py.File(results_filename, 'r+') as f:
            g = f.create_group(star.apogee_id)
            for key in samples_dict.keys():
                quantity_to_hdf5(g, key, samples_dict[key])

        result = StarResult()
        result.star = star
        result.jokerrun = run

        if all_samples.shape[0] >= run.requested_samples_per_star:
            result.status_id = 5 # completed

        else:
            result.status_id = 3 # needs mcmc
        session.add(result)
        session.commit()

        logger.debug("...done with star {} - {:.2f} seconds".format(star.apogee_id,
                                                                    time.time()-t0))

    pool.close()
    session.close()

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

    # Define parser object
    parser = ArgumentParser(description="")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')

    oc_group = parser.add_mutually_exclusive_group()
    oc_group.add_argument("--overwrite", dest="overwrite", default=False,
                          action="store_true", help="Overwrite any existing results for "
                                                    "this JokerRun.")
    oc_group.add_argument("--continue", dest="_continue", default=False,
                          action="store_true", help="Continue the the JokerRun.")

    parser.add_argument("-s", "--seed", dest="seed", default=None, type=int,
                        help="Random number seed")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--procs", dest="n_procs", default=1,
                       type=int, help="Number of processes.")
    group.add_argument("--mpi", dest="mpi", default=False,
                       action="store_true", help="Run with MPI.")

    #
    parser.add_argument("-c", "--config", dest="config_file", required=True,
                        type=str, help="Path to config file that specifies the parameters for "
                                       "this JokerRun.")

    args = parser.parse_args()

    loggers = [joker_logger, logger]

    for _logger in loggers:
        # Set logger level based on verbose flags
        if args.verbosity != 0:
            if args.verbosity == 1:
                _logger.setLevel(logging.DEBUG)
            else: # anything >= 2
                _logger.setLevel(1)

        elif args.quietness != 0:
            if args.quietness == 1:
                _logger.setLevel(logging.WARNING)
            else: # anything >= 2
                _logger.setLevel(logging.ERROR)

        else: # default
            _logger.setLevel(logging.INFO)

    pool_kwargs = dict(mpi=args.mpi, processes=args.n_procs)
    pool = choose_pool(**pool_kwargs)

    main(config_file=args.config_file, pool=pool, seed=args.seed,
         overwrite=args.overwrite, _continue=args._continue)
<|endoftext|>


<commit_before># Standard library
from os.path import abspath, expanduser, join
import os
import time

# Third-party
import astropy.units as u
import h5py
import numpy as np
from schwimmbad import choose_pool
from sqlalchemy import func
from thejoker.log import log as joker_logger
from thejoker.sampler import JokerParams, TheJoker, save_prior_samples
from thejoker.utils import quantity_to_hdf5
import yaml

# Project
from twoface.data import APOGEERVData
from twoface.log import log as logger
from twoface.db import Session, db_connect
from twoface.db import JokerRun, AllStar, AllVisit, StarResult, Status, AllVisitToAllStar
from twoface.config import TWOFACE_CACHE_PATH

def main(config_file, pool, seed, overwrite=False, _continue=False):
    config_file = abspath(expanduser(config_file))

    # parse config file
    with open(config_file, 'r') as f:
        config = yaml.load(f.read())

    # filename of sqlite database
    if 'database_file' not in config:
        database_file = None

    else:
        database_file = config['database_file']

    db_path = join(TWOFACE_CACHE_PATH, database_file)

    if not os.path.exists(db_path):
        raise IOError("sqlite database not found at '{}'\n Did you run scripts/initdb.py yet?"
                      .format(db_path))

    logger.debug("Connecting to sqlite database at '{}'".format(db_path))
    engine = db_connect(database_path=db_path, ensure_db_exists=False)
    session = Session()
    logger.debug("...connected!")

    # HACK: for testing
    session.query(StarResult).delete()
    session.query(JokerRun).delete()
    session.commit()

    # see if this run (by name) is in db already, if so, just grab
    _run = session.query(JokerRun).filter(JokerRun.name == config['name']).all()
    if len(_run) == 0:
        logger.info("JokerRun '{}' not found in database - creating new entry..."
                    .format(config['name']))

        # create a JokerRun for this run
        run = JokerRun()
        run.config_file = config_file
        run.name = config['name']
        run.P_min = u.Quantity(*config['hyperparams']['P_min'].split())
        run.P_max = u.Quantity(*config['hyperparams']['P_max'].split())
        run.requested_samples_per_star = int(config['hyperparams']['requested_samples_per_star'])
        run.max_prior_samples = int(config['prior']['max_samples'])
        run.prior_samples_file = join(TWOFACE_CACHE_PATH, config['prior']['samples_file'])

        if 'jitter' in config['hyperparams']:
            run.jitter = u.Quantity(*config['hyperparams']['jitter'].split())

        elif 'jitter_prior_mean' in config['hyperparams']:
            raise NotImplementedError("NOT YET SORRY")

        else:
            # no jitter specified, set to 0
            run.jitter = 0 * u.m/u.s

        # TODO: need a way to pass in velocity trend info

        # TODO: get all stars we're going to process - some filter on number of
        #   observations based on the number of linear parameters
        # TODO: need some way of specifying a constraint on the stars that are to be processed
        stars = session.query(AllStar).join(AllVisitToAllStar, AllVisit)\
                                      .group_by(AllStar.apogee_id)\
                                      .having(func.count(AllVisit.id) >= 3)\
                                      .all()

        run.stars = stars
        session.add(run)
        session.commit()

    elif len(_run) == 1:
        logger.info("JokerRun '{}' already found in database - retrieving".format(config['name']))
        run = _run[0]

    else:
        raise ValueError("Multiple JokerRun rows found for name '{}'".format(config['name']))

    # create TheJoker sampler instance
    rnd = np.random.RandomState(seed=seed)
    logger.debug("Creating TheJoker instance with {}, {}".format(rnd, pool))
    params = JokerParams(P_min=run.P_min, P_max=run.P_max, jitter=run.jitter,
                         anomaly_tol=1E-11)
    joker = TheJoker(params, random_state=rnd, pool=pool)

    # create prior samples cache, store to file and store filename in DB
    if not os.path.exists(run.prior_samples_file):
        logger.debug("Prior samples file not found - generating now...")

        prior_samples = joker.sample_prior(config['prior']['num_cache'])
        prior_units = save_prior_samples(run.prior_samples_file, prior_samples, u.km/u.s) # data in km/s
        del prior_samples
        # TODO: for large prior samples libraries, may want a way to write
        #   without storing all samples in memory

        logger.debug("...done")

    else:
        with h5py.File(run.prior_samples_file, 'r') as f:
            prior_units = [u.Unit(uu) for uu in f.attrs['units']]

    # build a query to get all stars associated with this JokerRun that need processing
    star_query = session.query(AllStar).join(StarResult, JokerRun, Status)\
                                       .filter(JokerRun.name == run.name)\
                                       .filter(Status.id == 0)

    # create a file to cache the results
    n_stars = star_query.count()
    logger.info("{} stars left to process for run '{}'".format(n_stars, run.name))

    results_filename = join(TWOFACE_CACHE_PATH, "{}.hdf5".format(run.name))

    # TODO: what should structure be? currently thinking /APOGEE_ID/key, e.g.,
    #       /2M00000222+5625359/P for period, etc.
    if not os.path.exists(results_filename):
        already_done = []
        with h5py.File(results_filename, 'w') as f:
            pass

    else:
        with h5py.File(results_filename, 'r') as f:
            already_done = list(f.keys())

    # TODO: grab a batch of targets to process
    for star in star_query.limit(1000).all():
        if star.apogee_id in already_done:
            logger.debug("Star {} already done".format(star.apogee_id))
            continue

        logger.debug("Starting star {}".format(star.apogee_id))
        t0 = time.time()

        data = APOGEERVData.from_visits(star.visits)

        # adaptive scheme for the rejection sampling:
        n_process = 64 * run.requested_samples_per_star
        n_samples = 0 # running total of good samples returned
        start_idx = 0
        all_samples = None
        for n in range(128): # magic number we should never hit
            logger.debug("Iteration {}, using {} prior samples".format(n, n_process))

            # process n_process samples from prior cache
            samples = joker._rejection_sample_from_cache(data, n_process,
                                                         run.prior_samples_file, start_idx)

            n_survive = samples.shape[0]
            if n_survive > 1:
                n_samples += n_survive

                if all_samples is None:
                    all_samples = samples
                else:
                    all_samples = np.vstack((all_samples, samples))

            if n_samples >= run.requested_samples_per_star:
                break

            start_idx += n_process
            n_process *= 2

            if n_process > run.max_prior_samples:
                logger.debug("Hit max prior samples limit!!")

        else: # hit maxiter
            # TODO: error, should never get here
            pass

        # for now, it's sufficient to write the run results to an HDF5 file
        n = run.requested_samples_per_star
        samples_dict = joker.unpack_full_samples(all_samples[:n], data.t_offset, prior_units)

        with h5py.File(results_filename, 'r+') as f:
            g = f.create_group(star.apogee_id)
            for key in samples_dict.keys():
                quantity_to_hdf5(g, key, samples_dict[key])

        result = StarResult()
        result.star = star
        result.jokerrun = run

        if all_samples.shape[0] >= run.requested_samples_per_star:
            result.status_id = 5 # completed

        elif all_samples.shape[0] == 1:
            result.status_id = 3 # needs mcmc

        else:
            result.status_id = 2 # needs more samples

        session.add(result)
        session.commit()

        logger.debug("...done with star {} - {:.2f} seconds".format(star.apogee_id,
                                                                    time.time()-t0))

    pool.close()
    session.close()

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

    # Define parser object
    parser = ArgumentParser(description="")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')

    oc_group = parser.add_mutually_exclusive_group()
    oc_group.add_argument("--overwrite", dest="overwrite", default=False,
                          action="store_true", help="Overwrite any existing results for "
                                                    "this JokerRun.")
    oc_group.add_argument("--continue", dest="_continue", default=False,
                          action="store_true", help="Continue the the JokerRun.")

    parser.add_argument("-s", "--seed", dest="seed", default=None, type=int,
                        help="Random number seed")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--procs", dest="n_procs", default=1,
                       type=int, help="Number of processes.")
    group.add_argument("--mpi", dest="mpi", default=False,
                       action="store_true", help="Run with MPI.")

    #
    parser.add_argument("-c", "--config", dest="config_file", required=True,
                        type=str, help="Path to config file that specifies the parameters for "
                                       "this JokerRun.")

    args = parser.parse_args()

    loggers = [joker_logger, logger]

    for _logger in loggers:
        # Set logger level based on verbose flags
        if args.verbosity != 0:
            if args.verbosity == 1:
                _logger.setLevel(logging.DEBUG)
            else: # anything >= 2
                _logger.setLevel(1)

        elif args.quietness != 0:
            if args.quietness == 1:
                _logger.setLevel(logging.WARNING)
            else: # anything >= 2
                _logger.setLevel(logging.ERROR)

        else: # default
            _logger.setLevel(logging.INFO)

    pool_kwargs = dict(mpi=args.mpi, processes=args.n_procs)
    pool = choose_pool(**pool_kwargs)

    main(config_file=args.config_file, pool=pool, seed=args.seed,
         overwrite=args.overwrite, _continue=args._continue)
<commit_msg>stupid, forgot to break<commit_after># Standard library
from os.path import abspath, expanduser, join
import os
import time

# Third-party
import astropy.units as u
import h5py
import numpy as np
from schwimmbad import choose_pool
from sqlalchemy import func
from thejoker.log import log as joker_logger
from thejoker.sampler import JokerParams, TheJoker, save_prior_samples
from thejoker.utils import quantity_to_hdf5
import yaml

# Project
from twoface.data import APOGEERVData
from twoface.log import log as logger
from twoface.db import Session, db_connect
from twoface.db import JokerRun, AllStar, AllVisit, StarResult, Status, AllVisitToAllStar
from twoface.config import TWOFACE_CACHE_PATH

def main(config_file, pool, seed, overwrite=False, _continue=False):
    config_file = abspath(expanduser(config_file))

    # parse config file
    with open(config_file, 'r') as f:
        config = yaml.load(f.read())

    # filename of sqlite database
    if 'database_file' not in config:
        database_file = None

    else:
        database_file = config['database_file']

    db_path = join(TWOFACE_CACHE_PATH, database_file)

    if not os.path.exists(db_path):
        raise IOError("sqlite database not found at '{}'\n Did you run scripts/initdb.py yet?"
                      .format(db_path))

    logger.debug("Connecting to sqlite database at '{}'".format(db_path))
    engine = db_connect(database_path=db_path, ensure_db_exists=False)
    session = Session()
    logger.debug("...connected!")

    # HACK: for testing
    session.query(StarResult).delete()
    session.query(JokerRun).delete()
    session.commit()

    # see if this run (by name) is in db already, if so, just grab
    _run = session.query(JokerRun).filter(JokerRun.name == config['name']).all()
    if len(_run) == 0:
        logger.info("JokerRun '{}' not found in database - creating new entry..."
                    .format(config['name']))

        # create a JokerRun for this run
        run = JokerRun()
        run.config_file = config_file
        run.name = config['name']
        run.P_min = u.Quantity(*config['hyperparams']['P_min'].split())
        run.P_max = u.Quantity(*config['hyperparams']['P_max'].split())
        run.requested_samples_per_star = int(config['hyperparams']['requested_samples_per_star'])
        run.max_prior_samples = int(config['prior']['max_samples'])
        run.prior_samples_file = join(TWOFACE_CACHE_PATH, config['prior']['samples_file'])

        if 'jitter' in config['hyperparams']:
            run.jitter = u.Quantity(*config['hyperparams']['jitter'].split())

        elif 'jitter_prior_mean' in config['hyperparams']:
            raise NotImplementedError("NOT YET SORRY")

        else:
            # no jitter specified, set to 0
            run.jitter = 0 * u.m/u.s

        # TODO: need a way to pass in velocity trend info

        # TODO: get all stars we're going to process - some filter on number of
        #   observations based on the number of linear parameters
        # TODO: need some way of specifying a constraint on the stars that are to be processed
        stars = session.query(AllStar).join(AllVisitToAllStar, AllVisit)\
                                      .group_by(AllStar.apogee_id)\
                                      .having(func.count(AllVisit.id) >= 3)\
                                      .all()

        run.stars = stars
        session.add(run)
        session.commit()

    elif len(_run) == 1:
        logger.info("JokerRun '{}' already found in database - retrieving".format(config['name']))
        run = _run[0]

    else:
        raise ValueError("Multiple JokerRun rows found for name '{}'".format(config['name']))

    # create TheJoker sampler instance
    rnd = np.random.RandomState(seed=seed)
    logger.debug("Creating TheJoker instance with {}, {}".format(rnd, pool))
    params = JokerParams(P_min=run.P_min, P_max=run.P_max, jitter=run.jitter,
                         anomaly_tol=1E-11)
    joker = TheJoker(params, random_state=rnd, pool=pool)

    # create prior samples cache, store to file and store filename in DB
    if not os.path.exists(run.prior_samples_file):
        logger.debug("Prior samples file not found - generating now...")

        prior_samples = joker.sample_prior(config['prior']['num_cache'])
        prior_units = save_prior_samples(run.prior_samples_file, prior_samples, u.km/u.s) # data in km/s
        del prior_samples
        # TODO: for large prior samples libraries, may want a way to write
        #   without storing all samples in memory

        logger.debug("...done")

    else:
        with h5py.File(run.prior_samples_file, 'r') as f:
            prior_units = [u.Unit(uu) for uu in f.attrs['units']]

    # build a query to get all stars associated with this JokerRun that need processing
    star_query = session.query(AllStar).join(StarResult, JokerRun, Status)\
                                       .filter(JokerRun.name == run.name)\
                                       .filter(Status.id == 0)

    # create a file to cache the results
    n_stars = star_query.count()
    logger.info("{} stars left to process for run '{}'".format(n_stars, run.name))

    results_filename = join(TWOFACE_CACHE_PATH, "{}.hdf5".format(run.name))

    # TODO: what should structure be? currently thinking /APOGEE_ID/key, e.g.,
    #       /2M00000222+5625359/P for period, etc.
    if not os.path.exists(results_filename):
        already_done = []
        with h5py.File(results_filename, 'w') as f:
            pass

    else:
        with h5py.File(results_filename, 'r') as f:
            already_done = list(f.keys())

    # TODO: grab a batch of targets to process
    for star in star_query.limit(1000).all():
        if star.apogee_id in already_done:
            logger.debug("Star {} already done".format(star.apogee_id))
            continue

        logger.debug("Starting star {}".format(star.apogee_id))
        t0 = time.time()

        data = APOGEERVData.from_visits(star.visits)

        # adaptive scheme for the rejection sampling:
        n_process = 64 * run.requested_samples_per_star
        n_samples = 0 # running total of good samples returned
        start_idx = 0
        all_samples = None
        for n in range(128): # magic number we should never hit
            logger.debug("Iteration {}, using {} prior samples".format(n, n_process))

            # process n_process samples from prior cache
            samples = joker._rejection_sample_from_cache(data, n_process,
                                                         run.prior_samples_file, start_idx)

            n_survive = samples.shape[0]
            if n_survive > 1:
                n_samples += n_survive

                if all_samples is None:
                    all_samples = samples
                else:
                    all_samples = np.vstack((all_samples, samples))

            if n_samples >= run.requested_samples_per_star:
                break

            start_idx += n_process
            n_process *= 2

            if n_process > run.max_prior_samples:
                logger.debug("Hit max prior samples limit!!")
                break

        else: # hit maxiter
            # TODO: error, should never get here
            pass

        # for now, it's sufficient to write the run results to an HDF5 file
        n = run.requested_samples_per_star
        samples_dict = joker.unpack_full_samples(all_samples[:n], data.t_offset, prior_units)

        with h5py.File(results_filename, 'r+') as f:
            g = f.create_group(star.apogee_id)
            for key in samples_dict.keys():
                quantity_to_hdf5(g, key, samples_dict[key])

        result = StarResult()
        result.star = star
        result.jokerrun = run

        if all_samples.shape[0] >= run.requested_samples_per_star:
            result.status_id = 5 # completed

        elif all_samples.shape[0] == 1:
            result.status_id = 3 # needs mcmc

        else:
            result.status_id = 2 # needs more samples

        session.add(result)
        session.commit()

        logger.debug("...done with star {} - {:.2f} seconds".format(star.apogee_id,
                                                                    time.time()-t0))

    pool.close()
    session.close()

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

    # Define parser object
    parser = ArgumentParser(description="")

    vq_group = parser.add_mutually_exclusive_group()
    vq_group.add_argument('-v', '--verbose', action='count', default=0, dest='verbosity')
    vq_group.add_argument('-q', '--quiet', action='count', default=0, dest='quietness')

    oc_group = parser.add_mutually_exclusive_group()
    oc_group.add_argument("--overwrite", dest="overwrite", default=False,
                          action="store_true", help="Overwrite any existing results for "
                                                    "this JokerRun.")
    oc_group.add_argument("--continue", dest="_continue", default=False,
                          action="store_true", help="Continue the the JokerRun.")

    parser.add_argument("-s", "--seed", dest="seed", default=None, type=int,
                        help="Random number seed")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--procs", dest="n_procs", default=1,
                       type=int, help="Number of processes.")
    group.add_argument("--mpi", dest="mpi", default=False,
                       action="store_true", help="Run with MPI.")

    #
    parser.add_argument("-c", "--config", dest="config_file", required=True,
                        type=str, help="Path to config file that specifies the parameters for "
                                       "this JokerRun.")

    args = parser.parse_args()

    loggers = [joker_logger, logger]

    for _logger in loggers:
        # Set logger level based on verbose flags
        if args.verbosity != 0:
            if args.verbosity == 1:
                _logger.setLevel(logging.DEBUG)
            else: # anything >= 2
                _logger.setLevel(1)

        elif args.quietness != 0:
            if args.quietness == 1:
                _logger.setLevel(logging.WARNING)
            else: # anything >= 2
                _logger.setLevel(logging.ERROR)

        else: # default
            _logger.setLevel(logging.INFO)

    pool_kwargs = dict(mpi=args.mpi, processes=args.n_procs)
    pool = choose_pool(**pool_kwargs)

    main(config_file=args.config_file, pool=pool, seed=args.seed,
         overwrite=args.overwrite, _continue=args._continue)
<|endoftext|>


<commit_before>import datetime
import markdown
import re
import sys

headers = re.compile(r'(=+)(.*?)\1')
checkboxes = re.compile(r'\[[ *x]\]')
newlines = re.compile(r'\n')
strike = re.compile(r'~~(.*?)~~')

states = {' ': r'☐',
          '*': r'☑',
          'x': r'☒'}

def prettify_checkbox(match):
    return '%s' % states[match.group()[1]]

intake = sys.stdin.read()
output = headers.sub(lambda x: '#'* max(1,(6-len(x.groups()[0]))) + x.groups()[1], intake)
output = checkboxes.sub(lambda x: '* ' + x.group(), output)

marked_down = markdown.markdown('[TOC]\n\n' + output, extensions=['markdown.extensions.toc'])
marked_down = strike.sub(r'\1', marked_down)
marked_down = marked_down.replace('
\n
[', '
\n
[')
marked_down = checkboxes.sub(prettify_checkbox, marked_down)

final = ['<link rel="stylesheet" type="text/css" href="style.css"><title>',
        datetime.datetime.today().strftime('%b %d %Y'),
        '</title>',
        marked_down,
        '']
print('\n'.join(final))
<commit_msg>add in superscript, subscript<commit_after>import datetime
import markdown
import re
import sys

headers = re.compile(r'(=+)(.*?)\1')
checkboxes = re.compile(r'\[[ *x]\]')
newlines = re.compile(r'\n')
strike = re.compile(r'~~(.*?)~~')
superscript = re.compile(r'\^\{(.*?)\}')
subscript = re.compile(r'_\{(.*?)\}')
verbatim = re.compile(r"''(.*?)''")

states = {' ': r'☐',
          '*': r'☑',
          'x': r'☒'}

def prettify_checkbox(match):
    return '%s' % states[match.group()[1]]

intake = sys.stdin.read()
output = headers.sub(lambda x: '#'* max(1,(6-len(x.groups()[0]))) + x.groups()[1], intake)
output = checkboxes.sub(lambda x: '* ' + x.group(), output)

marked_down = markdown.markdown('[TOC]\n\n' + output, extensions=['markdown.extensions.toc'])
marked_down = strike.sub(r'\1', marked_down)
marked_down = superscript.sub(r'\1', marked_down)
marked_down = subscript.sub(r'\1', marked_down)
marked_down = verbatim.sub(r'\1', marked_down)
marked_down = marked_down.replace('
\n
[', '
\n
[')
marked_down = checkboxes.sub(prettify_checkbox, marked_down)

final = ['<link rel="stylesheet" type="text/css" href="style.css"><title>',
        datetime.datetime.today().strftime('%b %d %Y'),
        '</title>',
        marked_down,
        '']
print('\n'.join(final))
<|endoftext|>


<commit_before>class Node(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def add(self, value):
        if value <= self.value:
            if self.left:
                self.left.add(value)
            else:
                self.left = Node(value)
        else:
            if self.right:
                self.right.add(value)
            else:
                self.right = Node(value)

class BinaryTree:
    def __init__(self):
        self.root = None

    def add(self, value):
        if self.root:
            self.root.add(value)
        else:
            self.root = Node(value)

if __name__ == '__main__':
    bt = BinaryTree()
    bt.add(10)
    assert bt.root.value == 10
    bt.add(7)
    assert bt.root.left.value == 7
    bt.add(12)
    assert bt.root.right.value == 12<commit_msg>BinaryTree recursive contains<commit_after>class Node(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def add(self, value):
        if value <= self.value:
            if self.left:
                self.left.add(value)
            else:
                self.left = Node(value)
        else:
            if self.right:
                self.right.add(value)
            else:
                self.right = Node(value)

    def contains(self, value):
        if value == self.value:
            return True
        elif value < self.value:
            if self.left:
                return self.left.contains(value)
            else:
                return False
        else:
            if self.right:
                return self.right.contains(value)
            else:
                return False

class BinaryTree:
    def __init__(self):
        self.root = None

    def add(self, value):
        if self.root:
            self.root.add(value)
        else:
            self.root = Node(value)

    def contains(self, value):
        return self.root.contains(value)

if __name__ == '__main__':
    bt = BinaryTree()
    bt.add(10)
    assert bt.root.value == 10
    bt.add(7)
    assert bt.root.left.value == 7
    bt.add(12)
    assert bt.root.right.value == 12
    # contains
    assert bt.contains(10)
    assert bt.contains(7)
    assert bt.contains(12)
    assert not bt.contains(42)

<|endoftext|>


<commit_before>from __future__ import unicode_literals

import os
import sys
from optparse import make_option

from django.core.files.storage import FileSystemStorage
from django.core.management.base import CommandError, NoArgsCommand
from django.utils.encoding import smart_text
from django.utils.datastructures import SortedDict
from django.utils.six.moves import input

from django.contrib.staticfiles import finders, storage

class Command(NoArgsCommand):
    """
    Command that allows to copy or symlink static files from different
    locations to the settings.STATIC_ROOT.
    """
    option_list = NoArgsCommand.option_list + (
        make_option('--noinput',
            action='store_false', dest='interactive', default=True,
            help="Do NOT prompt the user for input of any kind."),
        make_option('--no-post-process',
            action='store_false', dest='post_process', default=True,
            help="Do NOT post process collected files."),
        make_option('-i', '--ignore', action='append', default=[],
            dest='ignore_patterns', metavar='PATTERN',
            help="Ignore files or directories matching this glob-style "
                "pattern. Use multiple times to ignore more."),
        make_option('-n', '--dry-run',
            action='store_true', dest='dry_run', default=False,
            help="Do everything except modify the filesystem."),
        make_option('-c', '--clear',
            action='store_true', dest='clear', default=False,
            help="Clear the existing files using the storage "
                 "before trying to copy or link the original file."),
        make_option('-l', '--link',
            action='store_true', dest='link', default=False,
            help="Create a symbolic link to each file instead of copying."),
        make_option('--no-default-ignore', action='store_false',
            dest='use_default_ignore_patterns', default=True,
            help="Don't ignore the common private glob-style patterns 'CVS', "
                "'.*' and '*~'."),
    )
    help = "Collect static files in a single location."
    requires_model_validation = False

    def __init__(self, *args, **kwargs):
        super(NoArgsCommand, self).__init__(*args, **kwargs)
        self.copied_files = []
        self.symlinked_files = []
        self.unmodified_files = []
        self.post_processed_files = []
        self.storage = storage.staticfiles_storage
        try:
            self.storage.path('')
        except NotImplementedError:
            self.local = False
        else:
            self.local = True
        # Use ints for file times (ticket #14665), if supported
        if hasattr(os, 'stat_float_times'):
            os.stat_float_times(False)

    def set_options(self, **options):
        """
        Set instance variables based on an options dict
        """
        self.interactive = options['interactive']
        self.verbosity = int(options.get('verbosity', 1))
        self.symlink = options['link']
        self.clear = options['clear']
        self.dry_run = options['dry_run']
        ignore_patterns = options['ignore_patterns']
        if options['use_default_ignore_patterns']:
            ignore_patterns += ['CVS', '.*', '*~']
        self.ignore_patterns = list(set(ignore_patterns))
        self.post_process = options['post_process']

    def collect(self):
        """
        Perform the bulk of the work of collectstatic.

        Split off from handle_noargs() to facilitate testing.
        """
        if self.symlink:
            if sys.platform == 'win32':
                raise CommandError("Symlinking is not supported by this "
                                   "platform (%s)." % sys.platform)
            if not self.local:
                raise CommandError("Can't symlink to a remote destination.")

        if self.clear:
            self.clear_dir('')

        if self.symlink:
            handler = self.link_file
        else:
            handler = self.copy_file

        found_files = SortedDict()
        for finder in finders.get_finders():
            for path, storage in finder.list(self.ignore_patterns):
                # Prefix the relative path if the source storage contains it
                if getattr(storage, 'prefix', None):
                    prefixed_path = os.path.join(storage.prefix, path)
                else:
                    prefixed_path = path

                if prefixed_path not in found_files:
                    found_files[prefixed_path] = (storage, path)
                    handler(path, prefixed_path, storage)

        # Here we check if the storage backend has a post_process
        # method and pass it the list of modified files.
        if self.post_process and hasattr(self.storage, 'post_process'):
            processor = self.storage.post_process(found_files,
                                                  dry_run=self.dry_run)
            for original_path, processed_path, processed in processor:
                if processed:
                    self.log("Post-processed '%s' as '%s" %
                             (original_path, processed_path), level=1)
                    self.post_processed_files.append(original_path)
                else:
                    self.log("Skipped post-processing '%s'" % original_path)

        return {
            'modified': self.copied_files + self.symlinked_files,
            'unmodified': self.unmodified_files,
            'post_processed': self.post_processed_files,
        }

    def handle_noargs(self, **options):
        self.set_options(**options)
        # Warn before doing anything more.
        if (isinstance(self.storage, FileSystemStorage) and
                self.storage.location):
            destination_path = self.storage.location
            destination_display = ':\n\n    %s' % destination_path
        else:
            destination_path = None
            destination_display = '.'

        if self.clear:
            clear_display = 'This will DELETE EXISTING FILES!'
        else:
            clear_display = 'This will overwrite existing files!'

        if self.interactive:
            confirm = input("""
You have requested to collect static files at the destination
location as specified in your settings%s

%s
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: """
% (destination_display, clear_display))
            if confirm != 'yes':
                raise CommandError("Collecting static files cancelled.")

        collected = self.collect()
        modified_count = len(collected['modified'])
        unmodified_count = len(collected['unmodified'])
        post_processed_count = len(collected['post_processed'])

        if self.verbosity >= 1:
            template = ("\n%(modified_count)s %(identifier)s %(action)s"
                        "%(destination)s%(unmodified)s%(post_processed)s.\n")
            summary = template % {
                'modified_count': modified_count,
                'identifier': 'static file' + (modified_count != 1 and 's' or ''),
                'action': self.symlink and 'symlinked' or 'copied',
                'destination': (destination_path and " to '%s'"
                                % destination_path or ''),
                'unmodified': (collected['unmodified'] and ', %s unmodified'
                               % unmodified_count or ''),
                'post_processed': (collected['post_processed'] and
                                   ', %s post-processed'
                                   % post_processed_count or ''),
            }
            self.stdout.write(summary)

    def log(self, msg, level=2):
        """
        Small log helper
        """
        if self.verbosity >= level:
            self.stdout.write(msg)

    def clear_dir(self, path):
        """
        Deletes the given relative path using the destination storage backend.
        """
        dirs, files = self.storage.listdir(path)
        for f in files:
            fpath = os.path.join(path, f)
            if self.dry_run:
                self.log("Pretending to delete '%s'" %
                         smart_text(fpath), level=1)
            else:
                self.log("Deleting '%s'" % smart_text(fpath), level=1)
                self.storage.delete(fpath)
        for d in dirs:
            self.clear_dir(os.path.join(path, d))

    def delete_file(self, path, prefixed_path, source_storage):
        """
        Checks if the target file should be deleted if it already exists
        """
        if self.storage.exists(prefixed_path):
            try:
                # When was the target file modified last time?
                target_last_modified = \
                    self.storage.modified_time(prefixed_path)
            except (OSError, NotImplementedError, AttributeError):
                # The storage doesn't support ``modified_time`` or failed
                pass
            else:
                try:
                    # When was the source file modified last time?
                    source_last_modified = source_storage.modified_time(path)
                except (OSError, NotImplementedError, AttributeError):
                    pass
                else:
                    # The full path of the target file
                    if self.local:
                        full_path = self.storage.path(prefixed_path)
                    else:
                        full_path = None
                    # Skip the file if the source file is younger
                    if target_last_modified >= source_last_modified:
                        if not ((self.symlink and full_path
                                 and not os.path.islink(full_path)) or
                                (not self.symlink and full_path
                                 and os.path.islink(full_path))):
                            if prefixed_path not in self.unmodified_files:
                                self.unmodified_files.append(prefixed_path)
                            self.log("Skipping '%s' (not modified)" % path)
                            return False
            # Then delete the existing file if really needed
            if self.dry_run:
                self.log("Pretending to delete '%s'" % path)
            else:
                self.log("Deleting '%s'" % path)
                self.storage.delete(prefixed_path)
        return True

    def link_file(self, path, prefixed_path, source_storage):
        """
        Attempt to link ``path``
        """
        # Skip this file if it was already copied earlier
        if prefixed_path in self.symlinked_files:
            return self.log("Skipping '%s' (already linked earlier)" % path)
        # Delete the target file if needed or break
        if not self.delete_file(path, prefixed_path, source_storage):
            return
        # The full path of the source file
        source_path = source_storage.path(path)
        # Finally link the file
        if self.dry_run:
            self.log("Pretending to link '%s'" % source_path, level=1)
        else:
            self.log("Linking '%s'" % source_path, level=1)
            full_path = self.storage.path(prefixed_path)
            try:
                os.makedirs(os.path.dirname(full_path))
            except OSError:
                pass
            os.symlink(source_path, full_path)
        if prefixed_path not in self.symlinked_files:
            self.symlinked_files.append(prefixed_path)

    def copy_file(self, path, prefixed_path, source_storage):
        """
        Attempt to copy ``path`` with storage
        """
        # Skip this file if it was already copied earlier
        if prefixed_path in self.copied_files:
            return self.log("Skipping '%s' (already copied earlier)" % path)
        # Delete the target file if needed or break
        if not self.delete_file(path, prefixed_path, source_storage):
            return
        # The full path of the source file
        source_path = source_storage.path(path)
        # Finally start copying
        if self.dry_run:
            self.log("Pretending to copy '%s'" % source_path, level=1)
        else:
            self.log("Copying '%s'" % source_path, level=1)
            if self.local:
                full_path = self.storage.path(prefixed_path)
                try:
                    os.makedirs(os.path.dirname(full_path))
                except OSError:
                    pass
            with source_storage.open(path) as source_file:
                self.storage.save(prefixed_path, source_file)
        if not prefixed_path in self.copied_files:
            self.copied_files.append(prefixed_path)
<commit_msg>Update django/contrib/staticfiles/management/commands/collectstatic.py<commit_after>from __future__ import unicode_literals

import os
import sys
from optparse import make_option

from django.core.files.storage import FileSystemStorage
from django.core.management.base import CommandError, NoArgsCommand
from django.utils.encoding import smart_text
from django.utils.datastructures import SortedDict
from django.utils.six.moves import input

from django.contrib.staticfiles import finders, storage

class Command(NoArgsCommand):
    """
    Command that allows to copy or symlink static files from different
    locations to the settings.STATIC_ROOT.
    """
    option_list = NoArgsCommand.option_list + (
        make_option('--noinput',
            action='store_false', dest='interactive', default=True,
            help="Do NOT prompt the user for input of any kind."),
        make_option('--no-post-process',
            action='store_false', dest='post_process', default=True,
            help="Do NOT post process collected files."),
        make_option('-i', '--ignore', action='append', default=[],
            dest='ignore_patterns', metavar='PATTERN',
            help="Ignore files or directories matching this glob-style "
                "pattern. Use multiple times to ignore more."),
        make_option('-n', '--dry-run',
            action='store_true', dest='dry_run', default=False,
            help="Do everything except modify the filesystem."),
        make_option('-c', '--clear',
            action='store_true', dest='clear', default=False,
            help="Clear the existing files using the storage "
                 "before trying to copy or link the original file."),
        make_option('-l', '--link',
            action='store_true', dest='link', default=False,
            help="Create a symbolic link to each file instead of copying."),
        make_option('--no-default-ignore', action='store_false',
            dest='use_default_ignore_patterns', default=True,
            help="Don't ignore the common private glob-style patterns 'CVS', "
                "'.*' and '*~'."),
    )
    help = "Collect static files in a single location."
    requires_model_validation = False

    def __init__(self, *args, **kwargs):
        super(NoArgsCommand, self).__init__(*args, **kwargs)
        self.copied_files = []
        self.symlinked_files = []
        self.unmodified_files = []
        self.post_processed_files = []
        self.storage = storage.staticfiles_storage
        try:
            self.storage.path('')
        except NotImplementedError:
            self.local = False
        else:
            self.local = True
        # Use ints for file times (ticket #14665), if supported
        if hasattr(os, 'stat_float_times'):
            os.stat_float_times(False)

    def set_options(self, **options):
        """
        Set instance variables based on an options dict
        """
        self.interactive = options['interactive']
        self.verbosity = int(options.get('verbosity', 1))
        self.symlink = options['link']
        self.clear = options['clear']
        self.dry_run = options['dry_run']
        ignore_patterns = options['ignore_patterns']
        if options['use_default_ignore_patterns']:
            ignore_patterns += ['CVS', '.*', '*~']
        self.ignore_patterns = list(set(ignore_patterns))
        self.post_process = options['post_process']

    def collect(self):
        """
        Perform the bulk of the work of collectstatic.

        Split off from handle_noargs() to facilitate testing.
        """
        if self.symlink:
            if sys.platform == 'win32':
                raise CommandError("Symlinking is not supported by this "
                                   "platform (%s)." % sys.platform)
            if not self.local:
                raise CommandError("Can't symlink to a remote destination.")

        if self.clear:
            self.clear_dir('')

        if self.symlink:
            handler = self.link_file
        else:
            handler = self.copy_file

        found_files = SortedDict()
        for finder in finders.get_finders():
            for path, storage in finder.list(self.ignore_patterns):
                # Prefix the relative path if the source storage contains it
                if getattr(storage, 'prefix', None):
                    prefixed_path = os.path.join(storage.prefix, path)
                else:
                    prefixed_path = path

                if prefixed_path not in found_files:
                    found_files[prefixed_path] = (storage, path)
                    handler(path, prefixed_path, storage)

        # Here we check if the storage backend has a post_process
        # method and pass it the list of modified files.
        if self.post_process and hasattr(self.storage, 'post_process'):
            processor = self.storage.post_process(found_files,
                                                  dry_run=self.dry_run)
            for original_path, processed_path, processed in processor:
                if processed:
                    self.log("Post-processed '%s' as '%s'" %
                             (original_path, processed_path), level=1)
                    self.post_processed_files.append(original_path)
                else:
                    self.log("Skipped post-processing '%s'" % original_path)

        return {
            'modified': self.copied_files + self.symlinked_files,
            'unmodified': self.unmodified_files,
            'post_processed': self.post_processed_files,
        }

    def handle_noargs(self, **options):
        self.set_options(**options)
        # Warn before doing anything more.
        if (isinstance(self.storage, FileSystemStorage) and
                self.storage.location):
            destination_path = self.storage.location
            destination_display = ':\n\n    %s' % destination_path
        else:
            destination_path = None
            destination_display = '.'

        if self.clear:
            clear_display = 'This will DELETE EXISTING FILES!'
        else:
            clear_display = 'This will overwrite existing files!'

        if self.interactive:
            confirm = input("""
You have requested to collect static files at the destination
location as specified in your settings%s

%s
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: """
% (destination_display, clear_display))
            if confirm != 'yes':
                raise CommandError("Collecting static files cancelled.")

        collected = self.collect()
        modified_count = len(collected['modified'])
        unmodified_count = len(collected['unmodified'])
        post_processed_count = len(collected['post_processed'])

        if self.verbosity >= 1:
            template = ("\n%(modified_count)s %(identifier)s %(action)s"
                        "%(destination)s%(unmodified)s%(post_processed)s.\n")
            summary = template % {
                'modified_count': modified_count,
                'identifier': 'static file' + (modified_count != 1 and 's' or ''),
                'action': self.symlink and 'symlinked' or 'copied',
                'destination': (destination_path and " to '%s'"
                                % destination_path or ''),
                'unmodified': (collected['unmodified'] and ', %s unmodified'
                               % unmodified_count or ''),
                'post_processed': (collected['post_processed'] and
                                   ', %s post-processed'
                                   % post_processed_count or ''),
            }
            self.stdout.write(summary)

    def log(self, msg, level=2):
        """
        Small log helper
        """
        if self.verbosity >= level:
            self.stdout.write(msg)

    def clear_dir(self, path):
        """
        Deletes the given relative path using the destination storage backend.
        """
        dirs, files = self.storage.listdir(path)
        for f in files:
            fpath = os.path.join(path, f)
            if self.dry_run:
                self.log("Pretending to delete '%s'" %
                         smart_text(fpath), level=1)
            else:
                self.log("Deleting '%s'" % smart_text(fpath), level=1)
                self.storage.delete(fpath)
        for d in dirs:
            self.clear_dir(os.path.join(path, d))

    def delete_file(self, path, prefixed_path, source_storage):
        """
        Checks if the target file should be deleted if it already exists
        """
        if self.storage.exists(prefixed_path):
            try:
                # When was the target file modified last time?
                target_last_modified = \
                    self.storage.modified_time(prefixed_path)
            except (OSError, NotImplementedError, AttributeError):
                # The storage doesn't support ``modified_time`` or failed
                pass
            else:
                try:
                    # When was the source file modified last time?
                    source_last_modified = source_storage.modified_time(path)
                except (OSError, NotImplementedError, AttributeError):
                    pass
                else:
                    # The full path of the target file
                    if self.local:
                        full_path = self.storage.path(prefixed_path)
                    else:
                        full_path = None
                    # Skip the file if the source file is younger
                    if target_last_modified >= source_last_modified:
                        if not ((self.symlink and full_path
                                 and not os.path.islink(full_path)) or
                                (not self.symlink and full_path
                                 and os.path.islink(full_path))):
                            if prefixed_path not in self.unmodified_files:
                                self.unmodified_files.append(prefixed_path)
                            self.log("Skipping '%s' (not modified)" % path)
                            return False
            # Then delete the existing file if really needed
            if self.dry_run:
                self.log("Pretending to delete '%s'" % path)
            else:
                self.log("Deleting '%s'" % path)
                self.storage.delete(prefixed_path)
        return True

    def link_file(self, path, prefixed_path, source_storage):
        """
        Attempt to link ``path``
        """
        # Skip this file if it was already copied earlier
        if prefixed_path in self.symlinked_files:
            return self.log("Skipping '%s' (already linked earlier)" % path)
        # Delete the target file if needed or break
        if not self.delete_file(path, prefixed_path, source_storage):
            return
        # The full path of the source file
        source_path = source_storage.path(path)
        # Finally link the file
        if self.dry_run:
            self.log("Pretending to link '%s'" % source_path, level=1)
        else:
            self.log("Linking '%s'" % source_path, level=1)
            full_path = self.storage.path(prefixed_path)
            try:
                os.makedirs(os.path.dirname(full_path))
            except OSError:
                pass
            os.symlink(source_path, full_path)
        if prefixed_path not in self.symlinked_files:
            self.symlinked_files.append(prefixed_path)

    def copy_file(self, path, prefixed_path, source_storage):
        """
        Attempt to copy ``path`` with storage
        """
        # Skip this file if it was already copied earlier
        if prefixed_path in self.copied_files:
            return self.log("Skipping '%s' (already copied earlier)" % path)
        # Delete the target file if needed or break
        if not self.delete_file(path, prefixed_path, source_storage):
            return
        # The full path of the source file
        source_path = source_storage.path(path)
        # Finally start copying
        if self.dry_run:
            self.log("Pretending to copy '%s'" % source_path, level=1)
        else:
            self.log("Copying '%s'" % source_path, level=1)
            if self.local:
                full_path = self.storage.path(prefixed_path)
                try:
                    os.makedirs(os.path.dirname(full_path))
                except OSError:
                    pass
            with source_storage.open(path) as source_file:
                self.storage.save(prefixed_path, source_file)
        if not prefixed_path in self.copied_files:
            self.copied_files.append(prefixed_path)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function
import errno
import os
import re
import shutil
import string
import sys
import tvdb_api
from tvdb_exceptions import tvdb_error, tvdb_shownotfound

CONFIG_FILE = '/etc/sabnzbd/scripts.conf'
EXTENSIONS = ['avi', 'm4v', 'mkv', 'mp4']
PATTERNS = [
    re.compile('^(.*)[\._]s(\d+)e(\d+)(.*)', re.I),
    re.compile('^(.*)(\d)x(\d+)(.*)', re.I),
]

def get_config(cfgfile):
    config = {}
    with open(cfgfile) as fh:
        for line in fh:
            k, v = re.split('\s+=\s+', line, 1)
            config[k] = v.strip()
    return config

CONFIG = get_config(CONFIG_FILE)

def get_canonical_name(name):
    """
    Translates a dotted filename into the real name with spaces while
    preserving acronyms.
    """
    if not name.endswith('.'):
        name += '.'

    regex = '((?:[A-Z]\.)+)'
    acronyms = re.findall(regex, name)
    parts = re.split(regex, name)

    if not parts:
        return name

    canonical_name = []
    for part in parts:
        if part in acronyms:
            part = re.sub('\.$', '', part)
        else:
            part = part.replace('.', ' ')
        canonical_name.append(part.strip())
    canonical_name = ' '.join(canonical_name)

    # look up TheTVDB for the show name
    canonical_name = get_tvdb_name(canonical_name)

    # replace trailing dots
    canonical_name = re.sub('\.+$', '', canonical_name)

    return canonical_name

def get_tvdb_name(show):
    tvdb = tvdb_api.Tvdb()
    try:
        return tvdb[show]['seriesname']
    except (tvdb_error, tvdb_shownotfound):
        return show

def mkdirp(path, mode=0755):
    try:
        os.makedirs(path, mode=mode)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

    for r, _, _ in os.walk(path):
        os.chmod(os.path.join(path, r), mode)

def main(job_dir, nzb, clean, index_num, category, group, status):
    # Loop through the files in ``job_dir``, trying to find one which
    # matches PATTERN. Also build a list of files and their sizes.
    #
    # If no files match PATTERN but we do find one or more files with
    # a suitable extension we assume the largest file is the one we
    # want. We take this file and use the ``job_dir`` as the base for
    # the new file name.
    match = None
    files_by_size = []
    for fn in os.listdir(job_dir):
        ext = fn.rsplit('.', 1)[-1]
        if ext not in EXTENSIONS:
            continue

        abs = os.path.abspath(os.path.join(job_dir, fn))
        files_by_size.append((fn, os.path.getsize(abs)))
        match = None

        for pattern in PATTERNS:
            match = pattern.match(fn)
            if match:
                break

        if match:
            break

    files_by_size = sorted(files_by_size, key=lambda x: x[1])

    if not match:
        # no files at all? just exit
        if not files_by_size:
            print("No suitable files found")
            sys.exit()

        # found a potential match
        fn = files_by_size[0][0]
        match = None

        for pattern in PATTERNS:
            match = pattern.match(os.path.basename(job_dir))
            if match:
                break

        if not match:
            print("No suitable files found")
            sys.exit()

    fn = os.path.join(job_dir, fn)
    ext = fn.rsplit('.', 1)[-1]

    show_name, season_num, episode_num, _ = match.groups()

    dotted_name = string.capwords(show_name, '.').replace('_', '.')
    canonical_name = get_canonical_name(dotted_name)
    dotted_name = canonical_name.replace(' ', '.')
    show_dir = os.path.join(CONFIG['tv_directory'], canonical_name)
    season_dir = os.path.join(show_dir, 'Season %s' % int(match.group(2)))

    season_num = match.group(2).zfill(2)
    episode_num = match.group(3).zfill(2)

    vfile = "%s.s%se%s.%s" % (dotted_name, season_num, episode_num, ext)
    final_name = os.path.join(season_dir, vfile)
    if os.path.isfile(final_name):
        print("File already exists: %s" % final_name)
        sys.exit(1)

    # create the full path to the new location, move the file over
    # and remove the old directory
    mkdirp(show_dir)
    mkdirp(season_dir)
    shutil.move(os.path.join(job_dir, fn), final_name)
    shutil.rmtree(job_dir)

    print("New file: %s (%s)" % (final_name, os.path.basename(fn)))

    # try to remove the empty category directories
    parent_dirname = os.path.dirname(os.path.abspath(job_dir))
    parent_basename = os.path.basename(parent_dirname)
    if parent_basename.lower() == category.lower():
        try:
            os.rmdir(parent_dirname)
            print("Removed empty directory: %s" % parent_dirname)
        except OSError:
            print("Skipped non-empty directory: %s" % parent_dirname)
            pass

if __name__ == '__main__':
    main(*sys.argv[1:])
<commit_msg>Always strip trailing non-alphanumeric characters.<commit_after>#!/usr/bin/env python
from __future__ import print_function
import errno
import os
import re
import shutil
import string
import sys
import tvdb_api
from tvdb_exceptions import tvdb_error, tvdb_shownotfound

CONFIG_FILE = '/etc/sabnzbd/scripts.conf'
EXTENSIONS = ['avi', 'm4v', 'mkv', 'mp4']
PATTERNS = [
    re.compile('^(.*)[\._]s(\d+)e(\d+)(.*)', re.I),
    re.compile('^(.*)(\d)x(\d+)(.*)', re.I),
]

def get_config(cfgfile):
    config = {}
    with open(cfgfile) as fh:
        for line in fh:
            k, v = re.split('\s+=\s+', line, 1)
            config[k] = v.strip()
    return config

CONFIG = get_config(CONFIG_FILE)

def get_canonical_name(name):
    """
    Translates a dotted filename into the real name with spaces while
    preserving acronyms.
    """
    if not name.endswith('.'):
        name += '.'

    regex = '((?:[A-Z]\.)+)'
    acronyms = re.findall(regex, name)
    parts = re.split(regex, name)

    if not parts:
        return name

    canonical_name = []
    for part in parts:
        if part in acronyms:
            part = re.sub('\.$', '', part)
        else:
            part = part.replace('.', ' ')
        canonical_name.append(part.strip())
    canonical_name = ' '.join(canonical_name)

    # look up TheTVDB for the show name
    canonical_name = get_tvdb_name(canonical_name)

    # strip trailing non-alphanumeric characters
    canonical_name = re.sub('[^0-9a-zA-Z]+$', '', canonical_name)

    return canonical_name

def get_tvdb_name(show):
    tvdb = tvdb_api.Tvdb()
    try:
        return tvdb[show]['seriesname']
    except (tvdb_error, tvdb_shownotfound):
        return show

def mkdirp(path, mode=0755):
    try:
        os.makedirs(path, mode=mode)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

    for r, _, _ in os.walk(path):
        os.chmod(os.path.join(path, r), mode)

def main(job_dir, nzb, clean, index_num, category, group, status):
    # Loop through the files in ``job_dir``, trying to find one which
    # matches PATTERN. Also build a list of files and their sizes.
    #
    # If no files match PATTERN but we do find one or more files with
    # a suitable extension we assume the largest file is the one we
    # want. We take this file and use the ``job_dir`` as the base for
    # the new file name.
    match = None
    files_by_size = []
    for fn in os.listdir(job_dir):
        ext = fn.rsplit('.', 1)[-1]
        if ext not in EXTENSIONS:
            continue

        abs = os.path.abspath(os.path.join(job_dir, fn))
        files_by_size.append((fn, os.path.getsize(abs)))
        match = None

        for pattern in PATTERNS:
            match = pattern.match(fn)
            if match:
                break

        if match:
            break

    files_by_size = sorted(files_by_size, key=lambda x: x[1])

    if not match:
        # no files at all? just exit
        if not files_by_size:
            print("No suitable files found")
            sys.exit()

        # found a potential match
        fn = files_by_size[0][0]
        match = None

        for pattern in PATTERNS:
            match = pattern.match(os.path.basename(job_dir))
            if match:
                break

        if not match:
            print("No suitable files found")
            sys.exit()

    fn = os.path.join(job_dir, fn)
    ext = fn.rsplit('.', 1)[-1]

    show_name, season_num, episode_num, _ = match.groups()

    dotted_name = string.capwords(show_name, '.').replace('_', '.')
    canonical_name = get_canonical_name(dotted_name)
    dotted_name = canonical_name.replace(' ', '.')
    show_dir = os.path.join(CONFIG['tv_directory'], canonical_name)
    season_dir = os.path.join(show_dir, 'Season %s' % int(match.group(2)))

    season_num = match.group(2).zfill(2)
    episode_num = match.group(3).zfill(2)

    vfile = "%s.s%se%s.%s" % (dotted_name, season_num, episode_num, ext)
    final_name = os.path.join(season_dir, vfile)
    if os.path.isfile(final_name):
        print("File already exists: %s" % final_name)
        sys.exit(1)

    # create the full path to the new location, move the file over
    # and remove the old directory
    mkdirp(show_dir)
    mkdirp(season_dir)
    shutil.move(os.path.join(job_dir, fn), final_name)
    shutil.rmtree(job_dir)

    print("New file: %s (%s)" % (final_name, os.path.basename(fn)))

    # try to remove the empty category directories
    parent_dirname = os.path.dirname(os.path.abspath(job_dir))
    parent_basename = os.path.basename(parent_dirname)
    if parent_basename.lower() == category.lower():
        try:
            os.rmdir(parent_dirname)
            print("Removed empty directory: %s" % parent_dirname)
        except OSError:
            print("Skipped non-empty directory: %s" % parent_dirname)
            pass

if __name__ == '__main__':
    main(*sys.argv[1:])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Author: Bo Maryniuk <bo@suse.de>
#
# Copyright 2017 SUSE LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import logging
import importlib
import yaml
import fnmatch
import subprocess
import json

from salt.exceptions import LoaderError, CommandExecutionError
try:
    import ansible
    import ansible.constants
    import ansible.modules
    from ansible.plugins import module_loader
except ImportError:
    ansible = None

__virtualname__ = 'ansible'
log = logging.getLogger(__name__)

class AnsibleModuleResolver(object):
    '''
    This class is to resolve all available modules in Ansible.
    '''
    def __init__(self, opts):
        self._opts = opts
        self._modules_map = {}

    def _get_modules_map(self, path=None):
        '''
        Get installed Ansible modules
        :return:
        '''
        paths = {}
        root = ansible.modules.__path__[0]
        if not path:
            path = root
        for p_el in os.listdir(path):
            p_el_path = os.path.join(path, p_el)
            if os.path.islink(p_el_path): continue
            if os.path.isdir(p_el_path):
                paths.update(self._get_modules_map(p_el_path))
            else:
                if (any(p_el.startswith(elm) for elm in ['__', '.']) or
                        not p_el.endswith('.py') or
                        p_el in ansible.constants.IGNORE_FILES):
                    continue
                p_el_path = p_el_path.replace(root, '').split('.')[0]
                als_name = p_el_path.replace('.', '').replace('/', '', 1).replace('/', '.')
                paths[als_name] = p_el_path

        return paths

    def load_module(self, module):
        '''
        Introspect Ansible module.

        :param module:
        :return:
        '''
        m_ref = self._modules_map.get(module)
        if m_ref is None:
            raise LoaderError('Module "{0}" was not found'.format(module))
        mod = importlib.import_module('ansible.modules{0}'.format(
            '.'.join([elm.split('.')[0] for elm in m_ref.split(os.path.sep)])))

        return mod

    def get_modules_list(self, pattern=None):
        '''
        Return module map references.
        :return:
        '''
        if pattern and '*' not in pattern:
            pattern = '*{0}*'.format(pattern)
        modules = []
        for m_name, m_path in self._modules_map.items():
            m_path = m_path.split('.')[0]
            m_name = '.'.join([elm for elm in m_path.split(os.path.sep) if elm])
            if pattern and fnmatch.fnmatch(m_name, pattern) or not pattern:
                modules.append(m_name)
        return sorted(modules)

    def resolve(self):
        log.debug('Resolving Ansible modules')
        self._modules_map = self._get_modules_map()
        return self

    def install(self):
        log.debug('Installing Ansible modules')
        return self

class AnsibleModuleCaller(object):
    def __init__(self, resolver):
        self._resolver = resolver

    def call(self, module, *args, **kwargs):
        '''
        Call an Ansible module by invoking it.
        :param module: the name of the module.
        :param args: Arguments to the module
        :param kwargs: keywords to the module
        :return:
        '''

        module = self._resolver.load_module(module)
        if not hasattr(module, 'main'):
            raise CommandExecutionError('This module is not callable '
                                        '(see "ansible.help {0}")'.format(module.__name__.replace('ansible.modules.',
                                                                                '')))
        if args:
            kwargs['_raw_params'] = ' '.join(args)
        js_args = '{{"ANSIBLE_MODULE_ARGS": {args}}}'.format(args=json.dumps(kwargs))
        js_out = subprocess.Popen(["echo", "{0}".format(js_args)], stdout=subprocess.PIPE)
        md_exc = subprocess.Popen(['python', module.__file__],
                                  stdin=js_out.stdout, stdout=subprocess.PIPE)
        js_out.stdout.close()
        js_out, js_err = md_exc.communicate()

        try:
            out = json.loads(js_out)
        except ValueError as ex:
            out = {'Error': (js_err and (js_err + '.') or str(ex))}
            if js_out:
                out['Given JSON output'] = js_out
            return out

        if 'invocation' in out:
            del out['invocation']

        return out

_resolver = None
_caller = None

def _set_callables(modules):
    '''
    Set all Ansible modules callables
    :return:
    '''
    def _mkf(cmd_name, doc):
        '''
        Create a Salt function for the Ansible module.
        '''
        def _cmd(*args, **kw):
            '''
            Call an Ansible module as a function from the Salt.
            '''
            kwargs = {}
            if kw.get('__pub_arg'):
                for _kw in kw.get('__pub_arg', []):
                    if isinstance(_kw, dict):
                        kwargs = _kw
                        break

            global _caller
            return _caller.call(cmd_name, *args, **kwargs)
        _cmd.__doc__ = doc
        return _cmd

    for mod in modules:
        setattr(sys.modules[__name__], mod, _mkf(mod, 'Available'))

def __virtual__():
    '''
    Ansible module caller.
    :return:
    '''
    ret = ansible is not None
    msg = not ret and "Ansible is not installed on this system" or None
    if msg:
        log.warning(msg)
    else:
        global _resolver
        global _caller
        _resolver = AnsibleModuleResolver(__opts__).resolve().install()
        _caller = AnsibleModuleCaller(_resolver)
    _set_callables(list())

    return ret, msg

def help(module=None, *args):
    '''
    Display help on Ansible standard module.

    :param module:
    :return:
    '''
    if not module:
        raise CommandExecutionError('Please tell me what module you want to have helped with. '
                                    'Or call "ansible.list" to know what is available.')
    try:
        module = _resolver.load_module(module)
    except (ImportError, LoaderError) as err:
        raise CommandExecutionError('Module "{0}" is currently not functional on your system.'.format(module))

    doc = {}
    ret = {}
    for docset in module.DOCUMENTATION.split('---'):
        try:
            docset = yaml.load(docset)
            if docset:
                doc.update(docset)
        except Exception as err:
            log.error("Error parsing doc section: {0}".format(err))
    if not args:
        if 'description' in doc:
            description = doc.get('description') or ''
            del doc['description']
            ret['Description'] = description
        ret['Available sections on module "{}"'.format(module.__name__.replace('ansible.modules.', ''))] = doc.keys()
    else:
        for arg in args:
            info = doc.get(arg)
            if info is not None:
                ret[arg] = info

    return ret

def list(pattern=None):
    '''
    Lists available modules.
    :return:
    '''
    return _resolver.get_modules_list(pattern=pattern)
<commit_msg>Lintfix: W1699<commit_after># -*- coding: utf-8 -*-
#
# Author: Bo Maryniuk <bo@suse.de>
#
# Copyright 2017 SUSE LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
import os
import sys
import logging
import importlib
import yaml
import fnmatch
import subprocess
import json

from salt.exceptions import LoaderError, CommandExecutionError
try:
    import ansible
    import ansible.constants
    import ansible.modules
    from ansible.plugins import module_loader
except ImportError:
    ansible = None

__virtualname__ = 'ansible'
log = logging.getLogger(__name__)

class AnsibleModuleResolver(object):
    '''
    This class is to resolve all available modules in Ansible.
    '''
    def __init__(self, opts):
        self._opts = opts
        self._modules_map = {}

    def _get_modules_map(self, path=None):
        '''
        Get installed Ansible modules
        :return:
        '''
        paths = {}
        root = ansible.modules.__path__[0]
        if not path:
            path = root
        for p_el in os.listdir(path):
            p_el_path = os.path.join(path, p_el)
            if os.path.islink(p_el_path): continue
            if os.path.isdir(p_el_path):
                paths.update(self._get_modules_map(p_el_path))
            else:
                if (any(p_el.startswith(elm) for elm in ['__', '.']) or
                        not p_el.endswith('.py') or
                        p_el in ansible.constants.IGNORE_FILES):
                    continue
                p_el_path = p_el_path.replace(root, '').split('.')[0]
                als_name = p_el_path.replace('.', '').replace('/', '', 1).replace('/', '.')
                paths[als_name] = p_el_path

        return paths

    def load_module(self, module):
        '''
        Introspect Ansible module.

        :param module:
        :return:
        '''
        m_ref = self._modules_map.get(module)
        if m_ref is None:
            raise LoaderError('Module "{0}" was not found'.format(module))
        mod = importlib.import_module('ansible.modules{0}'.format(
            '.'.join([elm.split('.')[0] for elm in m_ref.split(os.path.sep)])))

        return mod

    def get_modules_list(self, pattern=None):
        '''
        Return module map references.
        :return:
        '''
        if pattern and '*' not in pattern:
            pattern = '*{0}*'.format(pattern)
        modules = []
        for m_name, m_path in self._modules_map.items():
            m_path = m_path.split('.')[0]
            m_name = '.'.join([elm for elm in m_path.split(os.path.sep) if elm])
            if pattern and fnmatch.fnmatch(m_name, pattern) or not pattern:
                modules.append(m_name)
        return sorted(modules)

    def resolve(self):
        log.debug('Resolving Ansible modules')
        self._modules_map = self._get_modules_map()
        return self

    def install(self):
        log.debug('Installing Ansible modules')
        return self

class AnsibleModuleCaller(object):
    def __init__(self, resolver):
        self._resolver = resolver

    def call(self, module, *args, **kwargs):
        '''
        Call an Ansible module by invoking it.
        :param module: the name of the module.
        :param args: Arguments to the module
        :param kwargs: keywords to the module
        :return:
        '''

        module = self._resolver.load_module(module)
        if not hasattr(module, 'main'):
            raise CommandExecutionError('This module is not callable '
                                        '(see "ansible.help {0}")'.format(module.__name__.replace('ansible.modules.',
                                                                                '')))
        if args:
            kwargs['_raw_params'] = ' '.join(args)
        js_args = '{{"ANSIBLE_MODULE_ARGS": {args}}}'.format(args=json.dumps(kwargs))
        js_out = subprocess.Popen(["echo", "{0}".format(js_args)], stdout=subprocess.PIPE)
        md_exc = subprocess.Popen(['python', module.__file__],
                                  stdin=js_out.stdout, stdout=subprocess.PIPE)
        js_out.stdout.close()
        js_out, js_err = md_exc.communicate()

        try:
            out = json.loads(js_out)
        except ValueError as ex:
            out = {'Error': (js_err and (js_err + '.') or str(ex))}
            if js_out:
                out['Given JSON output'] = js_out
            return out

        if 'invocation' in out:
            del out['invocation']

        return out

_resolver = None
_caller = None

def _set_callables(modules):
    '''
    Set all Ansible modules callables
    :return:
    '''
    def _mkf(cmd_name, doc):
        '''
        Create a Salt function for the Ansible module.
        '''
        def _cmd(*args, **kw):
            '''
            Call an Ansible module as a function from the Salt.
            '''
            kwargs = {}
            if kw.get('__pub_arg'):
                for _kw in kw.get('__pub_arg', []):
                    if isinstance(_kw, dict):
                        kwargs = _kw
                        break

            global _caller
            return _caller.call(cmd_name, *args, **kwargs)
        _cmd.__doc__ = doc
        return _cmd

    for mod in modules:
        setattr(sys.modules[__name__], mod, _mkf(mod, 'Available'))

def __virtual__():
    '''
    Ansible module caller.
    :return:
    '''
    ret = ansible is not None
    msg = not ret and "Ansible is not installed on this system" or None
    if msg:
        log.warning(msg)
    else:
        global _resolver
        global _caller
        _resolver = AnsibleModuleResolver(__opts__).resolve().install()
        _caller = AnsibleModuleCaller(_resolver)
    _set_callables(list())

    return ret, msg

def help(module=None, *args):
    '''
    Display help on Ansible standard module.

    :param module:
    :return:
    '''
    if not module:
        raise CommandExecutionError('Please tell me what module you want to have helped with. '
                                    'Or call "ansible.list" to know what is available.')
    try:
        module = _resolver.load_module(module)
    except (ImportError, LoaderError) as err:
        raise CommandExecutionError('Module "{0}" is currently not functional on your system.'.format(module))

    doc = {}
    ret = {}
    for docset in module.DOCUMENTATION.split('---'):
        try:
            docset = yaml.load(docset)
            if docset:
                doc.update(docset)
        except Exception as err:
            log.error("Error parsing doc section: {0}".format(err))
    if not args:
        if 'description' in doc:
            description = doc.get('description') or ''
            del doc['description']
            ret['Description'] = description
        ret['Available sections on module "{}"'.format(module.__name__.replace('ansible.modules.', ''))] = doc.keys()
    else:
        for arg in args:
            info = doc.get(arg)
            if info is not None:
                ret[arg] = info

    return ret

def list(pattern=None):
    '''
    Lists available modules.
    :return:
    '''
    return _resolver.get_modules_list(pattern=pattern)
<|endoftext|>