<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import re
import os
import socket
import subprocess

import yaml
import pystache
import requests

def parse_args():
    sections = ['all', 'patroni', 'patronictl', 'certificate', 'wal-e', 'crontab']
    argp = argparse.ArgumentParser(description='Configures Spilo',
                                   epilog="Choose from the following sections:\n\t{}".format('\n\t'.join(sections)),
                                   formatter_class=argparse.RawDescriptionHelpFormatter)

    argp.add_argument('sections', metavar='sections', type=str, nargs='+', choices=sections,
                      help='Which section to (re)configure')
    argp.add_argument('-l', '--loglevel', type=str, help='Explicitly set loglevel')
    argp.add_argument('-f', '--force', help='Overwrite files if they exist', default=False, action='store_true')

    args = vars(argp.parse_args())

    if 'all' in args['sections']:
        args['sections'] = sections
        args['sections'].remove('all')
    args['sections'] = set(args['sections'])

    return args

def write_certificates(environment, overwrite):
    """Write SSL certificate to files

    If certificates are specified, they are written, otherwise
    dummy certificates are generated and written"""

    ssl_keys = ['SSL_CERTIFICATE', 'SSL_PRIVATE_KEY']
    if set(ssl_keys) <= set(environment):
        for k in ssl_keys:
            write_file(environment[k], environment[k + '_FILE'], overwrite)
    else:
        if os.path.exists(environment['SSL_PRIVATE_KEY_FILE']) and not overwrite:
            logging.warning('Private key already exists, not overwriting. (Use option --force if necessary)')
            return
        openssl_cmd = [
            '/usr/bin/openssl',
            'req',
            '-nodes',
            '-new',
            '-x509',
            '-subj',
            '/CN=spilo.dummy.org',
            '-keyout',
            environment['SSL_PRIVATE_KEY_FILE'],
            '-out',
            environment['SSL_CERTIFICATE_FILE'],
        ]
        logging.info('Generating ssl certificate')
        p = subprocess.Popen(openssl_cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output, _ = p.communicate()
        logging.debug(output)

    gid = os.stat(environment['PGHOME']).st_gid
    os.chmod(environment['SSL_PRIVATE_KEY_FILE'], 0o640)
    os.chown(environment['SSL_PRIVATE_KEY_FILE'], -1, gid)

def deep_update(a, b):
    """Updates data structures

    Dicts are merged, recursively
    List b is appended to a (except duplicates)
    For anything else, the value of a is returned"""

    if type(a) is dict and type(b) is dict:
        for key in b:
            if key in a:
                a[key] = deep_update(a[key], b[key])
            else:
                a[key] = b[key]
        return a
    if type(a) is list and type(b) is list:
        return a + [i for i in b if i not in a]

    return a if a is not None else b

TEMPLATE = \
    '''
ttl: &ttl 30
loop_wait: &loop_wait 10
scope: &scope '{{SCOPE}}'
restapi:
  listen: 0.0.0.0:{{APIPORT}}
  connect_address: {{instance_data.ip}}:{{APIPORT}}
postgresql:
  name: {{instance_data.id}}
  scope: *scope
  listen: 0.0.0.0:{{PGPORT}}
  connect_address: {{instance_data.ip}}:{{PGPORT}}
  data_dir: {{PGDATA}}
  initdb:
    - encoding: UTF8
    - locale: en_US.UTF-8
  pg_hba:
    - hostssl all all 0.0.0.0/0 md5
    - host    all all 0.0.0.0/0 md5
  superuser:
    username: postgres
    password: {{PGPASSWORD_SUPERUSER}}
  admin:
    username: admin
    password: {{PGPASSWORD_ADMIN}}
  replication:
    username: standby
    password: {{PGPASSWORD_STANDBY}}
    network: 0.0.0.0/0
  callbacks:
    on_start: patroni_aws
    on_stop: patroni_aws
    on_restart: patroni_aws
    on_role_change: patroni_aws
  wal_e:
    command: patroni_wale_restore
    envdir: {{WALE_ENV_DIR}}
    threshold_megabytes: {{WALE_BACKUP_THRESHOLD_MEGABYTES}}
    threshold_backup_size_percentage: {{WALE_BACKUP_THRESHOLD_PERCENTAGE}}
    use_iam: 1
    retries: 2
  parameters:
    archive_mode: 'on'
    wal_level: hot_standby
    archive_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-push "%p" -p 1
    max_wal_senders: 5
    wal_keep_segments: 8
    archive_timeout: 1800s
    max_connections: {{postgresql.parameters.max_connections}}
    max_replication_slots: 5
    hot_standby: 'on'
    tcp_keepalives_idle: 900
    tcp_keepalives_interval: 100
    ssl: 'on'
    ssl_cert_file: {{SSL_CERTIFICATE_FILE}}
    ssl_key_file: {{SSL_PRIVATE_KEY_FILE}}
    shared_buffers: {{postgresql.parameters.shared_buffers}}
    wal_log_hints: 'on'
    log_line_prefix: '%t [%p]: [%l-1] %c %x %d %u %a %h '
    log_checkpoints: 'on'
    log_lock_waits: 'on'
    log_min_duration_statement: 500
    log_autovacuum_min_duration: 0
    log_connections: 'on'
    log_disconnections: 'on'
    log_statement: 'ddl'
    log_temp_files: 0
  recovery_conf:
    restore_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-fetch "%f" "%p" -p 1
'''

def get_instance_meta_data(key):
    try:
        result = requests.get('http://instance-data/latest/meta-data/{}'.format(key), timeout=2)
        if result.status_code != 200:
            raise Exception('Received status code {} ({})'.format(result.status_code, result.text))
        return result.text
    except Exception as e:
        logging.debug('Could not inspect instance metadata for key {}, error: {}'.format(key, e))

def get_placeholders():
    placeholders = dict(os.environ)

    placeholders.setdefault('PGHOME', os.path.expanduser('~'))
    placeholders.setdefault('APIPORT', '8008')
    placeholders.setdefault('BACKUP_SCHEDULE', '00 01 * * *')
    placeholders.setdefault('CRONTAB', '[]')
    placeholders.setdefault('PGDATA', os.path.join(placeholders['PGHOME'], 'pgdata'))
    placeholders.setdefault('PGPASSWORD_ADMIN', 'standby')
    placeholders.setdefault('PGPASSWORD_STANDBY', 'standby')
    placeholders.setdefault('PGPASSWORD_SUPERUSER', 'zalando')
    placeholders.setdefault('PGPORT', '5432')
    placeholders.setdefault('SCOPE', 'dummy')
    placeholders.setdefault('SSL_CERTIFICATE_FILE', os.path.join(placeholders['PGHOME'], 'server.crt'))
    placeholders.setdefault('SSL_PRIVATE_KEY_FILE', os.path.join(placeholders['PGHOME'], 'server.key'))
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_MEGABYTES', 1024)
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_PERCENTAGE', 30)
    placeholders.setdefault('WALE_ENV_DIR', os.path.join(placeholders['PGHOME'], 'etc', 'wal-e.d', 'env'))
    placeholders.setdefault('WAL_S3_BUCKET', 'spilo-example-com')

    placeholders.setdefault('postgresql', {})
    placeholders['postgresql'].setdefault('parameters', {})

    os_memory_mb = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / 1024 / 1024

    # # We take 1/4 of the memory, expressed in full MB's
    placeholders['postgresql']['parameters']['shared_buffers'] = '{}MB'.format(int(os_memory_mb/4))
    # # 1 connection per 30 MB, at least 100, at most 1000
    placeholders['postgresql']['parameters']['max_connections'] = min(max(100, int(os_memory_mb/30)), 1000)

    placeholders['instance_data'] = dict()
    placeholders['instance_data']['ip'] = get_instance_meta_data('local-ipv4') \
        or socket.gethostbyname(socket.gethostname())
    placeholders['instance_data']['id'] = re.sub(r'\W+', '_', get_instance_meta_data('instance-id')
                                                 or socket.gethostname())

    return placeholders

def write_file(config, filename, overwrite):
    if not overwrite and os.path.exists(filename):
        logging.warning('File {} already exists, not overwriting. (Use option --force if necessary)'.format(filename))
    with open(filename, 'w') as f:
        f.write(config)

def pystache_render(*args, **kwargs):
    render = pystache.Renderer(missing_tags='strict')
    return render.render(*args, **kwargs)

def get_dcs_config(config, placeholders):
    defaults = \
        yaml.load('''\
zookeeper:
  scope: {scope}
  session_timeout: {ttl}
  reconnect_timeout: {loop_wait}
etcd:
  scope: {scope}
  ttl: {ttl}'''.format(**config))

    config = {}

    if 'ZOOKEEPER_HOSTS' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['hosts'] = yaml.load(placeholders['ZOOKEEPER_HOSTS'])
    elif 'EXHIBITOR_HOSTS' in placeholders and 'EXHIBITOR_PORT' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['exhibitor'] = {'poll_interval': '300', 'port': placeholders['EXHIBITOR_PORT'],
                                            'hosts': yaml.load(placeholders['EXHIBITOR_HOSTS'])}
    elif 'ETCD_HOST' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['host'] = placeholders['ETCD_HOST']
    elif 'ETCD_DISCOVERY_DOMAIN' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['discovery_srv'] = placeholders['ETCD_DISCOVERY_DOMAIN']
    else:
        pass  # Configuration can also be specified using PATRONI_CONFIGURATION

    return config

def write_wale_command_environment(placeholders, overwrite):
    az = get_instance_meta_data('placement/availability-zone') or 'dummy-region'

    if not os.path.exists(placeholders['WALE_ENV_DIR']):
        os.makedirs(placeholders['WALE_ENV_DIR'])

    write_file('s3://{WAL_S3_BUCKET}/spilo/{SCOPE}/wal/'.format(**placeholders),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_PREFIX'), overwrite)
    write_file('https+path://s3-{}.amazonaws.com:443'.format(az[:-1]),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_ENDPOINT'), overwrite)

def write_crontab(placeholders, path, overwrite):

    if not overwrite:
        with open(os.devnull, 'w') as devnull:
            cron_exit = subprocess.call(['sudo', '-u', 'postgres', '-l'], stdout=devnull, stderr=devnull)
            if cron_exit == 0:
                logging.warning('Cron is already configured. (Use option --force to overwrite cron)')
                return

    lines = ['PATH={}'.format(path)]
    lines += ['{BACKUP_SCHEDULE} /postgres_backup.sh "{WALE_ENV_DIR}" "{PGDATA}"'.format(**placeholders)]
    lines += yaml.load(placeholders['CRONTAB'])
    lines += ['']  # EOF requires empty line for cron

    c = subprocess.Popen(['sudo', '-u', 'postgres', 'crontab'], stdin=subprocess.PIPE)
    c.communicate(input='\n'.join(lines).encode())

def main():
    debug = os.environ.get('DEBUG', '') in ['1', 'true', 'on', 'ON']
    args = parse_args()

    logging.basicConfig(format='%(asctime)s - bootstrapping - %(levelname)s - %(message)s', level=('DEBUG'
                         if debug else (args.get('loglevel') or 'INFO').upper()))

    placeholders = get_placeholders()

    config = yaml.load(pystache_render(TEMPLATE, placeholders))
    config.update(get_dcs_config(config, placeholders))

    user_config = yaml.load(os.environ.get('PATRONI_CONFIGURATION', '')) or {}
    if not isinstance(user_config, dict):
        raise ValueError('PATRONI_CONFIGURATION should contain a dict, yet it is a {}'.format(type(user_config)))

    config = deep_update(user_config, config)

    # Ensure replication is available
    if not any(['replication' in i for i in config['postgresql']['pg_hba']]):
        rep_hba = 'hostssl replication {} 0.0.0.0/0 md5'.format(config['postgresql']['replication']['username'])
        config['postgresql']['pg_hba'].insert(0, rep_hba)

    for section in args['sections']:
        logging.info('Configuring {}'.format(section))
        if section == 'patroni':
            patroni_configfile = os.path.join(placeholders['PGHOME'], 'postgres.yml')
            write_file(yaml.dump(config, default_flow_style=False, width=120), patroni_configfile, args['force'])
        elif section == 'patronictl':
            patronictl_config = {k: v for k, v in config.items() if k in ['zookeeper', 'etcd', 'consul']}
            patronictl_configfile = os.path.join(placeholders['PGHOME'], '.config', 'patroni', 'patronictl.yaml')
            if not os.path.exists(os.path.dirname(patronictl_configfile)):
                os.makedirs(os.path.dirname(patronictl_configfile))
            write_file(yaml.dump(patronictl_config), patronictl_configfile, args['force'])
        elif section == 'wal-e':
            write_wale_command_environment(placeholders, args['force'])
        elif section == 'certificate':
            write_certificates(placeholders, args['force'])
        elif section == 'crontab':
            write_crontab(placeholders, os.environ.get('PATH'), args['force'])
        else:
            raise Exception('Unknown section: {}'.format(section))

if __name__ == '__main__':
    main()
<commit_msg>Address code feedback<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import re
import os
import socket
import subprocess

import yaml
import pystache
import requests

def parse_args():
    sections = ['all', 'patroni', 'patronictl', 'certificate', 'wal-e', 'crontab']
    argp = argparse.ArgumentParser(description='Configures Spilo',
                                   epilog="Choose from the following sections:\n\t{}".format('\n\t'.join(sections)),
                                   formatter_class=argparse.RawDescriptionHelpFormatter)

    argp.add_argument('sections', metavar='sections', type=str, nargs='+', choices=sections,
                      help='Which section to (re)configure')
    argp.add_argument('-l', '--loglevel', type=str, help='Explicitly set loglevel')
    argp.add_argument('-f', '--force', help='Overwrite files if they exist', default=False, action='store_true')

    args = vars(argp.parse_args())

    if 'all' in args['sections']:
        args['sections'] = sections
        args['sections'].remove('all')
    args['sections'] = set(args['sections'])

    return args

def write_certificates(environment, overwrite):
    """Write SSL certificate to files

    If certificates are specified, they are written, otherwise
    dummy certificates are generated and written"""

    ssl_keys = ['SSL_CERTIFICATE', 'SSL_PRIVATE_KEY']
    if set(ssl_keys) <= set(environment):
        for k in ssl_keys:
            write_file(environment[k], environment[k + '_FILE'], overwrite)
    else:
        if os.path.exists(environment['SSL_PRIVATE_KEY_FILE']) and not overwrite:
            logging.warning('Private key already exists, not overwriting. (Use option --force if necessary)')
            return
        openssl_cmd = [
            '/usr/bin/openssl',
            'req',
            '-nodes',
            '-new',
            '-x509',
            '-subj',
            '/CN=spilo.dummy.org',
            '-keyout',
            environment['SSL_PRIVATE_KEY_FILE'],
            '-out',
            environment['SSL_CERTIFICATE_FILE'],
        ]
        logging.info('Generating ssl certificate')
        p = subprocess.Popen(openssl_cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output, _ = p.communicate()
        logging.debug(output)

    uid = os.stat(environment['PGHOME']).st_uid
    os.chmod(environment['SSL_PRIVATE_KEY_FILE'], 0o600)
    os.chown(environment['SSL_PRIVATE_KEY_FILE'], uid, -1)

def deep_update(a, b):
    """Updates data structures

    Dicts are merged, recursively
    List b is appended to a (except duplicates)
    For anything else, the value of a is returned"""

    if type(a) is dict and type(b) is dict:
        for key in b:
            if key in a:
                a[key] = deep_update(a[key], b[key])
            else:
                a[key] = b[key]
        return a
    if type(a) is list and type(b) is list:
        return a + [i for i in b if i not in a]

    return a if a is not None else b

TEMPLATE = \
    '''
ttl: &ttl 30
loop_wait: &loop_wait 10
scope: &scope '{{SCOPE}}'
restapi:
  listen: 0.0.0.0:{{APIPORT}}
  connect_address: {{instance_data.ip}}:{{APIPORT}}
postgresql:
  name: {{instance_data.id}}
  scope: *scope
  listen: 0.0.0.0:{{PGPORT}}
  connect_address: {{instance_data.ip}}:{{PGPORT}}
  data_dir: {{PGDATA}}
  initdb:
    - encoding: UTF8
    - locale: en_US.UTF-8
  pg_hba:
    - hostssl all all 0.0.0.0/0 md5
    - host    all all 0.0.0.0/0 md5
  superuser:
    username: postgres
    password: {{PGPASSWORD_SUPERUSER}}
  admin:
    username: admin
    password: {{PGPASSWORD_ADMIN}}
  replication:
    username: standby
    password: {{PGPASSWORD_STANDBY}}
    network: 0.0.0.0/0
  callbacks:
    on_start: patroni_aws
    on_stop: patroni_aws
    on_restart: patroni_aws
    on_role_change: patroni_aws
  wal_e:
    command: patroni_wale_restore
    envdir: {{WALE_ENV_DIR}}
    threshold_megabytes: {{WALE_BACKUP_THRESHOLD_MEGABYTES}}
    threshold_backup_size_percentage: {{WALE_BACKUP_THRESHOLD_PERCENTAGE}}
    use_iam: 1
    retries: 2
  parameters:
    archive_mode: 'on'
    wal_level: hot_standby
    archive_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-push "%p" -p 1
    max_wal_senders: 5
    wal_keep_segments: 8
    archive_timeout: 1800s
    max_connections: {{postgresql.parameters.max_connections}}
    max_replication_slots: 5
    hot_standby: 'on'
    tcp_keepalives_idle: 900
    tcp_keepalives_interval: 100
    ssl: 'on'
    ssl_cert_file: {{SSL_CERTIFICATE_FILE}}
    ssl_key_file: {{SSL_PRIVATE_KEY_FILE}}
    shared_buffers: {{postgresql.parameters.shared_buffers}}
    wal_log_hints: 'on'
    log_line_prefix: '%t [%p]: [%l-1] %c %x %d %u %a %h '
    log_checkpoints: 'on'
    log_lock_waits: 'on'
    log_min_duration_statement: 500
    log_autovacuum_min_duration: 0
    log_connections: 'on'
    log_disconnections: 'on'
    log_statement: 'ddl'
    log_temp_files: 0
  recovery_conf:
    restore_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-fetch "%f" "%p" -p 1
'''

def get_instance_meta_data(key):
    try:
        result = requests.get('http://instance-data/latest/meta-data/{}'.format(key), timeout=2)
        if result.status_code != 200:
            raise Exception('Received status code {} ({})'.format(result.status_code, result.text))
        return result.text
    except Exception as e:
        logging.debug('Could not inspect instance metadata for key {}, error: {}'.format(key, e))

def get_placeholders():
    placeholders = dict(os.environ)

    placeholders.setdefault('PGHOME', os.path.expanduser('~'))
    placeholders.setdefault('APIPORT', '8008')
    placeholders.setdefault('BACKUP_SCHEDULE', '00 01 * * *')
    placeholders.setdefault('CRONTAB', '[]')
    placeholders.setdefault('PGDATA', os.path.join(placeholders['PGHOME'], 'pgdata'))
    placeholders.setdefault('PGPASSWORD_ADMIN', 'standby')
    placeholders.setdefault('PGPASSWORD_STANDBY', 'standby')
    placeholders.setdefault('PGPASSWORD_SUPERUSER', 'zalando')
    placeholders.setdefault('PGPORT', '5432')
    placeholders.setdefault('SCOPE', 'dummy')
    placeholders.setdefault('SSL_CERTIFICATE_FILE', os.path.join(placeholders['PGHOME'], 'server.crt'))
    placeholders.setdefault('SSL_PRIVATE_KEY_FILE', os.path.join(placeholders['PGHOME'], 'server.key'))
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_MEGABYTES', 1024)
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_PERCENTAGE', 30)
    placeholders.setdefault('WALE_ENV_DIR', os.path.join(placeholders['PGHOME'], 'etc', 'wal-e.d', 'env'))
    placeholders.setdefault('WAL_S3_BUCKET', 'spilo-example-com')

    placeholders.setdefault('postgresql', {})
    placeholders['postgresql'].setdefault('parameters', {})

    os_memory_mb = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / 1024 / 1024

    # # We take 1/4 of the memory, expressed in full MB's
    placeholders['postgresql']['parameters']['shared_buffers'] = '{}MB'.format(int(os_memory_mb/4))
    # # 1 connection per 30 MB, at least 100, at most 1000
    placeholders['postgresql']['parameters']['max_connections'] = min(max(100, int(os_memory_mb/30)), 1000)

    placeholders['instance_data'] = dict()
    placeholders['instance_data']['ip'] = get_instance_meta_data('local-ipv4') \
        or socket.gethostbyname(socket.gethostname())
    placeholders['instance_data']['id'] = re.sub(r'\W+', '_', get_instance_meta_data('instance-id')
                                                 or socket.gethostname())

    return placeholders

def write_file(config, filename, overwrite):
    if not overwrite and os.path.exists(filename):
        logging.warning('File {} already exists, not overwriting. (Use option --force if necessary)'.format(filename))
    with open(filename, 'w') as f:
        f.write(config)

def pystache_render(*args, **kwargs):
    render = pystache.Renderer(missing_tags='strict')
    return render.render(*args, **kwargs)

def get_dcs_config(config, placeholders):
    defaults = \
        yaml.load('''\
zookeeper:
  scope: {scope}
  session_timeout: {ttl}
  reconnect_timeout: {loop_wait}
etcd:
  scope: {scope}
  ttl: {ttl}'''.format(**config))

    config = {}

    if 'ZOOKEEPER_HOSTS' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['hosts'] = yaml.load(placeholders['ZOOKEEPER_HOSTS'])
    elif 'EXHIBITOR_HOSTS' in placeholders and 'EXHIBITOR_PORT' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['exhibitor'] = {'poll_interval': '300', 'port': placeholders['EXHIBITOR_PORT'],
                                            'hosts': yaml.load(placeholders['EXHIBITOR_HOSTS'])}
    elif 'ETCD_HOST' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['host'] = placeholders['ETCD_HOST']
    elif 'ETCD_DISCOVERY_DOMAIN' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['discovery_srv'] = placeholders['ETCD_DISCOVERY_DOMAIN']
    else:
        pass  # Configuration can also be specified using PATRONI_CONFIGURATION

    return config

def write_wale_command_environment(placeholders, overwrite):
    az = get_instance_meta_data('placement/availability-zone') or 'dummy-region'

    if not os.path.exists(placeholders['WALE_ENV_DIR']):
        os.makedirs(placeholders['WALE_ENV_DIR'])

    write_file('s3://{WAL_S3_BUCKET}/spilo/{SCOPE}/wal/'.format(**placeholders),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_PREFIX'), overwrite)
    write_file('https+path://s3-{}.amazonaws.com:443'.format(az[:-1]),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_ENDPOINT'), overwrite)

def write_crontab(placeholders, path, overwrite):

    if not overwrite:
        with open(os.devnull, 'w') as devnull:
            cron_exit = subprocess.call(['sudo', '-u', 'postgres', 'crontab', '-l'], stdout=devnull, stderr=devnull)
            if cron_exit == 0:
                logging.warning('Cron is already configured. (Use option --force to overwrite cron)')
                return

    lines = ['PATH={}'.format(path)]
    lines += ['{BACKUP_SCHEDULE} /postgres_backup.sh "{WALE_ENV_DIR}" "{PGDATA}"'.format(**placeholders)]
    lines += yaml.load(placeholders['CRONTAB'])
    lines += ['']  # EOF requires empty line for cron

    c = subprocess.Popen(['sudo', '-u', 'postgres', 'crontab'], stdin=subprocess.PIPE)
    c.communicate(input='\n'.join(lines).encode())

def main():
    debug = os.environ.get('DEBUG', '') in ['1', 'true', 'on', 'ON']
    args = parse_args()

    logging.basicConfig(format='%(asctime)s - bootstrapping - %(levelname)s - %(message)s', level=('DEBUG'
                         if debug else (args.get('loglevel') or 'INFO').upper()))

    placeholders = get_placeholders()

    config = yaml.load(pystache_render(TEMPLATE, placeholders))
    config.update(get_dcs_config(config, placeholders))

    user_config = yaml.load(os.environ.get('PATRONI_CONFIGURATION', '')) or {}
    if not isinstance(user_config, dict):
        raise ValueError('PATRONI_CONFIGURATION should contain a dict, yet it is a {}'.format(type(user_config)))

    config = deep_update(user_config, config)

    # Ensure replication is available
    if not any(['replication' in i for i in config['postgresql']['pg_hba']]):
        rep_hba = 'hostssl replication {} 0.0.0.0/0 md5'.format(config['postgresql']['replication']['username'])
        config['postgresql']['pg_hba'].insert(0, rep_hba)

    for section in args['sections']:
        logging.info('Configuring {}'.format(section))
        if section == 'patroni':
            patroni_configfile = os.path.join(placeholders['PGHOME'], 'postgres.yml')
            write_file(yaml.dump(config, default_flow_style=False, width=120), patroni_configfile, args['force'])
        elif section == 'patronictl':
            patronictl_config = {k: v for k, v in config.items() if k in ['zookeeper', 'etcd', 'consul']}
            patronictl_configfile = os.path.join(placeholders['PGHOME'], '.config', 'patroni', 'patronictl.yaml')
            if not os.path.exists(os.path.dirname(patronictl_configfile)):
                os.makedirs(os.path.dirname(patronictl_configfile))
            write_file(yaml.dump(patronictl_config), patronictl_configfile, args['force'])
        elif section == 'wal-e':
            write_wale_command_environment(placeholders, args['force'])
        elif section == 'certificate':
            write_certificates(placeholders, args['force'])
        elif section == 'crontab':
            write_crontab(placeholders, os.environ.get('PATH'), args['force'])
        else:
            raise Exception('Unknown section: {}'.format(section))

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


<commit_before>import abc
from nalaf.features import FeatureGenerator
import re
from nalaf import print_debug, print_verbose

class EdgeFeatureGenerator(FeatureGenerator):
    """
    Abstract class for generating features for each edge in the dataset.
    Subclasses that inherit this class should:
    * Be named [Name]FeatureGenerator
    * Implement the abstract method generate
    * Append new items to the dictionary field "features" of each Edge in the dataset
    """

    @abc.abstractmethod
    def generate(self, dataset, feature_set, use_gold=True, use_pred=False):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        pass

    def add_to_feature_set(self, feature_set, edge, feature_name, value=1):
        """
        Return True if feature was added to feature_set. False, otherwise

        If the feature_name is None, the feature is not added in anycase. See: self.mk_feature_name
        """
        if feature_name is None:
            return False

        else:
            feature_name = self.__set_final_name(feature_name)

            if not feature_set.is_locked:
                feature_index = feature_set.get(feature_name, None)

                if feature_index is None:
                    feature_index = len(feature_set)
                    feature_set[feature_name] = feature_index
                    print_verbose("Feature map: {} == {} -- _1st_ value: {}".format(str(feature_key), feature_name, str(value)))

                edge.features[feature_index] = value
                return True

            else:
                feature_index = feature_set.get(feature_name, None)

                if feature_index is not None:
                    edge.features[feature_index] = value
                    return True
                else:
                    return False

    def __set_final_name(self, feature_name):
        if not re.search('\[-?[0-9]+\]$', feature_name):
            # Identify the window position --> TODO likely deletable from a edge feature generator
            feature_name = feature_name + "_[0]"

        if not feature_name.startswith(self.__class__.__name__):
            feature_name = self.__class__.__name__ + "::" + feature_name

        return feature_name

    def mk_feature_name(self, prefix, *args):
        if prefix is None:
            return None
        else:
            l = [str(x) for x in ([prefix] + list(args))]
            return "_".join(l)

    def gen_prefix_feat_name(self, field_prefix_feature, *args):
        prefix = self.__getattribute__(field_prefix_feature)
        pure_name = field_prefix_feature[field_prefix_feature.find("_") + 1:]  # Remove "prefix_"
        feature_name = self.mk_feature_name(prefix, pure_name, *args)
        # print_debug(feature_name, field_prefix_feature, args)
        return feature_name

    def add(self, feature_set, edge, field_prefix_feature, *args):
        feature_name = self.gen_prefix_feat_name(field_prefix_feature, *args)
        self.add_to_feature_set(feature_set, edge, feature_name)

    def add_with_value(self, feature_set, edge, field_prefix_feature, value, *args):
        feature_name = self.gen_prefix_feat_name(field_prefix_feature, *args)
        self.add_to_feature_set(feature_set, edge, feature_name, value=value)

from nalaf.features.relations import EdgeFeatureGenerator
from nltk.stem import PorterStemmer
from math import log2
from operator import itemgetter

class TokenFeatureGenerator(EdgeFeatureGenerator):
    """
    Token based features for each entity belonging to an edge
    """

    def __init__(
        self,
        prefix_txt=None,  # 73 in relna
        prefix_pos=None,  # 74
        prefix_masked_txt=None,  # 75
        prefix_stem_masked_txt=None,  # 76
        prefix_ann_type=None,  # 77
    ):
        self.stemmer = PorterStemmer()
        """an instance of the PorterStemmer()"""

        self.prefix_txt = prefix_txt
        self.prefix_pos = prefix_pos
        self.prefix_masked_txt = prefix_masked_txt
        self.prefix_stem_masked_txt = prefix_stem_masked_txt
        self.prefix_ann_type = prefix_ann_type

    @abc.abstractmethod
    def generate(self, dataset, feature_set, is_training_mode):
        """
        Does nothing directly
        """
        pass

    def token_features(self, token, addendum, edge, feature_set, is_training_mode):
        feature_name_1 = self.gen_prefix_feat_name("prefix_txt", addendum, token.word)
        self.add_to_feature_set(feature_set, edge, feature_name_1)

        feature_name_2 = self.gen_prefix_feat_name("prefix_pos", addendum, token.features['pos'])
        self.add_to_feature_set(feature_set, edge, feature_name_2)

        feature_name_3 = self.gen_prefix_feat_name("prefix_masked_txt", addendum, token.masked_text(edge.same_part))
        self.add_to_feature_set(feature_set, edge, feature_name_3)

        # TODO why stem of masked text? -- makes little sense -- See TODO in original loctext too
        feature_name_4 = self.gen_prefix_feat_name("prefix_stem_masked_txt", addendum, self.stemmer.stem(token.masked_text(edge.same_part)))
        self.add_to_feature_set(feature_set, edge, feature_name_4)

        ann_types = self.annotated_types(token, edge)
        for ann in ann_types:
            feature_name_5 = self.gen_prefix_feat_name("prefix_ann_type", addendum, ann)
            self.add_to_feature_set(feature_set, edge, feature_name_5)

    def annotated_types(self, token, edge):
        head1 = edge.entity1.head_token
        head2 = edge.entity2.head_token

        if not token.is_entity_part(edge.same_part):
            feature_name = 'no_ann_type'
            return [feature_name]
        else:
            ann_types = []
            if token.is_entity_part(edge.same_part):
                entity = token.get_entity(edge.same_part)
                feature_name_1 = entity.class_id
                ann_types.append(feature_name_1)
                if entity == edge.entity1:
                    feature_name_2 = 'entity1_' + edge.entity1.class_id
                    ann_types.append(feature_name_2)
                    return ann_types
                elif entity == edge.entity2:
                    feature_name_2 = 'entity2_' + edge.entity2.class_id
                    ann_types.append(feature_name_2)
                    return ann_types
                return ann_types

def calculateInformationGain(feature_set, dataset, output_file):
    number_pos_instances = 0
    number_neg_instances = 0

    for edge in dataset.edges():
        if edge.real_target == +1:
            number_pos_instances += 1
        else:
            number_neg_instances += 1

    number_total_instances = number_pos_instances + number_neg_instances
    percentage_pos_instances = number_pos_instances / number_total_instances
    percentage_neg_instances = number_neg_instances / number_total_instances

    first_ent_component = -1 * (percentage_pos_instances * log2(percentage_pos_instances) + percentage_neg_instances * log2(percentage_neg_instances))
    feature_list = []
    for key, value in feature_set.items():
        feature_present_in_pos = 0
        feature_present_in_neg = 0
        feature_absent_in_pos = 0
        feature_absent_in_neg = 0
        total_feature_present = 0
        total_feature_absent = 0

        for edge in dataset.edges():
            if edge.real_target == +1:
                if value in edge.features.keys():
                    feature_present_in_pos += 1
                    total_feature_present += 1
                else:
                    feature_absent_in_pos += 1
                    total_feature_absent +=1
            if edge.real_target == -1:
                if value in edge.features.keys():
                    feature_present_in_neg += 1
                    total_feature_present += 1
                else:
                    feature_absent_in_neg += 1
                    total_feature_absent += 1

        percentage_pos_given_feature = 0
        percentage_neg_given_feature = 0
        if (total_feature_present > 0):
            percentage_pos_given_feature = feature_present_in_pos / total_feature_present
            percentage_neg_given_feature = feature_present_in_neg / total_feature_present

        percentage_pos_given_feature_log = 0
        percentage_neg_given_feature_log = 0
        if percentage_pos_given_feature > 0:
            percentage_pos_given_feature_log = log2(percentage_pos_given_feature)
        if percentage_neg_given_feature > 0:
            percentage_neg_given_feature_log = log2(percentage_neg_given_feature)

        second_emp_component_factor = percentage_pos_given_feature * percentage_pos_given_feature_log + \
                            percentage_neg_given_feature * percentage_neg_given_feature_log

        percentage_feature_given_pos = feature_present_in_pos / number_pos_instances
        percentage_feature_given_neg = feature_present_in_pos / number_neg_instances
        percentage_feature = percentage_feature_given_pos * percentage_pos_instances + \
                    percentage_feature_given_neg * percentage_neg_instances

        second_ent_component = percentage_feature * second_emp_component_factor
        percentage_pos_given_feature_component = 0
        percentage_neg_given_feature_component = 0
        if total_feature_absent>0:
            percentage_pos_given_feature_component = feature_absent_in_pos / total_feature_absent
            percentage_neg_given_feature_component = feature_absent_in_neg / total_feature_absent

        percentage_pos_given_feature_component_log = 0
        percentage_neg_given_feature_component_log = 0
        if percentage_pos_given_feature_component>0:
            percentage_pos_given_feature_component_log = log2(percentage_pos_given_feature_component)
        if percentage_neg_given_feature_component>0:
            percentage_neg_given_feature_component_log = log2(percentage_neg_given_feature_component)

        third_component_multi_factor = percentage_pos_given_feature_component * percentage_pos_given_feature_component_log + \
                percentage_neg_given_feature_component * percentage_neg_given_feature_component_log

        percentage_feature_comp_given_pos = feature_absent_in_pos / number_pos_instances
        percentage_feature_comp_given_neg = feature_absent_in_neg / number_neg_instances
        percentage_feature_comp = percentage_feature_comp_given_pos * percentage_pos_instances + \
                    percentage_feature_comp_given_neg * percentage_neg_instances

        third_ent_component = percentage_feature_comp * third_component_multi_factor
        entropy = first_ent_component + second_ent_component + third_ent_component

        feature_list.append([key, value, entropy])

    feature_list = sorted(feature_list, key=itemgetter(2), reverse=True)

    return feature_list
<commit_msg>fix name<commit_after>import abc
from nalaf.features import FeatureGenerator
import re
from nalaf import print_debug, print_verbose

class EdgeFeatureGenerator(FeatureGenerator):
    """
    Abstract class for generating features for each edge in the dataset.
    Subclasses that inherit this class should:
    * Be named [Name]FeatureGenerator
    * Implement the abstract method generate
    * Append new items to the dictionary field "features" of each Edge in the dataset
    """

    @abc.abstractmethod
    def generate(self, dataset, feature_set, use_gold=True, use_pred=False):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        pass

    def add_to_feature_set(self, feature_set, edge, feature_name, value=1):
        """
        Return True if feature was added to feature_set. False, otherwise

        If the feature_name is None, the feature is not added in anycase. See: self.mk_feature_name
        """
        if feature_name is None:
            return False

        else:
            feature_name = self.__set_final_name(feature_name)

            if not feature_set.is_locked:
                feature_index = feature_set.get(feature_name, None)

                if feature_index is None:
                    feature_index = len(feature_set)
                    feature_set[feature_name] = feature_index
                    print_verbose("Feature map: {} == {} -- _1st_ value: {}".format(str(feature_index), feature_name, str(value)))

                edge.features[feature_index] = value
                return True

            else:
                feature_index = feature_set.get(feature_name, None)

                if feature_index is not None:
                    edge.features[feature_index] = value
                    return True
                else:
                    return False

    def __set_final_name(self, feature_name):
        if not re.search('\[-?[0-9]+\]$', feature_name):
            # Identify the window position --> TODO likely deletable from a edge feature generator
            feature_name = feature_name + "_[0]"

        if not feature_name.startswith(self.__class__.__name__):
            feature_name = self.__class__.__name__ + "::" + feature_name

        return feature_name

    def mk_feature_name(self, prefix, *args):
        if prefix is None:
            return None
        else:
            l = [str(x) for x in ([prefix] + list(args))]
            return "_".join(l)

    def gen_prefix_feat_name(self, field_prefix_feature, *args):
        prefix = self.__getattribute__(field_prefix_feature)
        pure_name = field_prefix_feature[field_prefix_feature.find("_") + 1:]  # Remove "prefix_"
        feature_name = self.mk_feature_name(prefix, pure_name, *args)
        # print_debug(feature_name, field_prefix_feature, args)
        return feature_name

    def add(self, feature_set, edge, field_prefix_feature, *args):
        feature_name = self.gen_prefix_feat_name(field_prefix_feature, *args)
        self.add_to_feature_set(feature_set, edge, feature_name)

    def add_with_value(self, feature_set, edge, field_prefix_feature, value, *args):
        feature_name = self.gen_prefix_feat_name(field_prefix_feature, *args)
        self.add_to_feature_set(feature_set, edge, feature_name, value=value)

from nalaf.features.relations import EdgeFeatureGenerator
from nltk.stem import PorterStemmer
from math import log2
from operator import itemgetter

class TokenFeatureGenerator(EdgeFeatureGenerator):
    """
    Token based features for each entity belonging to an edge
    """

    def __init__(
        self,
        prefix_txt=None,  # 73 in relna
        prefix_pos=None,  # 74
        prefix_masked_txt=None,  # 75
        prefix_stem_masked_txt=None,  # 76
        prefix_ann_type=None,  # 77
    ):
        self.stemmer = PorterStemmer()
        """an instance of the PorterStemmer()"""

        self.prefix_txt = prefix_txt
        self.prefix_pos = prefix_pos
        self.prefix_masked_txt = prefix_masked_txt
        self.prefix_stem_masked_txt = prefix_stem_masked_txt
        self.prefix_ann_type = prefix_ann_type

    @abc.abstractmethod
    def generate(self, dataset, feature_set, is_training_mode):
        """
        Does nothing directly
        """
        pass

    def token_features(self, token, addendum, edge, feature_set, is_training_mode):
        feature_name_1 = self.gen_prefix_feat_name("prefix_txt", addendum, token.word)
        self.add_to_feature_set(feature_set, edge, feature_name_1)

        feature_name_2 = self.gen_prefix_feat_name("prefix_pos", addendum, token.features['pos'])
        self.add_to_feature_set(feature_set, edge, feature_name_2)

        feature_name_3 = self.gen_prefix_feat_name("prefix_masked_txt", addendum, token.masked_text(edge.same_part))
        self.add_to_feature_set(feature_set, edge, feature_name_3)

        # TODO why stem of masked text? -- makes little sense -- See TODO in original loctext too
        feature_name_4 = self.gen_prefix_feat_name("prefix_stem_masked_txt", addendum, self.stemmer.stem(token.masked_text(edge.same_part)))
        self.add_to_feature_set(feature_set, edge, feature_name_4)

        ann_types = self.annotated_types(token, edge)
        for ann in ann_types:
            feature_name_5 = self.gen_prefix_feat_name("prefix_ann_type", addendum, ann)
            self.add_to_feature_set(feature_set, edge, feature_name_5)

    def annotated_types(self, token, edge):
        head1 = edge.entity1.head_token
        head2 = edge.entity2.head_token

        if not token.is_entity_part(edge.same_part):
            feature_name = 'no_ann_type'
            return [feature_name]
        else:
            ann_types = []
            if token.is_entity_part(edge.same_part):
                entity = token.get_entity(edge.same_part)
                feature_name_1 = entity.class_id
                ann_types.append(feature_name_1)
                if entity == edge.entity1:
                    feature_name_2 = 'entity1_' + edge.entity1.class_id
                    ann_types.append(feature_name_2)
                    return ann_types
                elif entity == edge.entity2:
                    feature_name_2 = 'entity2_' + edge.entity2.class_id
                    ann_types.append(feature_name_2)
                    return ann_types
                return ann_types

def calculateInformationGain(feature_set, dataset, output_file):
    number_pos_instances = 0
    number_neg_instances = 0

    for edge in dataset.edges():
        if edge.real_target == +1:
            number_pos_instances += 1
        else:
            number_neg_instances += 1

    number_total_instances = number_pos_instances + number_neg_instances
    percentage_pos_instances = number_pos_instances / number_total_instances
    percentage_neg_instances = number_neg_instances / number_total_instances

    first_ent_component = -1 * (percentage_pos_instances * log2(percentage_pos_instances) + percentage_neg_instances * log2(percentage_neg_instances))
    feature_list = []
    for key, value in feature_set.items():
        feature_present_in_pos = 0
        feature_present_in_neg = 0
        feature_absent_in_pos = 0
        feature_absent_in_neg = 0
        total_feature_present = 0
        total_feature_absent = 0

        for edge in dataset.edges():
            if edge.real_target == +1:
                if value in edge.features.keys():
                    feature_present_in_pos += 1
                    total_feature_present += 1
                else:
                    feature_absent_in_pos += 1
                    total_feature_absent +=1
            if edge.real_target == -1:
                if value in edge.features.keys():
                    feature_present_in_neg += 1
                    total_feature_present += 1
                else:
                    feature_absent_in_neg += 1
                    total_feature_absent += 1

        percentage_pos_given_feature = 0
        percentage_neg_given_feature = 0
        if (total_feature_present > 0):
            percentage_pos_given_feature = feature_present_in_pos / total_feature_present
            percentage_neg_given_feature = feature_present_in_neg / total_feature_present

        percentage_pos_given_feature_log = 0
        percentage_neg_given_feature_log = 0
        if percentage_pos_given_feature > 0:
            percentage_pos_given_feature_log = log2(percentage_pos_given_feature)
        if percentage_neg_given_feature > 0:
            percentage_neg_given_feature_log = log2(percentage_neg_given_feature)

        second_emp_component_factor = percentage_pos_given_feature * percentage_pos_given_feature_log + \
                            percentage_neg_given_feature * percentage_neg_given_feature_log

        percentage_feature_given_pos = feature_present_in_pos / number_pos_instances
        percentage_feature_given_neg = feature_present_in_pos / number_neg_instances
        percentage_feature = percentage_feature_given_pos * percentage_pos_instances + \
                    percentage_feature_given_neg * percentage_neg_instances

        second_ent_component = percentage_feature * second_emp_component_factor
        percentage_pos_given_feature_component = 0
        percentage_neg_given_feature_component = 0
        if total_feature_absent>0:
            percentage_pos_given_feature_component = feature_absent_in_pos / total_feature_absent
            percentage_neg_given_feature_component = feature_absent_in_neg / total_feature_absent

        percentage_pos_given_feature_component_log = 0
        percentage_neg_given_feature_component_log = 0
        if percentage_pos_given_feature_component>0:
            percentage_pos_given_feature_component_log = log2(percentage_pos_given_feature_component)
        if percentage_neg_given_feature_component>0:
            percentage_neg_given_feature_component_log = log2(percentage_neg_given_feature_component)

        third_component_multi_factor = percentage_pos_given_feature_component * percentage_pos_given_feature_component_log + \
                percentage_neg_given_feature_component * percentage_neg_given_feature_component_log

        percentage_feature_comp_given_pos = feature_absent_in_pos / number_pos_instances
        percentage_feature_comp_given_neg = feature_absent_in_neg / number_neg_instances
        percentage_feature_comp = percentage_feature_comp_given_pos * percentage_pos_instances + \
                    percentage_feature_comp_given_neg * percentage_neg_instances

        third_ent_component = percentage_feature_comp * third_component_multi_factor
        entropy = first_ent_component + second_ent_component + third_ent_component

        feature_list.append([key, value, entropy])

    feature_list = sorted(feature_list, key=itemgetter(2), reverse=True)

    return feature_list
<|endoftext|>


<commit_before>from base import Gateway
from StringIO import StringIO
import xml.etree.ElementTree as ET

from post import post

class AuthorizeNet(Gateway):
    API_VERSION = '3.1'

    APPROVED, DECLINED, ERROR, FRAUD_REVIEW = 1, 2, 3, 4
    RESPONSE_CODE, RESPONSE_REASON_CODE, RESPONSE_REASON_TEXT = 0, 2, 3
    AVS_RESULT_CODE, TRANSACTION_ID, CARD_CODE_RESPONSE_CODE = 5, 6, 38
    RECURRING_ACTIONS = {
      u'create': 'ARBCreateSubscription',
      u'update': 'ARBUpdateSubscription',
      u'cancel': 'ARBCancelSubscription'
    }
    test_url = "https://test.authorize.net/gateway/transact.dll"
    live_url = "https://secure.authorize.net/gateway/transact.dll"

    arb_test_url = 'https://apitest.authorize.net/xml/v1/request.api'
    arb_live_url = 'https://api.authorize.net/xml/v1/request.api'

    supported_countries = ['US']
    supported_cardtypes = ['visa', 'master', 'american_express', 'discover']
    homepage_url = 'http://www.authorize.net/'
    display_name = u'Authorize.Net'

    CARD_CODE_ERRORS = ('N', 'S')
    AVS_ERRORS = ('A', 'E', 'N', 'R', 'W', 'Z')

    AUTHORIZE_NET_ARB_NAMESPACE = 'AnetApi/xml/v1/schema/AnetApiSchema.xsd'

    def __init__(self, options, gateway_mode='live'):
        assert(options.has_key('login'))
        assert(options.has_key('password'))
        self.options = options
        super(AuthorizeNet, self).__init__(gateway_mode=gateway_mode)

    def recurring(self, money, creditcard, **kwargs):
        assert(kwargs.has_key('interval') and
               kwargs.has_key('duration') and
               kwargs.has_key('billing_address'))
        assert(kwargs['interval'].has_key('length') and
               kwargs['interval'].has_key('unit'))
        assert(kwargs['interval']['unit'] == 'days' or
               kwargs['interval']['unit'] == 'months')
        assert(kwargs['duration'].has_key('start_date') or
               kwargs['duration'].has_key('occurrences'))
        assert(kwargs['billing_address'].has_key('first_name') or
               kwargs['billing_address'].has_key('last_name'))

        #Create a copy of kwargs and call it options. This is not
        #intended to be self.options. ARB has different options
        #than the standard Authorize.Net payment gateway
        options = kwargs.copy()
        options['credit_card'] = creditcard
        options['amount'] = money

        request = self.build_recurring_request(u'create', options)
        return self.recurring_commit(u'create', request)

    def update_recurring(self, money, creditcard, **kwargs):
        assert(kwargs.has_key('subscription_id'))

        #Create a copy of kwargs and call it options. This is not
        #intended to be self.options. ARB has different options
        #than the standard Authorize.Net payment gateway
        options = kwargs.copy()
        options['credit_card'] = creditcard
        options['amount'] = money

        request = self.build_recurring_request(u'update', options)
        return self.recurring_commit(u'update', request)

    def cancel_recurring(self, subscription_id):
        request = self.build_recurring_request('cancel', {'subscription_id': subscription_id})
        return self.recurring_commit(u'cancel', request)

    def build_recurring_request(self, action, options):
        if not self.RECURRING_ACTIONS.has_key(action):
            raise StandardError, u"Invalid Automated Recurring Billing Action: " + unicode(action)

        requestbuff = StringIO()

        requestbuff.write('<?xml version="1.0" encoding="utf-8"?>')
        root = ET.Element(self.RECURRING_ACTIONS[action] + 'Request',
                          attrib={'xmlns': self.AUTHORIZE_NET_ARB_NAMESPACE})
        self.add_arb_merchant_authentication(root)
        if options.has_key('ref_id'):
            refId = ET.SubElement(root, 'refId')
            refId.text = options['ref_id']
        {'create': self.build_arb_create_subscription_request,
         'update': self.build_arb_update_subscription_request,
         'cancel': self.build_arb_cancel_subscription_request
        }[action](root, options)

        tree = ET.ElementTree(root)
        tree.write(requestbuff)

        return requestbuff.getvalue()

    def add_se(self, node, name, text):
        subnode = ET.SubElement(node, name)
        subnode.text = text
        return subnode

    def add_arb_merchant_authentication(self, root):
        merch_auth = ET.SubElement(root, 'merchantAuthentication')
        self.add_se(merch_auth, 'name', self.options['login'])
        self.add_se(merch_auth, 'transactionKey', self.options['password'])

    def build_arb_create_subscription_request(self, root, options):
        self.add_arb_subscription(root, options)

    def build_arb_update_subscription_request(self, root, options):
        subscription_id = self.add_se(root, 'subscriptionId', options['subscription_id'])
        self.add_arb_subscription(root, options)

    def build_arb_cancel_subscription_request(self, root, options):
        self.add_se(root, 'subscriptionId', options['subscription_id'])

    def add_arb_subscription(self, root, options):
        subscription = ET.SubElement(root, 'subscription')
        if options.has_key('subscription_name'):
            self.add_se(subscription, 'name', options['subscription_name'])
        self.add_arb_payment_schedule(subscription, options)
        if options.has_key('amount'):
            self.add_se(subscription, 'amount', self.amount(options['amount']))
        if options.setdefault('trial', None) and options['trial'].setdefault('amount', 0):
            self.add_se(subscription, 'trailAmount', options['trial']['amount'])
        self.add_arb_payment(subscription, options)
        self.add_arb_order(subscription, options)
        self.add_arb_customer(subscription, options)
        self.add_arb_address(subscription, 'billTo', options['billing_address'])
        if 'shipping_address' in options.keys():
            self.add_arb_address(subscription, 'shipTo', options['shipping_address'])

    def add_arb_interval(self, payment_schedule, options):
        if not 'interval' in options.keys() or not options['interval']:
            return
        interval = ET.SubElement(payment_schedule, 'interval')
        self.add_se(interval, 'length', options['interval']['length'])
        self.add_se(interval, 'unit', options['interval']['unit'])

    def add_arb_duration(self, payment_schedule, options):
        if not 'duration' in options.keys() or not options['duration']:
            return
        self.add_se(payment_schedule, 'startDate', options['duration']['start_date'])
        self.add_se(payment_schedule, 'totalOccurrences', options['duration']['occurrences'])

    def add_arb_payment_schedule(self, subscription, options):
        if 'interval' in options.keys() or 'duration' in options.keys():
            payment_schedule = ET.SubElement(subscription, 'paymentSchedule')
            self.add_arb_interval(payment_schedule, options)
            self.add_arb_duration(payment_schedule, options)
            if 'trial' in options.keys() and options['trial']:
                self.add_se(payment_schedule, 'trialOccurrences', options['trial']['trialOccurrences'])

    def add_arb_payment(self, subscription, options):
        #if not 'credit_card' in options.keys() or 'bank_account' in options.keys():
        #    return
        payment = ET.SubElement(subscription, 'payment')
        self.add_arb_credit_card(payment, options)
        self.add_arb_bank_account(payment, options)

    def add_arb_credit_card(self, payment, options):
        if not 'credit_card' in options.keys() or not options['credit_card']:
            return
        credit_card = ET.SubElement(payment, 'creditCard')
        self.add_se(credit_card, 'cardNumber', options['credit_card'].number)
        self.add_se(credit_card, 'expirationDate', self.arb_expdate(options['credit_card']))

    def add_arb_bank_account(self, payment, options):
        if not 'bank_account' in options.keys() or not options['bank_account']:
            return
        bank_account = ET.SubElement(payment, 'bankAccount')
        self.add_se(bank_account, 'accountType', options['bank_account']['account_type'])
        self.add_se(bank_account, 'routingNumber', options['bank_account']['routing_number'])
        self.add_se(bank_account, 'accountNumber', options['bank_account']['account_number'])
        self.add_se(bank_account, 'nameOfAccount', options['bank_account']['name_of_account'])
        if 'bank_name' in options['bank_account'].keys():
            self.add_se(bank_account, 'bankName', options['bank_account']['bank_name'])
        self.add_se(bank_account, 'echeckType', options['bank_account']['echeck_type'])

    def add_arb_order(self, root, options):
        if not 'order' in options.keys() or not options['order']:
            return
        order = ET.SubElement(root, 'order')
        self.add_se(order, 'invoiceNumber', options['order']['invoice_number'])
        self.add_se(order, 'description', options['order']['description'])

    def add_arb_customer(self, root, options):
        if not 'customer' in options.keys() or not options['customer']:
            return
        customer = ET.SubElement(root, 'customer')
        if 'type' in options['customer'].keys():
            self.add_se(customer, 'type', options['customer']['type'])
        if 'id' in options['customer'].keys():
            self.add_se(customer, 'id', options['customer']['id'])
        if 'email' in options['customer'].keys():
            self.add_se(customer, 'email', options['customer']['email'])
        if 'phone_number' in options['customer'].keys():
            self.add_se(customer, 'phoneNumber', options['customer']['phone_number'])
        if 'fax_number' in options['customer'].keys():
            self.add_se(customer, 'faxNumber', options['customer']['fax_number'])
        self.add_arb_drivers_license(customer, options)
        if 'tax_id' in options['customer'].keys():
            self.add_se(customer, 'taxId', options['customer']['tax_id'])

    def add_arb_drivers_license(self, customer, options):
        if not 'customer' in options.keys() or not options['customer']:
            return
        if not 'drivers_license' in options['customer'].keys() or not options['customer']['drivers_license']:
            return
        drivers_license = ET.SubElement(customer, 'driversLicense')
        self.add_se(drivers_license, 'number', options['drivers_license']['number'])
        self.add_se(drivers_license, 'state', options['drivers_license']['state'])
        self.add_se(drivers_license, 'dateOfBirth', options['drivers_license']['date_of_birth'])

    def add_arb_address(self, root, container_name, address):
        if not address.keys():
            return
        container = ET.SubElement(root, container_name)
        self.add_se(container, 'firstName', address['first_name'])
        self.add_se(container, 'lastName', address['last_name'])

        if 'company' in address.keys():
            self.add_se(container, 'company', address['company'])
        if 'address1' in address.keys():
            self.add_se(container, 'address', address['address1'])
        if 'city' in address.keys():
            self.add_se(container, 'city', address['city'])
        if 'state' in address.keys():
            self.add_se(container, 'state', address['state'])
        if 'zip' in address.keys():
            self.add_se(container, 'zip', address['zip'])
        if 'country' in address.keys():
            self.add_se(container, 'country', address['country'])

    def arb_expdate(self, credit_card):
        return "%04d-%02d" % (credit_card.year, credit_card.month)

    def recurring_commit(self, action, request):
        if self.is_test():
            url = self.arb_test_url
        else:
            url = self.arb_live_url
        xml = post(url, request, {"Content-Type": "text/xml"})
        print 'RESPONSE: ', xml
        response = self.recurring_parse('create', xml.decode('utf-8-sig'))
        return response

    def normalize(self, name):
        if name[0] == "{":
            uri, tag = name[1:].split("}")
            return tag
        else:
            return name

    def recurring_parse(self, action, xml):
        response = {}
        xml = ET.parse(StringIO(xml))
        root = xml.getroot()
        root_tag = self.normalize(root.tag)
        print "ROOT TAG: ", self.normalize(root_tag)

        if root_tag == "%sResponse"%(self.RECURRING_ACTIONS[action]) or root_tag == "ErrorResponse":
            response[root_tag] = True
            for child in root.getchildren():
                self.recurring_parse_element(response, child)
        return response

    def recurring_parse_element(self, response, node):
        if node.getchildren():
            for child in node.getchildren():
                self.recurring_parse_element(response, child)
        else:
            response[self.normalize(node.tag)] = node.text
<commit_msg>Fixed miss named key in the trial dictionary for authorize.net ARB gateway<commit_after>from base import Gateway
from StringIO import StringIO
import xml.etree.ElementTree as ET

from post import post

class AuthorizeNet(Gateway):
    API_VERSION = '3.1'

    APPROVED, DECLINED, ERROR, FRAUD_REVIEW = 1, 2, 3, 4
    RESPONSE_CODE, RESPONSE_REASON_CODE, RESPONSE_REASON_TEXT = 0, 2, 3
    AVS_RESULT_CODE, TRANSACTION_ID, CARD_CODE_RESPONSE_CODE = 5, 6, 38
    RECURRING_ACTIONS = {
      u'create': 'ARBCreateSubscription',
      u'update': 'ARBUpdateSubscription',
      u'cancel': 'ARBCancelSubscription'
    }
    test_url = "https://test.authorize.net/gateway/transact.dll"
    live_url = "https://secure.authorize.net/gateway/transact.dll"

    arb_test_url = 'https://apitest.authorize.net/xml/v1/request.api'
    arb_live_url = 'https://api.authorize.net/xml/v1/request.api'

    supported_countries = ['US']
    supported_cardtypes = ['visa', 'master', 'american_express', 'discover']
    homepage_url = 'http://www.authorize.net/'
    display_name = u'Authorize.Net'

    CARD_CODE_ERRORS = ('N', 'S')
    AVS_ERRORS = ('A', 'E', 'N', 'R', 'W', 'Z')

    AUTHORIZE_NET_ARB_NAMESPACE = 'AnetApi/xml/v1/schema/AnetApiSchema.xsd'

    def __init__(self, options, gateway_mode='live'):
        assert(options.has_key('login'))
        assert(options.has_key('password'))
        self.options = options
        super(AuthorizeNet, self).__init__(gateway_mode=gateway_mode)

    def recurring(self, money, creditcard, **kwargs):
        assert(kwargs.has_key('interval') and
               kwargs.has_key('duration') and
               kwargs.has_key('billing_address'))
        assert(kwargs['interval'].has_key('length') and
               kwargs['interval'].has_key('unit'))
        assert(kwargs['interval']['unit'] == 'days' or
               kwargs['interval']['unit'] == 'months')
        assert(kwargs['duration'].has_key('start_date') or
               kwargs['duration'].has_key('occurrences'))
        assert(kwargs['billing_address'].has_key('first_name') or
               kwargs['billing_address'].has_key('last_name'))

        #Create a copy of kwargs and call it options. This is not
        #intended to be self.options. ARB has different options
        #than the standard Authorize.Net payment gateway
        options = kwargs.copy()
        options['credit_card'] = creditcard
        options['amount'] = money

        request = self.build_recurring_request(u'create', options)
        return self.recurring_commit(u'create', request)

    def update_recurring(self, money, creditcard, **kwargs):
        assert(kwargs.has_key('subscription_id'))

        #Create a copy of kwargs and call it options. This is not
        #intended to be self.options. ARB has different options
        #than the standard Authorize.Net payment gateway
        options = kwargs.copy()
        options['credit_card'] = creditcard
        options['amount'] = money

        request = self.build_recurring_request(u'update', options)
        return self.recurring_commit(u'update', request)

    def cancel_recurring(self, subscription_id):
        request = self.build_recurring_request('cancel', {'subscription_id': subscription_id})
        return self.recurring_commit(u'cancel', request)

    def build_recurring_request(self, action, options):
        if not self.RECURRING_ACTIONS.has_key(action):
            raise StandardError, u"Invalid Automated Recurring Billing Action: " + unicode(action)

        requestbuff = StringIO()

        requestbuff.write('<?xml version="1.0" encoding="utf-8"?>')
        root = ET.Element(self.RECURRING_ACTIONS[action] + 'Request',
                          attrib={'xmlns': self.AUTHORIZE_NET_ARB_NAMESPACE})
        self.add_arb_merchant_authentication(root)
        if options.has_key('ref_id'):
            refId = ET.SubElement(root, 'refId')
            refId.text = options['ref_id']
        {'create': self.build_arb_create_subscription_request,
         'update': self.build_arb_update_subscription_request,
         'cancel': self.build_arb_cancel_subscription_request
        }[action](root, options)

        tree = ET.ElementTree(root)
        tree.write(requestbuff)

        return requestbuff.getvalue()

    def add_se(self, node, name, text):
        subnode = ET.SubElement(node, name)
        subnode.text = text
        return subnode

    def add_arb_merchant_authentication(self, root):
        merch_auth = ET.SubElement(root, 'merchantAuthentication')
        self.add_se(merch_auth, 'name', self.options['login'])
        self.add_se(merch_auth, 'transactionKey', self.options['password'])

    def build_arb_create_subscription_request(self, root, options):
        self.add_arb_subscription(root, options)

    def build_arb_update_subscription_request(self, root, options):
        subscription_id = self.add_se(root, 'subscriptionId', options['subscription_id'])
        self.add_arb_subscription(root, options)

    def build_arb_cancel_subscription_request(self, root, options):
        self.add_se(root, 'subscriptionId', options['subscription_id'])

    def add_arb_subscription(self, root, options):
        subscription = ET.SubElement(root, 'subscription')
        if options.has_key('subscription_name'):
            self.add_se(subscription, 'name', options['subscription_name'])
        self.add_arb_payment_schedule(subscription, options)
        if options.has_key('amount'):
            self.add_se(subscription, 'amount', self.amount(options['amount']))
        if options.setdefault('trial', None) and options['trial'].setdefault('amount', 0):
            self.add_se(subscription, 'trailAmount', options['trial']['amount'])
        self.add_arb_payment(subscription, options)
        self.add_arb_order(subscription, options)
        self.add_arb_customer(subscription, options)
        self.add_arb_address(subscription, 'billTo', options['billing_address'])
        if 'shipping_address' in options.keys():
            self.add_arb_address(subscription, 'shipTo', options['shipping_address'])

    def add_arb_interval(self, payment_schedule, options):
        if not 'interval' in options.keys() or not options['interval']:
            return
        interval = ET.SubElement(payment_schedule, 'interval')
        self.add_se(interval, 'length', options['interval']['length'])
        self.add_se(interval, 'unit', options['interval']['unit'])

    def add_arb_duration(self, payment_schedule, options):
        if not 'duration' in options.keys() or not options['duration']:
            return
        self.add_se(payment_schedule, 'startDate', options['duration']['start_date'])
        self.add_se(payment_schedule, 'totalOccurrences', options['duration']['occurrences'])

    def add_arb_payment_schedule(self, subscription, options):
        if 'interval' in options.keys() or 'duration' in options.keys():
            payment_schedule = ET.SubElement(subscription, 'paymentSchedule')
            self.add_arb_interval(payment_schedule, options)
            self.add_arb_duration(payment_schedule, options)
            if 'trial' in options.keys() and options['trial']:
                self.add_se(payment_schedule, 'trialOccurrences', options['trial']['occurrences'])

    def add_arb_payment(self, subscription, options):
        #if not 'credit_card' in options.keys() or 'bank_account' in options.keys():
        #    return
        payment = ET.SubElement(subscription, 'payment')
        self.add_arb_credit_card(payment, options)
        self.add_arb_bank_account(payment, options)

    def add_arb_credit_card(self, payment, options):
        if not 'credit_card' in options.keys() or not options['credit_card']:
            return
        credit_card = ET.SubElement(payment, 'creditCard')
        self.add_se(credit_card, 'cardNumber', options['credit_card'].number)
        self.add_se(credit_card, 'expirationDate', self.arb_expdate(options['credit_card']))

    def add_arb_bank_account(self, payment, options):
        if not 'bank_account' in options.keys() or not options['bank_account']:
            return
        bank_account = ET.SubElement(payment, 'bankAccount')
        self.add_se(bank_account, 'accountType', options['bank_account']['account_type'])
        self.add_se(bank_account, 'routingNumber', options['bank_account']['routing_number'])
        self.add_se(bank_account, 'accountNumber', options['bank_account']['account_number'])
        self.add_se(bank_account, 'nameOfAccount', options['bank_account']['name_of_account'])
        if 'bank_name' in options['bank_account'].keys():
            self.add_se(bank_account, 'bankName', options['bank_account']['bank_name'])
        self.add_se(bank_account, 'echeckType', options['bank_account']['echeck_type'])

    def add_arb_order(self, root, options):
        if not 'order' in options.keys() or not options['order']:
            return
        order = ET.SubElement(root, 'order')
        self.add_se(order, 'invoiceNumber', options['order']['invoice_number'])
        self.add_se(order, 'description', options['order']['description'])

    def add_arb_customer(self, root, options):
        if not 'customer' in options.keys() or not options['customer']:
            return
        customer = ET.SubElement(root, 'customer')
        if 'type' in options['customer'].keys():
            self.add_se(customer, 'type', options['customer']['type'])
        if 'id' in options['customer'].keys():
            self.add_se(customer, 'id', options['customer']['id'])
        if 'email' in options['customer'].keys():
            self.add_se(customer, 'email', options['customer']['email'])
        if 'phone_number' in options['customer'].keys():
            self.add_se(customer, 'phoneNumber', options['customer']['phone_number'])
        if 'fax_number' in options['customer'].keys():
            self.add_se(customer, 'faxNumber', options['customer']['fax_number'])
        self.add_arb_drivers_license(customer, options)
        if 'tax_id' in options['customer'].keys():
            self.add_se(customer, 'taxId', options['customer']['tax_id'])

    def add_arb_drivers_license(self, customer, options):
        if not 'customer' in options.keys() or not options['customer']:
            return
        if not 'drivers_license' in options['customer'].keys() or not options['customer']['drivers_license']:
            return
        drivers_license = ET.SubElement(customer, 'driversLicense')
        self.add_se(drivers_license, 'number', options['drivers_license']['number'])
        self.add_se(drivers_license, 'state', options['drivers_license']['state'])
        self.add_se(drivers_license, 'dateOfBirth', options['drivers_license']['date_of_birth'])

    def add_arb_address(self, root, container_name, address):
        if not address.keys():
            return
        container = ET.SubElement(root, container_name)
        self.add_se(container, 'firstName', address['first_name'])
        self.add_se(container, 'lastName', address['last_name'])

        if 'company' in address.keys():
            self.add_se(container, 'company', address['company'])
        if 'address1' in address.keys():
            self.add_se(container, 'address', address['address1'])
        if 'city' in address.keys():
            self.add_se(container, 'city', address['city'])
        if 'state' in address.keys():
            self.add_se(container, 'state', address['state'])
        if 'zip' in address.keys():
            self.add_se(container, 'zip', address['zip'])
        if 'country' in address.keys():
            self.add_se(container, 'country', address['country'])

    def arb_expdate(self, credit_card):
        return "%04d-%02d" % (credit_card.year, credit_card.month)

    def recurring_commit(self, action, request):
        if self.is_test():
            url = self.arb_test_url
        else:
            url = self.arb_live_url
        xml = post(url, request, {"Content-Type": "text/xml"})
        print 'RESPONSE: ', xml
        response = self.recurring_parse('create', xml.decode('utf-8-sig'))
        return response

    def normalize(self, name):
        if name[0] == "{":
            uri, tag = name[1:].split("}")
            return tag
        else:
            return name

    def recurring_parse(self, action, xml):
        response = {}
        xml = ET.parse(StringIO(xml))
        root = xml.getroot()
        root_tag = self.normalize(root.tag)
        print "ROOT TAG: ", self.normalize(root_tag)

        if root_tag == "%sResponse"%(self.RECURRING_ACTIONS[action]) or root_tag == "ErrorResponse":
            response[root_tag] = True
            for child in root.getchildren():
                self.recurring_parse_element(response, child)
        return response

    def recurring_parse_element(self, response, node):
        if node.getchildren():
            for child in node.getchildren():
                self.recurring_parse_element(response, child)
        else:
            response[self.normalize(node.tag)] = node.text
<|endoftext|>


<commit_before>import os
import io
import re
from contextlib import redirect_stdout

from nose.tools import assert_greater_equal, assert_true, assert_equal

import numpy as np
from scipy import sparse
from sklearn.neighbors import KDTree
from sklearn.preprocessing import normalize

from pynndescent import NNDescent, PyNNDescentTransformer

np.random.seed(42)
spatial_data = np.random.randn(10, 20)
spatial_data = np.vstack(
    [spatial_data, np.zeros((2, 20))]
)  # Add some all zero data for corner case test

nn_data = np.random.uniform(0, 1, size=(1000, 5))
nn_data = np.vstack(
    [nn_data, np.zeros((2, 5))]
)  # Add some all zero data for corner case test
# for_sparse_nn_data = np.random.uniform(0, 1, size=(1002, 500))
# binary_nn_data = np.random.choice(a=[False, True], size=(1000, 500), p=[0.1, 1 - 0.1])
# binary_nn_data = np.vstack(
#     [binary_nn_data, np.zeros((2, 500))]
# )  # Add some all zero data for corner case test
# sparse_nn_data = sparse.csr_matrix(for_sparse_nn_data * binary_nn_data)
sparse_nn_data = sparse.random(1000, 50, density=0.5, format="csr")
# sparse_nn_data = sparse.csr_matrix(nn_data)

def test_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        nn_data, "euclidean", {}, 10, random_state=np.random
    )._neighbor_graph

    tree = KDTree(nn_data)
    true_indices = tree.query(nn_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.98,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_angular_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        nn_data, "cosine", {}, 10, random_state=np.random
    )._neighbor_graph

    angular_data = normalize(nn_data, norm="l2")
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.98,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_sparse_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        sparse_nn_data, "euclidean", n_neighbors=20, random_state=np.random
    )._neighbor_graph

    tree = KDTree(sparse_nn_data.toarray())
    true_indices = tree.query(sparse_nn_data.toarray(), 10, return_distance=False)

    num_correct = 0.0
    for i in range(sparse_nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (sparse_nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.85,
        "Sparse NN-descent did not get 95% " "accuracy on nearest neighbors",
    )

def test_sparse_angular_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        sparse_nn_data, "cosine", {}, 20, random_state=np.random
    )._neighbor_graph

    angular_data = normalize(sparse_nn_data, norm="l2").toarray()
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(sparse_nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (sparse_nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.85,
        "Sparse angular NN-descent did not get 98% " "accuracy on nearest neighbors",
    )

def test_nn_descent_query_accuracy():
    nnd = NNDescent(nn_data[200:], "euclidean", n_neighbors=10, random_state=np.random)
    knn_indices, _ = nnd.query(nn_data[:200], k=10)

    tree = KDTree(nn_data[200:])
    true_indices = tree.query(nn_data[:200], 10, return_distance=False)

    num_correct = 0.0
    for i in range(true_indices.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (true_indices.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.99,
        "NN-descent query did not get 99% " "accuracy on nearest neighbors",
    )

def test_sparse_nn_descent_query_accuracy():
    nnd = NNDescent(
        sparse_nn_data[200:], "euclidean", n_neighbors=10, random_state=np.random
    )
    knn_indices, _ = nnd.query(sparse_nn_data[:200], k=10)

    tree = KDTree(sparse_nn_data[200:].toarray())
    true_indices = tree.query(sparse_nn_data[:200].toarray(), 10, return_distance=False)

    num_correct = 0.0
    for i in range(true_indices.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (true_indices.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.95,
        "Sparse NN-descent query did not get 97% " "accuracy on nearest neighbors",
    )

def test_random_state_none():
    knn_indices, _ = NNDescent(
        nn_data, "euclidean", {}, 10, random_state=None
    )._neighbor_graph

    tree = KDTree(nn_data)
    true_indices = tree.query(nn_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (spatial_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.99,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_deterministic():
    seed = np.random.RandomState(42)

    x1 = seed.normal(0, 100, (1000, 50))
    x2 = seed.normal(0, 100, (1000, 50))

    index1 = NNDescent(x1, random_state=np.random.RandomState(42))
    neighbors1, distances1 = index1.query(x2)

    index2 = NNDescent(x1, random_state=np.random.RandomState(42))
    neighbors2, distances2 = index2.query(x2)

    np.testing.assert_equal(neighbors1, neighbors2)
    np.testing.assert_equal(distances1, distances2)

# This tests a recursion error on cosine metric reported at:
# https://github.com/lmcinnes/umap/issues/99
# data used is a cut-down version of that provided by @scharron
# It contains lots of all-zero vectors and some other duplicates
def test_rp_trees_should_not_stack_overflow_with_duplicate_data():
    this_dir = os.path.dirname(os.path.abspath(__file__))
    data_path = os.path.join(this_dir, "test_data/cosine_hang.npy")
    data = np.load(data_path)

    n_neighbors = 10
    knn_indices, _ = NNDescent(
        data, "cosine", {}, n_neighbors, random_state=np.random, n_trees=20
    )._neighbor_graph

    for i in range(data.shape[0]):
        assert_equal(
            len(knn_indices[i]),
            len(np.unique(knn_indices[i])),
            "Duplicate indices in knn graph",
        )

def test_deduplicated_data_behaves_normally():
    this_dir = os.path.dirname(os.path.abspath(__file__))
    data_path = os.path.join(this_dir, "test_data/cosine_hang.npy")
    data = np.unique(np.load(data_path), axis=0)
    data = data[~np.all(data == 0, axis=1)]
    data = data[:1000]

    n_neighbors = 10
    knn_indices, _ = NNDescent(
        data, "cosine", {}, n_neighbors, random_state=np.random, n_trees=20
    )._neighbor_graph

    for i in range(data.shape[0]):
        assert_equal(
            len(knn_indices[i]),
            len(np.unique(knn_indices[i])),
            "Duplicate indices in knn graph",
        )

    angular_data = normalize(data, norm="l2")
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, n_neighbors, return_distance=False)

    num_correct = 0
    for i in range(data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    proportion_correct = num_correct / (data.shape[0] * n_neighbors)
    assert_greater_equal(
        proportion_correct,
        0.95,
        "NN-descent did not get 95%" " accuracy on nearest neighbors",
    )

def test_output_when_verbose_is_true():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = NNDescent(
            data=spatial_data,
            metric="euclidean",
            metric_kwds={},
            n_neighbors=4,
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=True,
        )
    output = out.getvalue()
    assert_true(re.match("^.*5 trees", output, re.DOTALL))
    assert_true(re.match("^.*2 iterations", output, re.DOTALL))

def test_no_output_when_verbose_is_false():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = NNDescent(
            data=spatial_data,
            metric="euclidean",
            metric_kwds={},
            n_neighbors=4,
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=False,
        )
    output = out.getvalue().strip()
    assert_equal(len(output), 0)

# same as the previous two test, but this time using the PyNNDescentTransformer
# interface
def test_transformer_output_when_verbose_is_true():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = PyNNDescentTransformer(
            n_neighbors=4,
            metric="euclidean",
            metric_kwds={},
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=True,
        ).fit_transform(spatial_data)
    output = out.getvalue()
    assert_true(re.match("^.*5 trees", output, re.DOTALL))
    assert_true(re.match("^.*2 iterations", output, re.DOTALL))

def test_transformer_output_when_verbose_is_false():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = PyNNDescentTransformer(
            n_neighbors=4,
            metric="euclidean",
            metric_kwds={},
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=False,
        ).fit_transform(spatial_data)
    output = out.getvalue().strip()
    assert_equal(len(output), 0)
<commit_msg>Test equivalence of NNDescent vs PyNNDescentTranformer<commit_after>import os
import io
import re
from contextlib import redirect_stdout

from nose.tools import assert_greater_equal, assert_true, assert_equal

import numpy as np
from scipy import sparse
from sklearn.neighbors import KDTree
from sklearn.preprocessing import normalize

from pynndescent import NNDescent, PyNNDescentTransformer

np.random.seed(42)
spatial_data = np.random.randn(10, 20)
spatial_data = np.vstack(
    [spatial_data, np.zeros((2, 20))]
)  # Add some all zero data for corner case test

nn_data = np.random.uniform(0, 1, size=(1000, 5))
nn_data = np.vstack(
    [nn_data, np.zeros((2, 5))]
)  # Add some all zero data for corner case test
# for_sparse_nn_data = np.random.uniform(0, 1, size=(1002, 500))
# binary_nn_data = np.random.choice(a=[False, True], size=(1000, 500), p=[0.1, 1 - 0.1])
# binary_nn_data = np.vstack(
#     [binary_nn_data, np.zeros((2, 500))]
# )  # Add some all zero data for corner case test
# sparse_nn_data = sparse.csr_matrix(for_sparse_nn_data * binary_nn_data)
sparse_nn_data = sparse.random(1000, 50, density=0.5, format="csr")
# sparse_nn_data = sparse.csr_matrix(nn_data)

def test_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        nn_data, "euclidean", {}, 10, random_state=np.random
    )._neighbor_graph

    tree = KDTree(nn_data)
    true_indices = tree.query(nn_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.98,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_angular_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        nn_data, "cosine", {}, 10, random_state=np.random
    )._neighbor_graph

    angular_data = normalize(nn_data, norm="l2")
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.98,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_sparse_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        sparse_nn_data, "euclidean", n_neighbors=20, random_state=np.random
    )._neighbor_graph

    tree = KDTree(sparse_nn_data.toarray())
    true_indices = tree.query(sparse_nn_data.toarray(), 10, return_distance=False)

    num_correct = 0.0
    for i in range(sparse_nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (sparse_nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.85,
        "Sparse NN-descent did not get 95% " "accuracy on nearest neighbors",
    )

def test_sparse_angular_nn_descent_neighbor_accuracy():
    knn_indices, _ = NNDescent(
        sparse_nn_data, "cosine", {}, 20, random_state=np.random
    )._neighbor_graph

    angular_data = normalize(sparse_nn_data, norm="l2").toarray()
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(sparse_nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (sparse_nn_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.85,
        "Sparse angular NN-descent did not get 98% " "accuracy on nearest neighbors",
    )

def test_nn_descent_query_accuracy():
    nnd = NNDescent(nn_data[200:], "euclidean", n_neighbors=10, random_state=np.random)
    knn_indices, _ = nnd.query(nn_data[:200], k=10)

    tree = KDTree(nn_data[200:])
    true_indices = tree.query(nn_data[:200], 10, return_distance=False)

    num_correct = 0.0
    for i in range(true_indices.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (true_indices.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.99,
        "NN-descent query did not get 99% " "accuracy on nearest neighbors",
    )

def test_sparse_nn_descent_query_accuracy():
    nnd = NNDescent(
        sparse_nn_data[200:], "euclidean", n_neighbors=10, random_state=np.random
    )
    knn_indices, _ = nnd.query(sparse_nn_data[:200], k=10)

    tree = KDTree(sparse_nn_data[200:].toarray())
    true_indices = tree.query(sparse_nn_data[:200].toarray(), 10, return_distance=False)

    num_correct = 0.0
    for i in range(true_indices.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (true_indices.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.95,
        "Sparse NN-descent query did not get 97% " "accuracy on nearest neighbors",
    )

def test_transformer_equivalence():
    N_NEIGHBORS = 15
    QUEUE_SIZE = 5.0
    train = nn_data[:400]
    test = nn_data[:200]

    nnd = NNDescent(
        data=train,
        n_neighbors=N_NEIGHBORS,
        random_state=42
    )
    indices, dists = nnd.query(test, k=N_NEIGHBORS, queue_size=QUEUE_SIZE)
    sort_idx = np.argsort(indices, axis=1)
    indices_sorted = np.vstack(
         [indices[i, sort_idx[i]] for i in range(sort_idx.shape[0])]
    )
    dists_sorted = np.vstack(
        [dists[i, sort_idx[i]] for i in range(sort_idx.shape[0])]
    )

    transformer = PyNNDescentTransformer(
        n_neighbors=N_NEIGHBORS,
        search_queue_size=QUEUE_SIZE,
        random_state=42
    ).fit(train)
    Xt = transformer.transform(test).sorted_indices()

    assert np.all(Xt.indices == indices_sorted.flat)
    assert np.allclose(Xt.data, dists_sorted.flat)

def test_random_state_none():
    knn_indices, _ = NNDescent(
        nn_data, "euclidean", {}, 10, random_state=None
    )._neighbor_graph

    tree = KDTree(nn_data)
    true_indices = tree.query(nn_data, 10, return_distance=False)

    num_correct = 0.0
    for i in range(nn_data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    percent_correct = num_correct / (spatial_data.shape[0] * 10)
    assert_greater_equal(
        percent_correct,
        0.99,
        "NN-descent did not get 99% " "accuracy on nearest neighbors",
    )

def test_deterministic():
    seed = np.random.RandomState(42)

    x1 = seed.normal(0, 100, (1000, 50))
    x2 = seed.normal(0, 100, (1000, 50))

    index1 = NNDescent(x1, random_state=np.random.RandomState(42))
    neighbors1, distances1 = index1.query(x2)

    index2 = NNDescent(x1, random_state=np.random.RandomState(42))
    neighbors2, distances2 = index2.query(x2)

    np.testing.assert_equal(neighbors1, neighbors2)
    np.testing.assert_equal(distances1, distances2)

# This tests a recursion error on cosine metric reported at:
# https://github.com/lmcinnes/umap/issues/99
# data used is a cut-down version of that provided by @scharron
# It contains lots of all-zero vectors and some other duplicates
def test_rp_trees_should_not_stack_overflow_with_duplicate_data():
    this_dir = os.path.dirname(os.path.abspath(__file__))
    data_path = os.path.join(this_dir, "test_data/cosine_hang.npy")
    data = np.load(data_path)

    n_neighbors = 10
    knn_indices, _ = NNDescent(
        data, "cosine", {}, n_neighbors, random_state=np.random, n_trees=20
    )._neighbor_graph

    for i in range(data.shape[0]):
        assert_equal(
            len(knn_indices[i]),
            len(np.unique(knn_indices[i])),
            "Duplicate indices in knn graph",
        )

def test_deduplicated_data_behaves_normally():
    this_dir = os.path.dirname(os.path.abspath(__file__))
    data_path = os.path.join(this_dir, "test_data/cosine_hang.npy")
    data = np.unique(np.load(data_path), axis=0)
    data = data[~np.all(data == 0, axis=1)]
    data = data[:1000]

    n_neighbors = 10
    knn_indices, _ = NNDescent(
        data, "cosine", {}, n_neighbors, random_state=np.random, n_trees=20
    )._neighbor_graph

    for i in range(data.shape[0]):
        assert_equal(
            len(knn_indices[i]),
            len(np.unique(knn_indices[i])),
            "Duplicate indices in knn graph",
        )

    angular_data = normalize(data, norm="l2")
    tree = KDTree(angular_data)
    true_indices = tree.query(angular_data, n_neighbors, return_distance=False)

    num_correct = 0
    for i in range(data.shape[0]):
        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))

    proportion_correct = num_correct / (data.shape[0] * n_neighbors)
    assert_greater_equal(
        proportion_correct,
        0.95,
        "NN-descent did not get 95%" " accuracy on nearest neighbors",
    )

def test_output_when_verbose_is_true():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = NNDescent(
            data=spatial_data,
            metric="euclidean",
            metric_kwds={},
            n_neighbors=4,
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=True,
        )
    output = out.getvalue()
    assert_true(re.match("^.*5 trees", output, re.DOTALL))
    assert_true(re.match("^.*2 iterations", output, re.DOTALL))

def test_no_output_when_verbose_is_false():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = NNDescent(
            data=spatial_data,
            metric="euclidean",
            metric_kwds={},
            n_neighbors=4,
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=False,
        )
    output = out.getvalue().strip()
    assert_equal(len(output), 0)

# same as the previous two test, but this time using the PyNNDescentTransformer
# interface
def test_transformer_output_when_verbose_is_true():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = PyNNDescentTransformer(
            n_neighbors=4,
            metric="euclidean",
            metric_kwds={},
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=True,
        ).fit_transform(spatial_data)
    output = out.getvalue()
    assert_true(re.match("^.*5 trees", output, re.DOTALL))
    assert_true(re.match("^.*2 iterations", output, re.DOTALL))

def test_transformer_output_when_verbose_is_false():
    out = io.StringIO()
    with redirect_stdout(out):
        _ = PyNNDescentTransformer(
            n_neighbors=4,
            metric="euclidean",
            metric_kwds={},
            random_state=np.random,
            n_trees=5,
            n_iters=2,
            verbose=False,
        ).fit_transform(spatial_data)
    output = out.getvalue().strip()
    assert_equal(len(output), 0)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright (c) 2015 Angel Terrones (<angelterrones@gmail.com>)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from myhdl import modbv

class Consts:
    # control signals
    Y          = True
    N          = False
    # PC Select Signal
    PC_4       = modbv(0)[2:]
    PC_BRJMP   = modbv(1)[2:]
    PC_JALR    = modbv(2)[2:]
    PC_EXC     = modbv(3)[2:]
    # Branch type
    BR_N       = modbv(0)[4:]
    BR_NE      = modbv(1)[4:]
    BR_EQ      = modbv(2)[4:]
    BR_GE      = modbv(3)[4:]
    BR_GEU     = modbv(4)[4:]
    BR_LT      = modbv(5)[4:]
    BR_LTU     = modbv(6)[4:]
    BR_J       = modbv(7)[4:]
    BR_JR      = modbv(8)[4:]
    # RS1 Operand Select Signal
    OP1_X      = modbv(0)[2:]
    OP1_RS1    = modbv(0)[2:]
    OP1_PC     = modbv(1)[2:]
    OP1_IMZ    = modbv(2)[2:]
    # RS2 Operand Select Signal
    OP2_X      = modbv(0)[3:]
    OP2_RS2    = modbv(1)[3:]
    OP2_ITYPE  = modbv(2)[3:]
    OP2_STYPE  = modbv(3)[3:]
    OP2_SBTYPE = modbv(4)[3:]
    OP2_UTYPE  = modbv(5)[3:]
    OP2_JUTYPE = modbv(6)[3:]
    # PRIV
    MTVEC      = 0x100
    START_ADDR = MTVEC + 0x100

# Local Variables:
# flycheck-flake8-maximum-line-length: 120
# flycheck-flake8rc: ".flake8rc"
# End:
<commit_msg>Fix constants, and add WB select values.<commit_after>#!/usr/bin/env python
# Copyright (c) 2015 Angel Terrones (<angelterrones@gmail.com>)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from myhdl import modbv

class Consts:
    # control signals
    Y          = True
    N          = False
    # PC Select Signal
    PC_4       = modbv(0)[2:]
    PC_BRJMP   = modbv(1)[2:]
    PC_JALR    = modbv(2)[2:]
    PC_EXC     = modbv(3)[2:]
    # Branch type
    BR_N       = modbv(0)[4:]
    BR_NE      = modbv(1)[4:]
    BR_EQ      = modbv(2)[4:]
    BR_GE      = modbv(3)[4:]
    BR_GEU     = modbv(4)[4:]
    BR_LT      = modbv(5)[4:]
    BR_LTU     = modbv(6)[4:]
    BR_J       = modbv(7)[4:]
    BR_JR      = modbv(8)[4:]
    # RS1 Operand Select Signal
    OP1_X      = modbv(0)[2:]
    OP1_RS1    = modbv(0)[2:]
    OP1_PC     = modbv(1)[2:]
    OP1_IMZ    = modbv(2)[2:]
    # RS2 Operand Select Signal
    OP2_X      = modbv(0)[3:]
    OP2_RS2    = modbv(0)[3:]
    OP2_ITYPE  = modbv(1)[3:]
    OP2_STYPE  = modbv(2)[3:]
    OP2_SBTYPE = modbv(3)[3:]
    OP2_UTYPE  = modbv(4)[3:]
    OP2_UJTYPE = modbv(5)[3:]
    # Forwarding
    FWD_X      = modbv(0)[2:]
    FWD_N      = modbv(0)[2:]
    FWD_EX     = modbv(1)[2:]
    FWD_MEM    = modbv(2)[2:]
    FWD_WB     = modbv(3)[2:]
    # WB signals
    WB_X       = 0
    WB_ALU     = 0
    WB_MEM     = 1
    WB_PC4     = 2
    WB_CSR     = 3
    # PRIV
    MTVEC      = 0x100
    START_ADDR = MTVEC + 0x100

# Local Variables:
# flycheck-flake8-maximum-line-length: 120
# flycheck-flake8rc: ".flake8rc"
# End:
<|endoftext|>