<commit_before>#!/usr/bin/env python

"""Simple script to take an HDF5 file and chunk it into objects,
   and populate a specified (or managed) Redis instance to hold
   telescope state and other metadata.

   The primary use is for testing the coming katdal access layer
   that will wrap up such an object+Redis dataset and provide
   the usual standard access.

   Objects are stored in chunks of all baselines for obj_chunk_size
   frequencies. The chunking is chosen to produce objects of order 1MB.
   The schema used is as follows:

        <obj_basename>_<dump_index>_<chunk_offset>

   obj_basename: by default the obs start time in integer seconds
   dump_index: integer dump index of this object
   chunk_offset: integer start index for this frequency chunk

   The following useful object parameters are stored in telstate:
        obj_basename: as above
        obj_chunk_size: as above
        obj_size: size per object (needed for reads)
        obj_count: estimated object count (set before writing in this case)
        obj_pool: the name of the CEPH pool used
        obj_ceph_conf: copy of ceph.conf used to connect to target ceph cluster
"""

import struct
import logging
import sys
import time
import shlex
import subprocess
import os

import numpy as np
import h5py
import katsdptelstate
import katsdpservices

try:
    import rados
except ImportError:
    # librados (i.e. direct Ceph access) is only really available on Linux
    rados = None

try:
    import botocore
except ImportError:
    botocore = None
else:
    import botocore.config
    import botocore.session

logging.basicConfig()

logger = logging.getLogger('h5toobj')
logger.setLevel(logging.INFO)

def parse_args():
    parser = katsdpservices.ArgumentParser()
    parser.add_argument('file', type=str, nargs=1,
                        metavar='FILE', help='h5 file to process.')
    parser.add_argument('--basename', type=str, metavar='BASENAME',
                        help='Basename to use for object naming. '
                             'Default is to use file start time.')
    parser.add_argument('--obj-size', type=int, default=20,
                        help='Target obj size as a power of 2. '
                             'Default: 2**20 (1 MB)')
    parser.add_argument('--max-dumps', type=int, default=0,
                        help='Number of dumps to process. Default is all.')
    parser.add_argument('--ceph-conf', type=str, default="/etc/ceph/ceph.conf",
                        metavar='CEPHCONF',
                        help='CEPH configuration file used for cluster connect.')
    parser.add_argument('--ceph-pool', type=str, metavar='POOL',
                        help='CEPH pool to use for object storage.')
    parser.add_argument('--s3-access', type=str,
                        help='AWS access key ID')
    parser.add_argument('--s3-secret', type=str,
                        help='AWS secret access key')
    parser.add_argument('--s3-url', type=str,
                        help='S3 endpoint URL (includes leading "http")')
    parser.add_argument('--redis', type=str,
                        help='Redis host to connect to as Telescope State. '
                             'Default is to start a new local instance.')
    parser.add_argument('--redis-port', type=int, default=6379,
                        help='Port to use when connecting to Redis instance '
                             '(or creating a new one).')
    parser.add_argument('--redis-only', action='store_true',
                        help='Only (re)build Redis DB - no object creation')
    args = parser.parse_args()
    if not rados and not botocore:
        logger.warning("No Ceph or S3 installation found - storing Redis DB only")
        args.redis_only = True
    if not args.redis_only:
        use_s3 = None not in (args.s3_access, args.s3_secret, args.s3_url)
        use_rados = args.ceph_pool is not None
        if use_rados and use_s3:
            parser.error('Please specify either --ceph-pool or --s3-*')
        if use_rados and not rados:
            parser.error('Cannot use Ceph as rados Python library is not installed')
        if use_s3 and not botocore:
            parser.error('Cannot use S3 as botocore Python library is not installed')
    if args.basename is None:
        args.basename = args.file[0].split(".")[0]
    return args

def redis_gen_proto(*args):
    proto = ['*%d\r\n' % (len(args),)]
    proto += ['$%d\r\n%s\r\n' % (len(arg), arg) for arg in args]
    return ''.join(proto)

def redis_bulk_str(r_str, host, port):
    bulk_cmd = "redis-cli --pipe -h {} -p {}".format(host, port)
    bulk_redis = subprocess.Popen(shlex.split(bulk_cmd), stdin=subprocess.PIPE,
                                  stderr=subprocess.PIPE, stdout=subprocess.PIPE)
    (retout, reterr) = bulk_redis.communicate(input=r_str)
    if bulk_redis.returncode:
        logger.error("Failed on bulk key insert. Retcode: %d, Stderr: %s, "
                     "Stdout: %s", bulk_redis.returncode, retout, reterr)
        sys.exit()
    logger.debug("Bulk insert r_str of len %d completed: %s", len(r_str), retout)

def pack_chunk(x):
    # return x.data
    return x.tobytes()

def unpack_chunk(s, dtype, shape):
    # return np.frombuffer(s, dtype).reshape(shape)
    return np.ndarray(shape, dtype, s)

def obj_name(obj_base, dump_index, channel_index):
    return '{}_{}_{}'.format(obj_base, dump_index, channel_index)

def open_rados(ceph_conf, ceph_pool):
    cluster = rados.Rados(conffile=ceph_conf)
    cluster.connect()
    available_pools = cluster.list_pools()
    if ceph_pool not in available_pools:
        logger.error("Specified pool %s not available in this cluster (%s)",
                     ceph_pool, available_pools)
        sys.exit()
    ioctx = cluster.open_ioctx(ceph_pool)
    pool_stats = ioctx.get_stats()
    logger.info("Connected to pool %s. Currently holds %d objects totalling %g GB",
                ceph_pool, pool_stats['num_objects'],
                pool_stats['num_bytes'] / 2 ** 30)
    return ioctx

def write_dump_rados(ioctx, dump_index, dump_data, obj_base, freq_chunksize):
    total_write = 0
    chans = dump_data.shape[0]
    completions = []
    freq_chunks = np.arange(0, chans, freq_chunksize)
    for ch in freq_chunks:
        to_write = pack_chunk(dump_data[ch:ch + freq_chunksize])
        total_write += len(to_write)
        c = ioctx.aio_write_full(obj_name(obj_base, dump_index, ch), to_write,
                                 lambda completion: None, lambda completion: None)
        completions.append(c)
    for c in completions:
        c.wait_for_safe_and_cb()
    return total_write, len(freq_chunks)

def open_s3(access, secret, url):
    settings = dict(service_name='s3', region_name=None, endpoint_url=url,
                    aws_access_key_id=access, aws_secret_access_key=secret,
                    config=botocore.config.Config(max_pool_connections=200))
    session = botocore.session.get_session()
    client = session.create_client(**settings)
    return client

def write_dump_s3(client, dump_index, dump_data, obj_base, freq_chunksize):
    total_write = 0
    chans = dump_data.shape[0]
    freq_chunks = np.arange(0, chans, freq_chunksize)
    for ch in freq_chunks:
        to_write = pack_chunk(dump_data[ch:ch + freq_chunksize])
        total_write += len(to_write)
        bucket, key = os.path.split(obj_name(obj_base, dump_index, ch))
        response = client.put_object(Bucket=bucket, Key=key, Body=to_write)
    return total_write, len(freq_chunks)

def get_freq_chunk(data_shape, target_obj_size=20):
    """Get a frequency chunking that results in slices into this array
       being as close to 1MB as possible. Baselines are always grouped
       into a single object, so outliers in baseline number may produce
       objects much smaller or larger than 1MB"""
    bytes_per_baseline = data_shape[2] * data_shape[3] * 4
    logger.info("Bytes per baseline: %d", bytes_per_baseline)
    channels = data_shape[1]
    for chunk_power in range(int(np.log2(channels)) + 1):
        chunk_size = 2 ** chunk_power
        if (chunk_size * bytes_per_baseline) >= 2 ** target_obj_size:
            break
    real_obj_size = len(np.zeros((chunk_size, data_shape[2], data_shape[3]),
                                 dtype=np.float32).dumps())
     # figure out what the real written size of each dumps'ed object will be
    logger.info("Using chunk size %d giving obj size %d",
                chunk_size, real_obj_size)
    return chunk_size, real_obj_size

def main():
    args = parse_args()
    try:
        h5_file = h5py.File(args.file[0])
    except Exception as e:
        logger.error("Failed to open specified HDF5 file. %s", e)
        sys.exit()
    try:
        data = h5_file['Data/correlator_data']
    except KeyError:
        logger.error("This does not appear to be a valid MeerKAT HDF5 file")
        sys.exit()

    if args.redis is None:
        logger.info("Launching local Redis instance")
        try:
            launch_cmd = "/usr/bin/redis-server --port {}".format(args.redis_port)
            local_redis = subprocess.Popen(shlex.split(launch_cmd),
                                           stderr=subprocess.PIPE,
                                           stdout=subprocess.PIPE)
        except OSError:
            launch_cmd = "/usr/local/bin/redis-server --port {}".format(args.redis_port)
            local_redis = subprocess.Popen(shlex.split(launch_cmd),
                                           stderr=subprocess.PIPE,
                                           stdout=subprocess.PIPE)
        time.sleep(3)
        if local_redis.poll():
            logger.error("Failed to launch local Redis instance, terminating. %s",
                         local_redis.communicate())
            sys.exit()
        logger.info("Local Redis instance launched successfully")
        redis_host = 'localhost'
    else:
        redis_host = args.redis
    ts = katsdptelstate.TelescopeState(endpoint='{}:{}'.format(redis_host,
                                                               args.redis_port))
    logger.info("Connected to Redis on %s:%d. DB has %d existing keys",
                redis_host, args.redis_port, len(ts.keys()))

    r_str = ""
    for attr in h5_file['TelescopeState'].attrs:
        r_str += redis_gen_proto("SET", attr, h5_file['TelescopeState'].attrs[attr])
    redis_bulk_str(r_str, redis_host, args.redis_port)

    for d_count, dset in enumerate(h5_file['TelescopeState'].keys()):
        st = time.time()
        r_str = ""
        d_val = h5_file['TelescopeState'][dset].value
         # much quicker to read it first and then iterate
        for (timestamp, pval) in d_val:
            packed_ts = struct.pack('>d', float(timestamp))
            r_str += redis_gen_proto("ZADD", str(dset), "0", packed_ts + pval)
        bss = time.time()
        redis_bulk_str(r_str, redis_host, args.redis_port)
        logger.info("Added %d items in %gs to key %s. Bulk insert time: %g",
                    len(d_val), time.time() - st, dset, time.time() - bss)
    logger.info("Added %d ranged keys to TelescopeState", d_count + 1)

    if args.redis_only and args.redis is None:
        logger.warning("Terminating locally launched redis instance "
                       "(also saves telstate to local dump.rdb)")
        try:
            cli_cmd = "/usr/bin/redis-cli -p {} SHUTDOWN SAVE".format(args.redis_port)
            subprocess.call(shlex.split(cli_cmd))
        except OSError:
            cli_cmd = "/usr/local/bin/redis-cli -p {} SHUTDOWN SAVE".format(args.redis_port)
            subprocess.call(shlex.split(cli_cmd))
        local_redis.terminate()

    if args.redis_only:
        logger.warning("Building Redis DB only - no data will be written...")
        sys.exit(0)

    # Split visibilities into fixed-size chunks
    freq_chunksize, obj_size = get_freq_chunk(data.shape, args.obj_size)
    max_dumps = args.max_dumps if args.max_dumps > 0 else data.shape[0]
    obj_count_per_dump = data.shape[1] // freq_chunksize
    obj_count = max_dumps * obj_count_per_dump
    ts.add("obj_basename", args.basename)
    ts.add("obj_chunk_size", freq_chunksize)
    ts.add("obj_size", obj_size)
    ts.add("obj_count", obj_count)

    use_rados = args.ceph_pool is not None
    if use_rados:
        store = open_rados(args.ceph_conf, args.ceph_pool)
        ts.add("obj_pool", args.ceph_pool)
        with open(args.ceph_conf, "r") as ceph_conf:
            ts.add("obj_ceph_conf", ceph_conf.readlines())
    else:
        store = open_s3(args.s3_access, args.s3_secret, args.s3_url)
        ts.add("obj_s3_access", args.s3_access)
        ts.add("obj_s3_url", args.s3_url)
    logger.info("Inserted obj schema metadata into telstate")

    logger.info("Processing %d dumps into %d objects of size %d with basename %s",
                max_dumps, obj_count, obj_size, args.basename)
    for dump_index, dump_data in enumerate(data):
        st = time.time()
        if use_rados:
            bytes_written, objs_written = write_dump_rados(store, dump_index, dump_data,
                                                           args.basename, freq_chunksize)
        else:
            bytes_written, objs_written = write_dump_s3(store, dump_index, dump_data,
                                                        args.basename, freq_chunksize)
        et = time.time() - st
        if objs_written == obj_count_per_dump and \
           bytes_written == obj_size * obj_count_per_dump:
            logger.info("Stored dump index %d in %gs (%d objects totalling %gMBps)",
                        dump_index, et, obj_count_per_dump,
                        obj_count_per_dump * obj_size / (1024*1024) / et)
        else:
            logger.error("Failed to full write ts index %d. "
                         "Wrote %d/%d objects and %d/%d bytes",
                         dump_index, objs_written, obj_count_per_dump,
                         bytes_written, obj_size * obj_count_per_dump)
        if dump_index >= max_dumps - 1:
            logger.info("Reached specified ts limit (%d).", max_dumps)
            break
    logger.info("Staging complete...")
    if args.redis is None:
        raw_input("You have started a local Redis server. "
                  "Hit enter to kill this and cleanup.")
        local_redis.terminate()

if __name__ == '__main__':
    main()
<commit_msg>Use dask to chunk the data and write in parallel<commit_after>#!/usr/bin/env python

"""Simple script to take an HDF5 file and chunk it into objects,
   and populate a specified (or managed) Redis instance to hold
   telescope state and other metadata.

   The primary use is for testing the coming katdal access layer
   that will wrap up such an object+Redis dataset and provide
   the usual standard access.

   Objects are stored in chunks of all baselines for obj_chunk_size
   frequencies. The chunking is chosen to produce objects of order 1MB.
   The schema used is as follows:

        <obj_basename>_<dump_index>_<chunk_offset>

   obj_basename: by default the obs start time in integer seconds
   dump_index: integer dump index of this object
   chunk_offset: integer start index for this frequency chunk

   The following useful object parameters are stored in telstate:
        obj_basename: as above
        obj_chunk_size: as above
        obj_size: size per object (needed for reads)
        obj_count: estimated object count (set before writing in this case)
        obj_pool: the name of the CEPH pool used
        obj_ceph_conf: copy of ceph.conf used to connect to target ceph cluster
"""

import struct
import logging
import sys
import time
import shlex
import subprocess
import os

import numpy as np
import h5py
import katsdptelstate
import katsdpservices
import dask
import dask.array as da
from dask.diagnostics import ProgressBar

try:
    import rados
except ImportError:
    # librados (i.e. direct Ceph access) is only really available on Linux
    rados = None

try:
    import botocore
except ImportError:
    botocore = None
else:
    import botocore.config
    import botocore.session

logging.basicConfig()

logger = logging.getLogger('h5toobj')
logger.setLevel(logging.INFO)

def parse_args():
    parser = katsdpservices.ArgumentParser()
    parser.add_argument('file', type=str, nargs=1,
                        metavar='FILE', help='h5 file to process.')
    parser.add_argument('--basename', type=str, metavar='BASENAME',
                        help='Basename to use for object naming. '
                             'Default is to use file start time.')
    parser.add_argument('--obj-size', type=int, default=20,
                        help='Target obj size as a power of 2. '
                             'Default: 2**20 (1 MB)')
    parser.add_argument('--max-dumps', type=int, default=0,
                        help='Number of dumps to process. Default is all.')
    parser.add_argument('--ceph-conf', type=str, default="/etc/ceph/ceph.conf",
                        metavar='CEPHCONF',
                        help='CEPH configuration file used for cluster connect.')
    parser.add_argument('--ceph-pool', type=str, metavar='POOL',
                        help='CEPH pool to use for object storage.')
    parser.add_argument('--s3-access', type=str,
                        help='AWS access key ID')
    parser.add_argument('--s3-secret', type=str,
                        help='AWS secret access key')
    parser.add_argument('--s3-url', type=str,
                        help='S3 endpoint URL (includes leading "http")')
    parser.add_argument('--redis', type=str,
                        help='Redis host to connect to as Telescope State. '
                             'Default is to start a new local instance.')
    parser.add_argument('--redis-port', type=int, default=6379,
                        help='Port to use when connecting to Redis instance '
                             '(or creating a new one).')
    parser.add_argument('--redis-only', action='store_true',
                        help='Only (re)build Redis DB - no object creation')
    args = parser.parse_args()
    if not rados and not botocore:
        logger.warning("No Ceph or S3 installation found - storing Redis DB only")
        args.redis_only = True
    if not args.redis_only:
        use_s3 = None not in (args.s3_access, args.s3_secret, args.s3_url)
        use_rados = args.ceph_pool is not None
        if use_rados and use_s3:
            parser.error('Please specify either --ceph-pool or --s3-*')
        if use_rados and not rados:
            parser.error('Cannot use Ceph as rados Python library is not installed')
        if use_s3 and not botocore:
            parser.error('Cannot use S3 as botocore Python library is not installed')
    if args.basename is None:
        args.basename = args.file[0].split(".")[0]
    return args

def redis_gen_proto(*args):
    proto = ['*%d\r\n' % (len(args),)]
    proto += ['$%d\r\n%s\r\n' % (len(arg), arg) for arg in args]
    return ''.join(proto)

def redis_bulk_str(r_str, host, port):
    bulk_cmd = "redis-cli --pipe -h {} -p {}".format(host, port)
    bulk_redis = subprocess.Popen(shlex.split(bulk_cmd), stdin=subprocess.PIPE,
                                  stderr=subprocess.PIPE, stdout=subprocess.PIPE)
    (retout, reterr) = bulk_redis.communicate(input=r_str)
    if bulk_redis.returncode:
        logger.error("Failed on bulk key insert. Retcode: %d, Stderr: %s, "
                     "Stdout: %s", bulk_redis.returncode, retout, reterr)
        sys.exit()
    logger.debug("Bulk insert r_str of len %d completed: %s", len(r_str), retout)

def pack_chunk(x):
    # return x.data
    return x.tobytes()

def unpack_chunk(s, dtype, shape):
    # return np.frombuffer(s, dtype).reshape(shape)
    return np.ndarray(shape, dtype, s)

def obj_name(obj_base, dump_index, channel_index):
    return '{}_{}_{}'.format(obj_base, dump_index, channel_index)

def key_to_name(key, freq_chunksize):
    return obj_name(key[0], key[1], key[2] * freq_chunksize)

def open_rados(ceph_conf, ceph_pool):
    cluster = rados.Rados(conffile=ceph_conf)
    cluster.connect()
    available_pools = cluster.list_pools()
    if ceph_pool not in available_pools:
        logger.error("Specified pool %s not available in this cluster (%s)",
                     ceph_pool, available_pools)
        sys.exit()
    ioctx = cluster.open_ioctx(ceph_pool)
    pool_stats = ioctx.get_stats()
    logger.info("Connected to pool %s. Currently holds %d objects totalling %g GB",
                ceph_pool, pool_stats['num_objects'],
                pool_stats['num_bytes'] / 2 ** 30)
    return ioctx

def write_chunk_rados(ioctx, name, chunk):
    ioctx.write_full(name, pack_chunk(chunk))

def open_s3(access, secret, url):
    settings = dict(service_name='s3', region_name=None, endpoint_url=url,
                    aws_access_key_id=access, aws_secret_access_key=secret,
                    config=botocore.config.Config(max_pool_connections=200))
    session = botocore.session.get_session()
    client = session.create_client(**settings)
    return client

def write_chunk_s3(client, name, chunk):
    bucket, key = os.path.split(name)
    client.put_object(Bucket=bucket, Key=key, Body=pack_chunk(chunk))

def get_freq_chunk(data_shape, target_obj_size=20):
    """Get a frequency chunking that results in slices into this array
       being as close to 1MB as possible. Baselines are always grouped
       into a single object, so outliers in baseline number may produce
       objects much smaller or larger than 1MB"""
    bytes_per_baseline = data_shape[2] * data_shape[3] * 4
    logger.info("Bytes per baseline: %d", bytes_per_baseline)
    channels = data_shape[1]
    for chunk_power in range(int(np.log2(channels)) + 1):
        chunk_size = 2 ** chunk_power
        if (chunk_size * bytes_per_baseline) >= 2 ** target_obj_size:
            break
    real_obj_size = len(np.zeros((chunk_size, data_shape[2], data_shape[3]),
                                 dtype=np.float32).dumps())
     # figure out what the real written size of each dumps'ed object will be
    logger.info("Using chunk size %d giving obj size %d",
                chunk_size, real_obj_size)
    return chunk_size, real_obj_size

def main():
    args = parse_args()
    try:
        h5_file = h5py.File(args.file[0])
    except Exception as e:
        logger.error("Failed to open specified HDF5 file. %s", e)
        sys.exit()
    try:
        data = h5_file['Data/correlator_data']
    except KeyError:
        logger.error("This does not appear to be a valid MeerKAT HDF5 file")
        sys.exit()

    if args.redis is None:
        logger.info("Launching local Redis instance")
        try:
            launch_cmd = "/usr/bin/redis-server --port {}".format(args.redis_port)
            local_redis = subprocess.Popen(shlex.split(launch_cmd),
                                           stderr=subprocess.PIPE,
                                           stdout=subprocess.PIPE)
        except OSError:
            launch_cmd = "/usr/local/bin/redis-server --port {}".format(args.redis_port)
            local_redis = subprocess.Popen(shlex.split(launch_cmd),
                                           stderr=subprocess.PIPE,
                                           stdout=subprocess.PIPE)
        time.sleep(3)
        if local_redis.poll():
            logger.error("Failed to launch local Redis instance, terminating. %s",
                         local_redis.communicate())
            sys.exit()
        logger.info("Local Redis instance launched successfully")
        redis_host = 'localhost'
    else:
        redis_host = args.redis
    ts = katsdptelstate.TelescopeState(endpoint='{}:{}'.format(redis_host,
                                                               args.redis_port))
    logger.info("Connected to Redis on %s:%d. DB has %d existing keys",
                redis_host, args.redis_port, len(ts.keys()))

    r_str = ""
    for attr in h5_file['TelescopeState'].attrs:
        r_str += redis_gen_proto("SET", attr, h5_file['TelescopeState'].attrs[attr])
    redis_bulk_str(r_str, redis_host, args.redis_port)

    for d_count, dset in enumerate(h5_file['TelescopeState'].keys()):
        st = time.time()
        r_str = ""
        d_val = h5_file['TelescopeState'][dset].value
         # much quicker to read it first and then iterate
        for (timestamp, pval) in d_val:
            packed_ts = struct.pack('>d', float(timestamp))
            r_str += redis_gen_proto("ZADD", str(dset), "0", packed_ts + pval)
        bss = time.time()
        redis_bulk_str(r_str, redis_host, args.redis_port)
        logger.info("Added %d items in %gs to key %s. Bulk insert time: %g",
                    len(d_val), time.time() - st, dset, time.time() - bss)
    logger.info("Added %d ranged keys to TelescopeState", d_count + 1)

    if args.redis_only and args.redis is None:
        logger.warning("Terminating locally launched redis instance "
                       "(also saves telstate to local dump.rdb)")
        try:
            cli_cmd = "/usr/bin/redis-cli -p {} SHUTDOWN SAVE".format(args.redis_port)
            subprocess.call(shlex.split(cli_cmd))
        except OSError:
            cli_cmd = "/usr/local/bin/redis-cli -p {} SHUTDOWN SAVE".format(args.redis_port)
            subprocess.call(shlex.split(cli_cmd))
        local_redis.terminate()

    if args.redis_only:
        logger.warning("Building Redis DB only - no data will be written...")
        sys.exit(0)

    # Split visibilities into fixed-size chunks
    freq_chunksize, obj_size = get_freq_chunk(data.shape, args.obj_size)
    max_dumps = args.max_dumps if args.max_dumps > 0 else data.shape[0]
    obj_count_per_dump = data.shape[1] // freq_chunksize
    obj_count = max_dumps * obj_count_per_dump
    chunk_shape = list(data.shape)
    chunk_shape[0] = 1
    chunk_shape[1] = freq_chunksize
    chunk_shape = tuple(chunk_shape)
    ts.add("obj_basename", args.basename)
    ts.add("obj_chunk_shape", chunk_shape)
    ts.add("obj_size", obj_size)
    ts.add("obj_count", obj_count)

    use_rados = args.ceph_pool is not None
    if use_rados:
        store = open_rados(args.ceph_conf, args.ceph_pool)
        ts.add("obj_pool", args.ceph_pool)
        with open(args.ceph_conf, "r") as ceph_conf:
            ts.add("obj_ceph_conf", ceph_conf.readlines())
    else:
        store = open_s3(args.s3_access, args.s3_secret, args.s3_url)
        ts.add("obj_s3_access", args.s3_access)
        ts.add("obj_s3_url", args.s3_url)
    logger.info("Inserted obj schema metadata into telstate")

    logger.info("Processing %d dumps into %d objects of size %d with basename %s",
                max_dumps, obj_count, obj_size, args.basename)
    dask_data = da.from_array(data[:max_dumps], chunk_shape, args.basename)
    h5_chunks = dask_data.to_delayed().ravel().tolist()
    save = write_chunk_rados if use_rados else write_chunk_s3
    out_chunks = [dask.delayed(save)(store, key_to_name(chunk.key, freq_chunksize), chunk)
                  for chunk in h5_chunks]
    with ProgressBar():
        da.compute(out_chunks, num_workers=8)
    logger.info("Staging complete...")

    if args.redis is None:
        raw_input("You have started a local Redis server. "
                  "Hit enter to kill this and cleanup.")
        local_redis.terminate()

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


<commit_before>import sys
from .testrunner import TestRunner
from .config.toml import TomlConfig

def cli():
	try:
		config = TomlConfig("airbag.toml")
	except ValueError:
		print('Error during configuration file parsing. Aborting...')
		sys.exit(1)

	runner = TestRunner(config)
	if runner.launch() != 0:
		sys.exit(1)<commit_msg>Require at least Python 3.2 to run the program<commit_after>from sys import exit, version_info

if (version_info < (3, 2, 0)):
	print('Python 3.2 is required to run this program')
	exit(1)

from .testrunner import TestRunner
from .config.toml import TomlConfig

def cli():
	try:
		config = TomlConfig("airbag.toml")
	except ValueError:
		print('Error during configuration file parsing. Aborting...')
		exit(1)

	runner = TestRunner(config)
	if runner.launch() != 0:
		exit(1)<|endoftext|>


<commit_before>import glob as g
# import ujson
import os
import os.path as op
import ocrolib
import re

def change_path(path, cat=None, ext=None, remove_type=False, rel_path=None):
    """change catalog paths to a simpler folder structure"""
    norm = os.path.normpath(path)
    path_sep = norm.split(os.sep)
    filename = op.splitext(path_sep[-1])[0]
    # splitext may not remove the extension from names like page1.pseg.bin
    filename = re.sub('\.pseg|\.bin', '', filename)
    # Throw away Filetype dir and filename
    if remove_type:
        path_sep = path_sep[:-2]
    else:
        del path_sep[-1]
    if cat:
        path_sep[-2] = cat
    if ext:
        path_sep.append(filename + ext)
    else:
        path_sep.append(filename)
    path = op.sep.join(path_sep)
    if rel_path:
        path = op.relpath(path, rel_path)
    return path

def convert_path(page, conversion):
    return {'path': change_path(page['path'], **conversion)}

def page_dir(page_path):
    """directory named like the image"""
    pagename = op.basename(change_path(page_path))
    return op.join(op.split(page_path)[0], pagename)

def catalog_pages(cat_path, batch='*', ext='.png', amount=None):
    pattern = op.join(cat_path, '{}/????????{}'.format(batch, ext))
    return g.glob(pattern)

def batches(cat_path):
    pattern = op.join(cat_path, '*')
    return map(op.basename,g.glob(pattern))

def line_index_to_name(idx):
    return '0'+hex((0x010000+(idx)))[2:]

def read_line_boxes(page):
    """read the dimensions of each text line"""
    path = change_path(page['path'], ext='.pseg.png')
    path = path.encode()
    try:
        pseg = ocrolib.read_page_segmentation(path)
    except IOError:
        return []
    regions = ocrolib.RegionExtractor()
    regions.setPageLines(pseg)
    lines = []
    for i in range(1, regions.length()):
        y0, x0, y1, x1 = regions.bboxMath(i)
        lines.append({'name': line_index_to_name(i),
                      'position': [x0, y0, x1, y1]})
    return lines

def load_box_positions(page):
    segments = read_line_boxes(page)
    page['lines'] = segments
    return page

def pages_from_paths(path):
    return {'path': path}

def read_text(page):
    if not 'lines' in page:
        return page
    text_dir = page_dir(page['path'])
    for line in page['lines']:
        path = op.join(text_dir, line['name']) + '.txt'
        text = ""
        if op.isfile(path):
            with open(path, 'rb') as sfile:
                text = sfile.read()
        line['text'] = text

<commit_msg>fix batch ordering<commit_after>import glob as g
# import ujson
import os
import os.path as op
import ocrolib
import re

def change_path(path, cat=None, ext=None, remove_type=False, rel_path=None):
    """change catalog paths to a simpler folder structure"""
    norm = os.path.normpath(path)
    path_sep = norm.split(os.sep)
    filename = op.splitext(path_sep[-1])[0]
    # splitext may not remove the extension from names like page1.pseg.bin
    filename = re.sub('\.pseg|\.bin', '', filename)
    # Throw away Filetype dir and filename
    if remove_type:
        path_sep = path_sep[:-2]
    else:
        del path_sep[-1]
    if cat:
        path_sep[-2] = cat
    if ext:
        path_sep.append(filename + ext)
    else:
        path_sep.append(filename)
    path = op.sep.join(path_sep)
    if rel_path:
        path = op.relpath(path, rel_path)
    return path

def convert_path(page, conversion):
    return {'path': change_path(page['path'], **conversion)}

def page_dir(page_path):
    """directory named like the image"""
    pagename = op.basename(change_path(page_path))
    return op.join(op.split(page_path)[0], pagename)

def catalog_pages(cat_path, batch='*', ext='.png', amount=None):
    pattern = op.join(cat_path, '{}/????????{}'.format(batch, ext))
    return g.glob(pattern)

def batches(cat_path):
    pattern = op.join(cat_path, '*')
    batches = sorted(map(op.basename,g.glob(pattern)))
    return batches

def line_index_to_name(idx):
    return '0'+hex((0x010000+(idx)))[2:]

def read_line_boxes(page):
    """read the dimensions of each text line"""
    path = change_path(page['path'], ext='.pseg.png')
    path = path.encode()
    try:
        pseg = ocrolib.read_page_segmentation(path)
    except IOError:
        return []
    regions = ocrolib.RegionExtractor()
    regions.setPageLines(pseg)
    lines = []
    for i in range(1, regions.length()):
        y0, x0, y1, x1 = regions.bboxMath(i)
        lines.append({'name': line_index_to_name(i),
                      'position': [x0, y0, x1, y1]})
    return lines

def load_box_positions(page):
    segments = read_line_boxes(page)
    page['lines'] = segments
    return page

def pages_from_paths(path):
    return {'path': path}

def read_text(page):
    if not 'lines' in page:
        return page
    text_dir = page_dir(page['path'])
    for line in page['lines']:
        path = op.join(text_dir, line['name']) + '.txt'
        text = ""
        if op.isfile(path):
            with open(path, 'rb') as sfile:
                text = sfile.read()
        line['text'] = text

<|endoftext|>


<commit_before># Copyright 2014 MongoDB, Inc.
#
# 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.

"""Represent the topology of servers."""

from collections import namedtuple

from pymongo import common
from pymongo.server_type import SERVER_TYPE
from pymongo.errors import ConfigurationError
from pymongo.server_description import ServerDescription

TOPOLOGY_TYPE = namedtuple('TopologyType', ['Single', 'ReplicaSetNoPrimary',
                                            'ReplicaSetWithPrimary', 'Sharded',
                                            'Unknown'])(*range(5))

class TopologyDescription(object):
    def __init__(self, topology_type, server_descriptions, replica_set_name):
        """Represent a topology of servers.

        :Parameters:
          - `topology_type`: initial type
          - `server_descriptions`: dict of (address, ServerDescription) for
            all seeds
          - `replica_set_name`: replica set name or None
        """
        self._topology_type = topology_type
        self._replica_set_name = replica_set_name
        self._server_descriptions = server_descriptions

        # Is PyMongo compatible with all servers' wire protocols?
        self._incompatible_err = None

        for s in self._server_descriptions.values():
            # s.min/max_wire_version is the server's wire protocol.
            # MIN/MAX_SUPPORTED_WIRE_VERSION is what PyMongo supports.
            server_too_new = (
                # Server too new.
                s.min_wire_version is not None
                and s.min_wire_version > common.MAX_SUPPORTED_WIRE_VERSION)

            server_too_old = (
                # Server too old.
                s.max_wire_version is not None
                and s.max_wire_version < common.MIN_SUPPORTED_WIRE_VERSION)

            if server_too_new or server_too_old:
                self._incompatible_err = (
                    "Server at %s:%d "
                    "uses wire protocol versions %d through %d, "
                    "but PyMongo only supports %d through %d"
                    % (s.address[0], s.address[1],
                       s.min_wire_version, s.max_wire_version,
                       common.MIN_SUPPORTED_WIRE_VERSION,
                       common.MAX_SUPPORTED_WIRE_VERSION))

                break

    def check_compatible(self):
        """Raise ConfigurationError if any server is incompatible.

        A server is incompatible if its wire protocol version range does not
        overlap with PyMongo's.
        """
        if self._incompatible_err:
            raise ConfigurationError(self._incompatible_err)

    def has_server(self, address):
        return address in self._server_descriptions

    def reset_server(self, address):
        """A copy of this description, with one server marked Unknown."""
        sds = self.server_descriptions()

        # The default ServerDescription's type is Unknown.
        sds[address] = ServerDescription(address)

        if self._topology_type == TOPOLOGY_TYPE.ReplicaSetWithPrimary:
            topology_type = _check_has_primary(sds)
        else:
            topology_type = self._topology_type

        return TopologyDescription(topology_type, sds, self._replica_set_name)

    def reset(self):
        """A copy of this description, with all servers marked Unknown."""
        if self._topology_type == TOPOLOGY_TYPE.ReplicaSetWithPrimary:
            topology_type = TOPOLOGY_TYPE.ReplicaSetNoPrimary
        else:
            topology_type = self._topology_type

        # The default ServerDescription's type is Unknown.
        sds = dict((address, ServerDescription(address))
                   for address in self._server_descriptions)

        return TopologyDescription(topology_type, sds, self._replica_set_name)

    def server_descriptions(self):
        """Dict of (address, ServerDescription)."""
        return self._server_descriptions.copy()

    @property
    def topology_type(self):
        return self._topology_type

    @property
    def replica_set_name(self):
        """The replica set name."""
        return self._replica_set_name

    @property
    def known_servers(self):
        """List of Servers of types besides Unknown."""
        return [s for s in self._server_descriptions.values()
                if s.is_server_type_known]

# If topology type is Unknown and we receive an ismaster response, what should
# the new topology type be?
_SERVER_TYPE_TO_TOPOLOGY_TYPE = {
    SERVER_TYPE.Mongos: TOPOLOGY_TYPE.Sharded,
    SERVER_TYPE.RSPrimary: TOPOLOGY_TYPE.ReplicaSetWithPrimary,
    SERVER_TYPE.RSSecondary: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
    SERVER_TYPE.RSArbiter: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
    SERVER_TYPE.RSOther: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
}

def updated_topology_description(topology_description, server_description):
    """Return an updated copy of a TopologyDescription.

    :Parameters:
      - `topology_description`: the current TopologyDescription
      - `server_description`: a new ServerDescription that resulted from
        an ismaster call

    Called after attempting (successfully or not) to call ismaster on the
    server at server_description.address. Does not modify topology_description.
    """
    address = server_description.address

    # These values will be updated, if necessary, to form the new
    # TopologyDescription.
    topology_type = topology_description.topology_type
    set_name = topology_description.replica_set_name
    server_type = server_description.server_type

    # Don't mutate the original dict of server descriptions; copy it.
    sds = topology_description.server_descriptions()

    # Replace this server's description with the new one.
    sds[address] = server_description

    if topology_type == TOPOLOGY_TYPE.Single:
        # Single type never changes.
        return TopologyDescription(TOPOLOGY_TYPE.Single, sds, set_name)

    if topology_type == TOPOLOGY_TYPE.Unknown:
        if server_type == SERVER_TYPE.Standalone:
            sds.pop(address)

        elif server_type not in (SERVER_TYPE.Unknown, SERVER_TYPE.RSGhost):
            topology_type = _SERVER_TYPE_TO_TOPOLOGY_TYPE[server_type]

    if topology_type == TOPOLOGY_TYPE.Sharded:
        if server_type != SERVER_TYPE.Mongos:
            sds.pop(address)

    elif topology_type == TOPOLOGY_TYPE.ReplicaSetNoPrimary:
        if server_type in (SERVER_TYPE.Standalone, SERVER_TYPE.Mongos):
            sds.pop(address)

        elif server_type == SERVER_TYPE.RSPrimary:
            topology_type, set_name = _update_rs_from_primary(
                sds, set_name, server_description)

        elif server_type in (
                SERVER_TYPE.RSSecondary,
                SERVER_TYPE.RSArbiter,
                SERVER_TYPE.RSOther):
            topology_type, set_name = _update_rs_no_primary_from_member(
                sds, set_name, server_description)

    elif topology_type == TOPOLOGY_TYPE.ReplicaSetWithPrimary:
        if server_type in (SERVER_TYPE.Standalone, SERVER_TYPE.Mongos):
            sds.pop(address)
            topology_type = _check_has_primary(sds)

        elif server_type == SERVER_TYPE.RSPrimary:
            topology_type, set_name = _update_rs_from_primary(
                sds, set_name, server_description)

        elif server_type in (
                SERVER_TYPE.RSSecondary,
                SERVER_TYPE.RSArbiter,
                SERVER_TYPE.RSOther):
            topology_type = _update_rs_with_primary_from_member(
                sds, set_name, server_description)

        else:
            # Server type is Unknown or RSGhost: did we just lose the primary?
            topology_type = _check_has_primary(sds)

    # Return updated copy.
    return TopologyDescription(topology_type, sds, set_name)

def _update_rs_from_primary(sds, replica_set_name, server_description):
    """Update topology description from a primary's ismaster response.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns (new topology type, new replica_set_name).
    """
    if replica_set_name is None:
        replica_set_name = server_description.replica_set_name

    elif replica_set_name != server_description.replica_set_name:
        # We found a primary but it doesn't have the replica_set_name
        # provided by the user.
        sds.pop(server_description.address)
        return _check_has_primary(sds), replica_set_name

    # We've heard from the primary. Is it the same primary as before?
    for server in sds.values():
        if (server.server_type is SERVER_TYPE.RSPrimary
                and server.address != server_description.address):

            # Reset old primary's type to Unknown.
            sds[server.address] = ServerDescription(server.address)

            # There can be only one prior primary.
            break

    # Discover new hosts from this primary's response.
    for new_address in server_description.all_hosts:
        if new_address not in sds:
            sds[new_address] = ServerDescription(new_address)

    # Remove hosts not in the response.
    for addr in set(sds) - server_description.all_hosts:
        sds.pop(addr)

    # If the host list differs from the seed list, we may not have a primary
    # after all.
    return _check_has_primary(sds), replica_set_name

def _update_rs_with_primary_from_member(
        sds,
        replica_set_name,
        server_description):
    """RS with known primary. Process a response from a non-primary.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns new topology type.
    """
    assert replica_set_name is not None

    if replica_set_name != server_description.replica_set_name:
        sds.pop(server_description.address)

    # Had this member been the primary?
    return _check_has_primary(sds)

def _update_rs_no_primary_from_member(
        sds,
        replica_set_name,
        server_description):
    """RS without known primary. Update from a non-primary's response.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns (new topology type, new replica_set_name).
    """
    topology_type = TOPOLOGY_TYPE.ReplicaSetNoPrimary
    if replica_set_name is None:
        replica_set_name = server_description.replica_set_name

    elif replica_set_name != server_description.replica_set_name:
        sds.pop(server_description.address)
        return topology_type, replica_set_name

    # This isn't the primary's response, so don't remove any servers
    # it doesn't report. Only add new servers.
    for address in server_description.all_hosts:
        if address not in sds:
            sds[address] = ServerDescription(address)

    return topology_type, replica_set_name

def _check_has_primary(sds):
    """Current topology type is ReplicaSetWithPrimary. Is primary still known?

    Pass in a dict of ServerDescriptions.

    Returns new topology type.
    """
    for s in sds.values():
        if s.server_type == SERVER_TYPE.RSPrimary:
            return TOPOLOGY_TYPE.ReplicaSetWithPrimary
    else:
        return TOPOLOGY_TYPE.ReplicaSetNoPrimary
<commit_msg>Simplify TopologyDescription.reset_server.<commit_after># Copyright 2014 MongoDB, Inc.
#
# 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.

"""Represent the topology of servers."""

from collections import namedtuple

from pymongo import common
from pymongo.server_type import SERVER_TYPE
from pymongo.errors import ConfigurationError
from pymongo.server_description import ServerDescription

TOPOLOGY_TYPE = namedtuple('TopologyType', ['Single', 'ReplicaSetNoPrimary',
                                            'ReplicaSetWithPrimary', 'Sharded',
                                            'Unknown'])(*range(5))

class TopologyDescription(object):
    def __init__(self, topology_type, server_descriptions, replica_set_name):
        """Represent a topology of servers.

        :Parameters:
          - `topology_type`: initial type
          - `server_descriptions`: dict of (address, ServerDescription) for
            all seeds
          - `replica_set_name`: replica set name or None
        """
        self._topology_type = topology_type
        self._replica_set_name = replica_set_name
        self._server_descriptions = server_descriptions

        # Is PyMongo compatible with all servers' wire protocols?
        self._incompatible_err = None

        for s in self._server_descriptions.values():
            # s.min/max_wire_version is the server's wire protocol.
            # MIN/MAX_SUPPORTED_WIRE_VERSION is what PyMongo supports.
            server_too_new = (
                # Server too new.
                s.min_wire_version is not None
                and s.min_wire_version > common.MAX_SUPPORTED_WIRE_VERSION)

            server_too_old = (
                # Server too old.
                s.max_wire_version is not None
                and s.max_wire_version < common.MIN_SUPPORTED_WIRE_VERSION)

            if server_too_new or server_too_old:
                self._incompatible_err = (
                    "Server at %s:%d "
                    "uses wire protocol versions %d through %d, "
                    "but PyMongo only supports %d through %d"
                    % (s.address[0], s.address[1],
                       s.min_wire_version, s.max_wire_version,
                       common.MIN_SUPPORTED_WIRE_VERSION,
                       common.MAX_SUPPORTED_WIRE_VERSION))

                break

    def check_compatible(self):
        """Raise ConfigurationError if any server is incompatible.

        A server is incompatible if its wire protocol version range does not
        overlap with PyMongo's.
        """
        if self._incompatible_err:
            raise ConfigurationError(self._incompatible_err)

    def has_server(self, address):
        return address in self._server_descriptions

    def reset_server(self, address):
        """A copy of this description, with one server marked Unknown."""
        return updated_topology_description(self, ServerDescription(address))

    def reset(self):
        """A copy of this description, with all servers marked Unknown."""
        if self._topology_type == TOPOLOGY_TYPE.ReplicaSetWithPrimary:
            topology_type = TOPOLOGY_TYPE.ReplicaSetNoPrimary
        else:
            topology_type = self._topology_type

        # The default ServerDescription's type is Unknown.
        sds = dict((address, ServerDescription(address))
                   for address in self._server_descriptions)

        return TopologyDescription(topology_type, sds, self._replica_set_name)

    def server_descriptions(self):
        """Dict of (address, ServerDescription)."""
        return self._server_descriptions.copy()

    @property
    def topology_type(self):
        return self._topology_type

    @property
    def replica_set_name(self):
        """The replica set name."""
        return self._replica_set_name

    @property
    def known_servers(self):
        """List of Servers of types besides Unknown."""
        return [s for s in self._server_descriptions.values()
                if s.is_server_type_known]

# If topology type is Unknown and we receive an ismaster response, what should
# the new topology type be?
_SERVER_TYPE_TO_TOPOLOGY_TYPE = {
    SERVER_TYPE.Mongos: TOPOLOGY_TYPE.Sharded,
    SERVER_TYPE.RSPrimary: TOPOLOGY_TYPE.ReplicaSetWithPrimary,
    SERVER_TYPE.RSSecondary: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
    SERVER_TYPE.RSArbiter: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
    SERVER_TYPE.RSOther: TOPOLOGY_TYPE.ReplicaSetNoPrimary,
}

def updated_topology_description(topology_description, server_description):
    """Return an updated copy of a TopologyDescription.

    :Parameters:
      - `topology_description`: the current TopologyDescription
      - `server_description`: a new ServerDescription that resulted from
        an ismaster call

    Called after attempting (successfully or not) to call ismaster on the
    server at server_description.address. Does not modify topology_description.
    """
    address = server_description.address

    # These values will be updated, if necessary, to form the new
    # TopologyDescription.
    topology_type = topology_description.topology_type
    set_name = topology_description.replica_set_name
    server_type = server_description.server_type

    # Don't mutate the original dict of server descriptions; copy it.
    sds = topology_description.server_descriptions()

    # Replace this server's description with the new one.
    sds[address] = server_description

    if topology_type == TOPOLOGY_TYPE.Single:
        # Single type never changes.
        return TopologyDescription(TOPOLOGY_TYPE.Single, sds, set_name)

    if topology_type == TOPOLOGY_TYPE.Unknown:
        if server_type == SERVER_TYPE.Standalone:
            sds.pop(address)

        elif server_type not in (SERVER_TYPE.Unknown, SERVER_TYPE.RSGhost):
            topology_type = _SERVER_TYPE_TO_TOPOLOGY_TYPE[server_type]

    if topology_type == TOPOLOGY_TYPE.Sharded:
        if server_type != SERVER_TYPE.Mongos:
            sds.pop(address)

    elif topology_type == TOPOLOGY_TYPE.ReplicaSetNoPrimary:
        if server_type in (SERVER_TYPE.Standalone, SERVER_TYPE.Mongos):
            sds.pop(address)

        elif server_type == SERVER_TYPE.RSPrimary:
            topology_type, set_name = _update_rs_from_primary(
                sds, set_name, server_description)

        elif server_type in (
                SERVER_TYPE.RSSecondary,
                SERVER_TYPE.RSArbiter,
                SERVER_TYPE.RSOther):
            topology_type, set_name = _update_rs_no_primary_from_member(
                sds, set_name, server_description)

    elif topology_type == TOPOLOGY_TYPE.ReplicaSetWithPrimary:
        if server_type in (SERVER_TYPE.Standalone, SERVER_TYPE.Mongos):
            sds.pop(address)
            topology_type = _check_has_primary(sds)

        elif server_type == SERVER_TYPE.RSPrimary:
            topology_type, set_name = _update_rs_from_primary(
                sds, set_name, server_description)

        elif server_type in (
                SERVER_TYPE.RSSecondary,
                SERVER_TYPE.RSArbiter,
                SERVER_TYPE.RSOther):
            topology_type = _update_rs_with_primary_from_member(
                sds, set_name, server_description)

        else:
            # Server type is Unknown or RSGhost: did we just lose the primary?
            topology_type = _check_has_primary(sds)

    # Return updated copy.
    return TopologyDescription(topology_type, sds, set_name)

def _update_rs_from_primary(sds, replica_set_name, server_description):
    """Update topology description from a primary's ismaster response.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns (new topology type, new replica_set_name).
    """
    if replica_set_name is None:
        replica_set_name = server_description.replica_set_name

    elif replica_set_name != server_description.replica_set_name:
        # We found a primary but it doesn't have the replica_set_name
        # provided by the user.
        sds.pop(server_description.address)
        return _check_has_primary(sds), replica_set_name

    # We've heard from the primary. Is it the same primary as before?
    for server in sds.values():
        if (server.server_type is SERVER_TYPE.RSPrimary
                and server.address != server_description.address):

            # Reset old primary's type to Unknown.
            sds[server.address] = ServerDescription(server.address)

            # There can be only one prior primary.
            break

    # Discover new hosts from this primary's response.
    for new_address in server_description.all_hosts:
        if new_address not in sds:
            sds[new_address] = ServerDescription(new_address)

    # Remove hosts not in the response.
    for addr in set(sds) - server_description.all_hosts:
        sds.pop(addr)

    # If the host list differs from the seed list, we may not have a primary
    # after all.
    return _check_has_primary(sds), replica_set_name

def _update_rs_with_primary_from_member(
        sds,
        replica_set_name,
        server_description):
    """RS with known primary. Process a response from a non-primary.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns new topology type.
    """
    assert replica_set_name is not None

    if replica_set_name != server_description.replica_set_name:
        sds.pop(server_description.address)

    # Had this member been the primary?
    return _check_has_primary(sds)

def _update_rs_no_primary_from_member(
        sds,
        replica_set_name,
        server_description):
    """RS without known primary. Update from a non-primary's response.

    Pass in a dict of ServerDescriptions, current replica set name, and the
    ServerDescription we are processing.

    Returns (new topology type, new replica_set_name).
    """
    topology_type = TOPOLOGY_TYPE.ReplicaSetNoPrimary
    if replica_set_name is None:
        replica_set_name = server_description.replica_set_name

    elif replica_set_name != server_description.replica_set_name:
        sds.pop(server_description.address)
        return topology_type, replica_set_name

    # This isn't the primary's response, so don't remove any servers
    # it doesn't report. Only add new servers.
    for address in server_description.all_hosts:
        if address not in sds:
            sds[address] = ServerDescription(address)

    return topology_type, replica_set_name

def _check_has_primary(sds):
    """Current topology type is ReplicaSetWithPrimary. Is primary still known?

    Pass in a dict of ServerDescriptions.

    Returns new topology type.
    """
    for s in sds.values():
        if s.server_type == SERVER_TYPE.RSPrimary:
            return TOPOLOGY_TYPE.ReplicaSetWithPrimary
    else:
        return TOPOLOGY_TYPE.ReplicaSetNoPrimary
<|endoftext|>


<commit_before>#!/usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#

import sys
from datetime import datetime, timedelta, tzinfo

try:
    from datetime import timezone

    utc = timezone.utc
except ImportError:
    # Hi there python2 user
    class UTC(tzinfo):
        def utcoffset(self, dt):
            return timedelta(0)

        def tzname(self, dt):
            return "UTC"

        def dst(self, dt):
            return timedelta(0)

    utc = UTC()
from configlib import getConfig, OptionParser
import json
import mozdef_client as mozdef
# TODO: to write/vendor this client
# import uptycs_client
import pickle

def normalize(details):
    # Normalizes fields to conform to http://mozdef.readthedocs.io/en/latest/usage.html#mandatory-fields
    # This is mainly used for common field names to put inside the details structure
    # There might be faster ways to do this
    normalized = {}

    for f in details:
        if f in ("ip", "ip_address", "client_ip"):
            normalized["sourceipaddress"] = details[f]
            continue
        if f == "result":
            if details[f].lower() == "success":
                normalized["success"] = True
            else:
                normalized["success"] = False
        normalized[f] = details[f]

    if "user" in normalized and type(normalized["user"]) is dict:
        if "name" in normalized["user"]:
            normalized["username"] = normalized["user"]["name"]
        if "key" in normalized["user"]:
            normalized["userkey"] = normalized["user"]["key"]
        del (normalized["user"])

    return normalized

def process_alerts(mozmsg, uptycs_alerts):
    """
    Data format of duo_events in api_version == 2 (str):
    duo_events.metadata = {u'total_objects': 49198, u'next_offset': [u'1547244648000', u'4da7180c-b1e5-47b4-9f4d-ee10dc3b5ac8']}
    duo_events.authlogs = [{...}, {...}, ...]
    authlogs entry = {u'access_device': {u'ip': u'a.b.c.d', u'location': {u'city': None, u'state': u'Anhui', u'country':
    u'China'}}, u'event_type': u'authentication', u'timestamp': 1547244800, u'factor': u'not_available', u'reason':
    u'deny_unenrolled_user', u'txid': u'68b33dd3-d341-46c6-a985-0640592fb7b0', u'application': {u'name': u'Integration
    Name Here', u'key': u'SOME KEY HERE'}, u'host': u'api-blah.duosecurity.com', u'result': u'denied', u'eventtype': u'authentication', u'auth_device': {u'ip': None, u'location': {u'city': None, u'state': None, u'country': None}, u'name': None}, u'user': {u'name': u'root', u'key': None}}
    """
    # There are some key fields that we use as MozDef fields, those are set to "noconsume"
    # After processing these fields, we just pour everything into the "details" fields of Mozdef, except for the
    # noconsume fields.

    if etype == "administration":
        noconsume = ["timestamp", "host", "action"]
    elif etype == "telephony":
        noconsume = ["timestamp", "host", "context"]
    elif etype == "authentication":
        noconsume = ["timestamp", "host", "eventtype"]
    else:
        return

    # Care for API v2
    if isinstance(duo_events, dict) and "authlogs" in duo_events:
        offset = duo_events["metadata"]["next_offset"]
        if offset is not None:
            state["{}_offset".format(etype)] = offset
        duo_events = duo_events["authlogs"]
        api_version = 2
    else:
        api_version = 1

    for e in duo_events:
        details = {}
        # Timestamp format: http://mozdef.readthedocs.io/en/latest/usage.html#mandatory-fields
        # Duo logs come as a UTC timestamp
        dt = datetime.utcfromtimestamp(e["timestamp"])
        mozmsg.timestamp = dt.replace(tzinfo=utc).isoformat()
        mozmsg.log["hostname"] = e["host"]
        for i in e:
            if i in noconsume:
                continue

            # Duo client doesn't translate inner dicts to dicts for some reason - its just a string, so we have to process and parse it
            if e[i] is not None and type(e[i]) == str and e[i].startswith("{"):
                j = json.loads(e[i])
                for x in j:
                    details[x] = j[x]
                continue

            details[i] = e[i]
        mozmsg.set_category(etype)
        localdetails = normalize(details)
        if "access_device" in localdetails:
            if "ip" in localdetails["access_device"]:
                localdetails["sourceipaddress"] = localdetails["access_device"]["ip"]
        mozmsg.details = localdetails
        del(localdetails)
        if etype == "administration":
            mozmsg.summary = e["action"]
        elif etype == "telephony":
            mozmsg.summary = e["context"]
        elif etype == "authentication":
            if api_version == 1:
                mozmsg.summary = (
                    e["eventtype"] + " " + e["result"] + " for " + e["username"]
                )
            else:
                mozmsg.summary = (
                    e["eventtype"] + " " + e["result"] + " for " + e["user"]["name"]
                )

        mozmsg.send()

    # last event timestamp record is stored and returned so that we can save our last position in the log.
    try:
        state[etype] = e["timestamp"]
    except UnboundLocalError:
        # duo_events was empty, no new event
        pass
    return state

def main():
    # FYI: This doesn't actually work yet, just mocking out as a mix of real-code and pseudo code
    api_call = UptycsRest(api_json_key=keyfile,
                          verify_ssl=verify_ssl,
                          suffix=options.DOMAIN,
                          api='/alerts',
                          method='GET',
                          post_data=None,
                          post_data_file=None)

    api_response = api_call.call_api()

    alerts = api_response['items']

    mozmsg = mozdef.MozDefEvent(options.MOZDEF_URL)
    mozmsg.tags = ["uptycs"]
    if options.update_tags != "":
        mozmsg.tags.append(options.update_tags)
    mozmsg.set_category("uptycs")
    mozmsg.source = "UptycsAPI"
    if options.DEBUG:
        mozmsg.debug = options.DEBUG
        mozmsg.set_send_to_syslog(True, only_syslog=True)

    # This will process events for all 3 log types and send them to MozDef. the state stores the last position in the
    # log when this script was last called.
    # NOTE: If administration and telephone logs support a "v2" API in the future it will most likely need to have the
    # same code with `next_offset` as authentication uses.
    process_events(
        mozmsg,
        duo.get_administrator_log(mintime=state["administration"] + 1),
        "administration",
        state,
    )

def initConfig():
    options.DOMAIN = getConfig("DOMAIN", "", options.configfile)
    options.USERID = getConfig("USERID", "", options.configfile)
    options.SECRET = getConfig("SECRET", "", options.configfile)
    options.CUSTOMERID = getConfig("CUSTOMERID", "", options.configfile)
    options.ID = getConfig("ID", "", options.configfile)
    options.KEY = getConfig("KEY", "", options.configfile)
    options.MOZDEF_URL = getConfig("MOZDEF_URL", "", options.configfile)
    options.DEBUG = getConfig("DEBUG", True, options.configfile)
    options.statepath = getConfig("statepath", "", options.configfile)
    options.update_tags = getConfig("addtag", "", options.configfile)

if __name__ == "__main__":
    parser = OptionParser()
    defaultconfigfile = sys.argv[0].replace(".py", ".conf")
    parser.add_option(
        "-c",
        dest="configfile",
        default=defaultconfigfile,
        help="configuration file to use",
    )
    (options, args) = parser.parse_args()
    initConfig()
    main()
<commit_msg>Add initial Uptycs POC logic to actual cron<commit_after>#!/usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#

import sys
from datetime import datetime, timedelta, tzinfo

try:
    from datetime import timezone

    utc = timezone.utc
except ImportError:
    # Hi there python2 user
    class UTC(tzinfo):
        def utcoffset(self, dt):
            return timedelta(0)

        def tzname(self, dt):
            return "UTC"

        def dst(self, dt):
            return timedelta(0)

    utc = UTC()
from configlib import getConfig, OptionParser
import json
import mozdef_client as mozdef
import pickle
import datetime
import pdb
import json
import jwt
import datetime
import requests
import urllib3
import click
import os
import sys

class UptycsFilter(object):
    def __init__(self):
        utc_now = datetime.datetime.utcnow()
        self.start_time = utc_now - datetime.timedelta(days=2)
        self.end_time = utc_now
        self.sort_style = 'alertTime:desc'
        self.limit = 100
        self.offset = 0

    def filter_string(self):
        return json.dumps({
            "alertTime": {
                "between": [
                    self.start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
                    self.end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
                ]
            },
            "status": {"in": ["open", "assigned"]}
        }, sort_keys=False, separators=(',', ':'))

class UptycsClient(object):
    def __init__(self, path):
        self.api_json = json.load(open(path))
        self.url = 'https://' + self.api_json['domain'] + '.uptycs.io'
        self.verify_ssl = True
        self.customer_id = self.api_json['customerId']
        self.key = self.api_json['key']
        self.secret = self.api_json['secret']

    def auth_headers(self):
        headers = {}
        utcnow = datetime.datetime.utcnow()
        date = utcnow.strftime("%a, %d %b %Y %H:%M:%S GMT")
        authVar = jwt.encode({'iss': self.key}, self.secret, algorithm='HS256')
        authorization = "Bearer %s" % (authVar.decode('utf-8'))
        headers['date'] = date
        headers['Authorization'] = authorization
        return headers

    def get(self, path, params={}):
        final_url = ("%s/public/api/customers/%s%s" %
                     (self.url, self.customer_id, path))
        try:
            response = requests.get(url=final_url,
                                    headers=self.auth_headers(),
                                    verify=self.verify_ssl)
            if (response and response.status_code in
                    [requests.codes.ok, requests.codes.bad]):
                return response.json()
        except requests.exceptions.RequestException as e:
            print("ERROR: {}".format(e))

    def alerts(self, filter, sort='alertTime:desc&limit'):
        path = '/alerts?filters=' + filter.filter_string() + "&sort=" + sort
        alerts_json = self.get(path)
        return alerts_json['items']

def normalize(details):
    # Normalizes fields to conform to http://mozdef.readthedocs.io/en/latest/usage.html#mandatory-fields
    # This is mainly used for common field names to put inside the details structure
    # There might be faster ways to do this
    normalized = {}

    # for f in details:
    #     if f in ("ip", "ip_address", "client_ip"):
    #         normalized["sourceipaddress"] = details[f]
    #         continue
    #     if f == "result":
    #         if details[f].lower() == "success":
    #             normalized["success"] = True
    #         else:
    #             normalized["success"] = False
    #     normalized[f] = details[f]

    # if "user" in normalized and type(normalized["user"]) is dict:
    #     if "name" in normalized["user"]:
    #         normalized["username"] = normalized["user"]["name"]
    #     if "key" in normalized["user"]:
    #         normalized["userkey"] = normalized["user"]["key"]
    #     del (normalized["user"])

    # return normalized
    return details

def process_alerts(mozmsg, uptycs_alerts):
    for alert in uptycs_alerts:
        details = {}
        # Timestamp format: http://mozdef.readthedocs.io/en/latest/usage.html#mandatory-fields
        # Duo logs come as a UTC timestamp
        dt = datetime.utcfromtimestamp(e["timestamp"])
        mozmsg.timestamp = dt.replace(tzinfo=utc).isoformat()

        # TODO: Fix this so hostName is properly set
        mozmsg.log["hostname"] = alert["host"]
        for item in alert:
            details[item] = alert[item]

        localdetails = normalize(details)
        mozmsg.details = localdetails
        mozmsg.summary = (
            alert["eventtype"] + " " + alert["result"] + " for " + alert["username"]
        )
        mozmsg.send()

def main():
    mozmsg = mozdef.MozDefEvent(options.MOZDEF_URL)
    mozmsg.tags = ["uptycs"]
    mozmsg.set_category("uptycs")
    mozmsg.source = "UptycsAPI"
    if options.DEBUG:
        mozmsg.debug = options.DEBUG
        mozmsg.set_send_to_syslog(True, only_syslog=True)

    client = UptycsClient(options.UPTYCS_API_JSON_FILE)

    # Adjust the filter as needed to set proper search window
    filter = UptycsFilter()
    utc_now = datetime.datetime.utcnow()

    # If an existing state file exists, set our start window to last run
    if os.path.exists(options.statepath):
        state = pickle.load(open(options.statepath, 'rb'))
        filter.start_time = state["last_run"]
        last_alert_ids = state["last_alert_ids"]
    else:
        last_alert_ids = []
        filter.start_time = utc_now - datetime.timedelta(days=30)

    filter.end_time = utc_now

    # Query alerts that match the filter
    alerts = client.alerts(filter)

    # Process all these alerts in MozDef
    process_alerts(mozmsg, alerts)

    state = {
        "last_run": filter.end_time,
        "last_alert_ids": alert_ids,
    }
    pickle.dump(state, open("poc.state", "wb"))

def initConfig():
    options.UPTYCS_API_JSON_FILE = getConfig("UPTYCS_API_JSON_FILE", "", options.configfile)
    options.MOZDEF_URL = getConfig("MOZDEF_URL", "", options.configfile)
    options.DEBUG = getConfig("DEBUG", True, options.configfile)
    options.statepath = getConfig("statepath", "", options.configfile)

if __name__ == "__main__":
    parser = OptionParser()
    defaultconfigfile = sys.argv[0].replace(".py", ".conf")
    parser.add_option(
        "-c",
        dest="configfile",
        default=defaultconfigfile,
        help="configuration file to use",
    )
    (options, args) = parser.parse_args()
    initConfig()
    main()
<|endoftext|>


<commit_before>try:
    from sys import exit, version_info, stdout, stderr
    import argparse
    import pkg_resources
    from os import getcwd, chdir, path
    from .error import write as writerr
    from .testrunner import TestRunner
    from .messagebag import MessageBag
except:
    if (version_info < (3, 4, 0)):
        stderr.write('Python 3.4+ is required to run this program')
        exit(1)

def cli_parse_args():
    parser = argparse.ArgumentParser(
        description='Runs functional tests on your programs'
    )
    parser.add_argument(
        '-d', '--working-dir',
        type=str,
        metavar='DIR',
        default=getcwd(),
        help='Changes the working directory'
    )
    parser.add_argument(
        '-f', '--input-file',
        type=argparse.FileType('r'),
        metavar='FILE',
        default='airbag.toml',
        help='Changes the configuration used'
    )
    parser.add_argument(
        '-F', '--formatter',
        type=str,
        metavar='FORMATTER',
        default='cli',
        help='Formatter used to display the output. Default is `cli`'
    )
    parser.add_argument(
        '-O', '--out-fmt',
        type=str,
        metavar='OUTPUTS',
        default=[],
        nargs='*',
        help='Additional formatters to use with their file destination'
    )
    parser.add_argument(
        '-W', '--max-workers',
        type=int,
        metavar='WORKERS',
        default=1,
        help='Max concurrent workers to run processes'
    )
    parser.add_argument(
        '-V', '--version',
        action='version',
        version='Airbag 0.3',
        help='Displays the current program\'s version and exit'
    )
    return parser.parse_args()

def get_parsers():
    parsers = dict()
    for v in pkg_resources.iter_entry_points(group='airbag.parsers'):
        v = v.load()
        if v.get_extension() in parsers:
            writerr(
                '{0}: duplicate parsers for this file type'.format(
                    v.get_extension()
                )
            )
        else:
            parsers[v.get_extension()] = v
    if not parsers:
        writerr('no config parsers available. Aborting...')
        exit(1)
    return parsers

def get_runners():
    runners = dict()
    for v in pkg_resources.iter_entry_points(group='airbag.runners'):
        v = v.load()
        if v.get_type() in runners:
            writerr(
                '{0}: duplicate runners for this test type'.format(
                    v.get_type()
                )
            )
        else:
            runners[v.get_type()] = v
    if not runners:
        writerr('no test runners available. Aborting...')
        exit(1)
    return runners

def get_message_bag():
    return MessageBag()

def get_formatter(formatter, message_bag, stream):
    for v in pkg_resources.iter_entry_points(group='airbag.formatters'):
        v = v.load()
        if formatter == v.get_type():
            return v(message_bag, stream)
    raise ValueError("No formatter named {} found".format(formatter))

def get_outputs(outputs, message_bag):
    formatters = []
    for output in outputs:
        try:
            [formatter_name, outfile] = output.split(':', maxsplit=1)
            stream = open(outfile, 'w+')
            formatters.append(
                get_formatter(formatter_name, message_bag, stream)
            )
        except ValueError as e:
            writerr('Error in formatter `{}`: {}'.format(output, e))
        except OSError as e:
            writerr('Couldn\'t open file `{}`: {}'.format(outfile, e))
    return formatters

def get_config(input_file, parsers):
    try:
        filename, fileext = path.splitext(input_file.name)
        config = parsers[fileext[1:]](input_file)
    except ValueError:
        writerr('{0}: error during parsing'.format(input_file.name))
        exit(1)
    return config

def change_working_dir(directory=None):
    if directory is not None:
        if path.exists(directory):
            chdir(directory)
        else:
            writerr(
                '{0}: directory does not exist'.format(directory)
            )
            exit(1)

def get_tests(config, runners):
    tests = []
    rtests = config.parse()
    for rtest in rtests:
        runner = None
        if 'type' in rtest and rtest['type'] is not None:
            try:
                runner = runners[rtest['type']]
            except KeyError:
                writerr(
                    '{0}: no runner defined for this test type'.format(
                        rtest['type']
                    )
                )
                exit(1)
        elif rtest['type'] is None and len(runners) == 1:
            runner = list(runners.values())[0]
        else:
            writerr(
                'no type defined for test \'{0}\''.format(
                    rtest['name'] if 'name' in rtest else '[unknown]'
                )
            )
            exit(1)
        del rtest['type']
        tests.append(runner(**rtest))
    return tests

def cli():
    args = cli_parse_args()
    parsers = get_parsers()
    runners = get_runners()
    message_bag = get_message_bag()
    formatter = get_formatter(args.formatter, message_bag, stdout)
    outputs = get_outputs(args.out_fmt, message_bag)
    outputs.append(formatter)
    config = get_config(args.input_file, parsers)
    change_working_dir(args.working_dir)
    tests = get_tests(config, runners)
    testrunner = TestRunner(tests, outputs, workers=args.max_workers)
    if testrunner.launch() != 0:
        exit(1)
    for formatter in outputs:
        formatter.stream.close()
<commit_msg>Don't forget to close the streams<commit_after>try:
    from sys import exit, version_info, stdout, stderr
    import argparse
    import pkg_resources
    from os import getcwd, chdir, path
    from .error import write as writerr
    from .testrunner import TestRunner
    from .messagebag import MessageBag
except:
    if (version_info < (3, 4, 0)):
        stderr.write('Python 3.4+ is required to run this program')
        exit(1)

def cli_parse_args():
    parser = argparse.ArgumentParser(
        description='Runs functional tests on your programs'
    )
    parser.add_argument(
        '-d', '--working-dir',
        type=str,
        metavar='DIR',
        default=getcwd(),
        help='Changes the working directory'
    )
    parser.add_argument(
        '-f', '--input-file',
        type=argparse.FileType('r'),
        metavar='FILE',
        default='airbag.toml',
        help='Changes the configuration used'
    )
    parser.add_argument(
        '-F', '--formatter',
        type=str,
        metavar='FORMATTER',
        default='cli',
        help='Formatter used to display the output. Default is `cli`'
    )
    parser.add_argument(
        '-O', '--out-fmt',
        type=str,
        metavar='OUTPUTS',
        default=[],
        nargs='*',
        help='Additional formatters to use with their file destination'
    )
    parser.add_argument(
        '-W', '--max-workers',
        type=int,
        metavar='WORKERS',
        default=1,
        help='Max concurrent workers to run processes'
    )
    parser.add_argument(
        '-V', '--version',
        action='version',
        version='Airbag 0.3',
        help='Displays the current program\'s version and exit'
    )
    return parser.parse_args()

def get_parsers():
    parsers = dict()
    for v in pkg_resources.iter_entry_points(group='airbag.parsers'):
        v = v.load()
        if v.get_extension() in parsers:
            writerr(
                '{0}: duplicate parsers for this file type'.format(
                    v.get_extension()
                )
            )
        else:
            parsers[v.get_extension()] = v
    if not parsers:
        writerr('no config parsers available. Aborting...')
        exit(1)
    return parsers

def get_runners():
    runners = dict()
    for v in pkg_resources.iter_entry_points(group='airbag.runners'):
        v = v.load()
        if v.get_type() in runners:
            writerr(
                '{0}: duplicate runners for this test type'.format(
                    v.get_type()
                )
            )
        else:
            runners[v.get_type()] = v
    if not runners:
        writerr('no test runners available. Aborting...')
        exit(1)
    return runners

def get_message_bag():
    return MessageBag()

def get_formatter(formatter, message_bag, stream):
    for v in pkg_resources.iter_entry_points(group='airbag.formatters'):
        v = v.load()
        if formatter == v.get_type():
            return v(message_bag, stream)
    raise ValueError("No formatter named {} found".format(formatter))

def get_outputs(outputs, message_bag):
    formatters = []
    for output in outputs:
        try:
            [formatter_name, outfile] = output.split(':', maxsplit=1)
            stream = open(outfile, 'w+')
            formatters.append(
                get_formatter(formatter_name, message_bag, stream)
            )
        except ValueError as e:
            writerr('Error in formatter `{}`: {}'.format(output, e))
        except OSError as e:
            writerr('Couldn\'t open file `{}`: {}'.format(outfile, e))
    return formatters

def get_config(input_file, parsers):
    try:
        filename, fileext = path.splitext(input_file.name)
        config = parsers[fileext[1:]](input_file)
    except ValueError:
        writerr('{0}: error during parsing'.format(input_file.name))
        exit(1)
    return config

def change_working_dir(directory=None):
    if directory is not None:
        if path.exists(directory):
            chdir(directory)
        else:
            writerr(
                '{0}: directory does not exist'.format(directory)
            )
            exit(1)

def get_tests(config, runners):
    tests = []
    rtests = config.parse()
    for rtest in rtests:
        runner = None
        if 'type' in rtest and rtest['type'] is not None:
            try:
                runner = runners[rtest['type']]
            except KeyError:
                writerr(
                    '{0}: no runner defined for this test type'.format(
                        rtest['type']
                    )
                )
                exit(1)
        elif rtest['type'] is None and len(runners) == 1:
            runner = list(runners.values())[0]
        else:
            writerr(
                'no type defined for test \'{0}\''.format(
                    rtest['name'] if 'name' in rtest else '[unknown]'
                )
            )
            exit(1)
        del rtest['type']
        tests.append(runner(**rtest))
    return tests

def cli():
    args = cli_parse_args()
    parsers = get_parsers()
    runners = get_runners()
    message_bag = get_message_bag()
    formatter = get_formatter(args.formatter, message_bag, stdout)
    outputs = get_outputs(args.out_fmt, message_bag)
    outputs.append(formatter)
    config = get_config(args.input_file, parsers)
    change_working_dir(args.working_dir)
    tests = get_tests(config, runners)
    testrunner = TestRunner(tests, outputs, workers=args.max_workers)
    status = testrunner.launch()
    for formatter in outputs:
        formatter.stream.close()
    if status != 0:
        exit(1)
<|endoftext|>


<commit_before><commit_msg>Added analysis<commit_after># -*- coding: UTF-8 -*-
"""
This module provides functionality for data set analysis.

Functions include volcano plots, sorted tables, and plotting sequence levels.
"""

from __future__ import division

# Built-ins
import logging
import os

# IPython
# from IPython.display import display

# Core data analysis libraries
# from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns

from . import fetch_data, utils

LOGGER = logging.getLogger("pyproteome.analysis")

def hierarchical_heatmap(
    data,
    baseline_channels=None,
    metric="euclidean", method="centroid",
    row_cluster=True, col_cluster=True
):
    """
    Plot a hierarhically-clustered heatmap of a data set.

    Parameters
    ----------
    data : :class:`DataSet<pyproteome.data_sets.DataSet>`
    baseline_channels : list of str, optional
        List of channels to average and use as baseline for each row.
    metric : str, optional
        Hierarchical clustering distance metric.
    method : str, optional
        Hierarchical clustering method.
    row_cluster : bool, optional
    col_cluster : bool, optional
    """
    channels = [
        channel
        for channels in data.groups.values()
        for channel in channels
    ]
    channel_names = [
        data.channels[i]
        for i in channels
    ]

    if baseline_channels is None:
        baseline_channels = list(data.groups.values())[-1]

    psms = data.psms[
        data.psms["Fold Change"].apply(lambda x: max([x, 1/x])) > 1.5
    ]

    raw = psms[channels].as_matrix()
    raw_baseline = psms[baseline_channels].as_matrix().mean(axis=1)
    raw = np.log2((raw.T / raw_baseline).T)

    sns.clustermap(
        raw,
        method=method,
        metric=metric,
        row_cluster=row_cluster,
        col_cluster=col_cluster,
        xticklabels=channel_names,
        yticklabels=False,
    )

def find_tfs(data, folder_name=None, csv_name=None):
    """
    Scan over a data set to find proteins annotated as transcription factors.

    Parameters
    ----------
    data : :class:`DataSet<pyproteome.data_sets.DataSet>`
    folder_name : str, optional
    csv_name : str, optional
    """
    if folder_name is None:
        folder_name = data.name

    utils.make_folder(folder_name)

    if csv_name is None:
        csv_name = "Changing TFs.csv"

    if folder_name and csv_name:
        csv_name = os.path.join(folder_name, csv_name)

    def _is_tf(prots):
        go_terms = (
            "DNA binding",
            "double-stranded DNA binding",
            "transcription factor binding",
            "transcription, DNA-templated",
        )
        return any(
            go_term in go
            for prot in prots
            for go in fetch_data.get_uniprot_data(prot.accession).get(
                "go", [],
            )
            for go_term in go_terms
        )

    tfs = data.psms[data.psms["Proteins"].apply(_is_tf)].copy()
    tfs.sort_values(by="Fold Change", ascending=False, inplace=True)

    if csv_name:
        tfs[["Proteins", "Sequence", "Modifications", "Fold Change"]].to_csv(
            csv_name,
            index=False,
        )

    tfs = tfs[["Proteins", "Sequence", "Fold Change", "p-value"]]

    return tfs.style.set_table_styles(  # Hide index and "Validated" columns
        [
            {"selector": "th:first-child", "props": [("display", "none")]},
            {"selector": "*", "props": [("text-align", "left")]},
        ]
    )
<|endoftext|>