<commit_before>import json
import os

from cStringIO import StringIO

from . import ansible

from ..config import config as teuth_config
from ..misc import get_scratch_devices

class CephAnsible(ansible.Ansible):
    name = 'ceph_ansible'

    _default_playbook = [
        dict(
            hosts='mons',
            become=True,
            roles=['ceph-mon'],
        ),
        dict(
            hosts='osds',
            become=True,
            roles=['ceph-osd'],
        ),
        dict(
            hosts='mdss',
            become=True,
            roles=['ceph-mds'],
        ),
        dict(
            hosts='rgws',
            become=True,
            roles=['ceph-rgw'],
        ),
        dict(
            hosts='restapis',
            become=True,
            roles=['ceph-restapi'],
        ),
    ]

    __doc__ = """
    A subclass of Ansible that defaults to:

    - ansible:
        repo: {git_base}ceph-ansible.git
        playbook: {playbook}

    It always uses a dynamic inventory.

    It will optionally do the following automatically based on ``vars`` that
    are passed in:
        * Set ``devices`` for each host if ``osd_auto_discovery`` is not True
        * Set ``monitor_interface`` for each host if ``monitor_interface`` is
          unset
        * Set ``public_network`` for each host if ``public_network`` is unset
    """.format(
        git_base=teuth_config.ceph_git_base_url,
        playbook=_default_playbook,
    )

    def __init__(self, ctx, config):
        config = config or dict()
        if 'playbook' not in config:
            config['playbook'] = self._default_playbook
        if 'repo' not in config:
            config['repo'] = os.path.join(teuth_config.ceph_git_base_url,
                                          'ceph-ansible.git')
        super(CephAnsible, self).__init__(ctx, config)

    def get_inventory(self):
        """
        Stub this method so we always generate the hosts file
        """
        pass

    def generate_hosts_file(self):
        groups_to_roles = dict(
            mons='mon',
            mdss='mds',
            osds='osd',
        )
        hosts_dict = dict()
        for group in sorted(groups_to_roles.keys()):
            role_prefix = groups_to_roles[group]
            want = lambda role: role.startswith(role_prefix)
            for (remote, roles) in self.cluster.only(want).remotes.iteritems():
                hostname = remote.hostname
                host_vars = self.get_host_vars(remote)
                if group not in hosts_dict:
                    hosts_dict[group] = {hostname: host_vars}
                elif hostname not in hosts_dict[group]:
                    hosts_dict[group][hostname] = host_vars

        hosts_stringio = StringIO()
        for group in sorted(hosts_dict.keys()):
            hosts_stringio.write('[%s]\n' % group)
            for hostname in sorted(hosts_dict[group].keys()):
                vars = hosts_dict[group][hostname]
                if vars:
                    vars_list = []
                    for key in sorted(vars.keys()):
                        vars_list.append(
                            "%s='%s'" % (key, json.dumps(vars[key]).strip('"'))
                        )
                    host_line = "{hostname} {vars}".format(
                        hostname=hostname,
                        vars=' '.join(vars_list),
                    )
                else:
                    host_line = hostname
                hosts_stringio.write('%s\n' % host_line)
            hosts_stringio.write('\n')
        hosts_stringio.seek(0)
        self.inventory = self._write_hosts_file(hosts_stringio.read().strip())
        self.generated_inventory = True

    def get_host_vars(self, remote):
        extra_vars = self.config.get('vars', dict())
        host_vars = dict()
        if not extra_vars.get('osd_auto_discovery', False):
            host_vars['devices'] = get_scratch_devices(remote)
        if 'monitor_interface' not in extra_vars:
            host_vars['monitor_interface'] = remote.interface
        if 'public_network' not in extra_vars:
            host_vars['public_network'] = remote.cidr
        return host_vars

task = CephAnsible
<commit_msg>use Task subclass instead of Ansible<commit_after>import json
import os

from cStringIO import StringIO

from . import ansible
from . import Task
from ..config import config as teuth_config
from ..misc import get_scratch_devices

class CephAnsible(Task):
    name = 'ceph_ansible'

    _default_playbook = [
        dict(
            hosts='mons',
            become=True,
            roles=['ceph-mon'],
        ),
        dict(
            hosts='osds',
            become=True,
            roles=['ceph-osd'],
        ),
        dict(
            hosts='mdss',
            become=True,
            roles=['ceph-mds'],
        ),
        dict(
            hosts='rgws',
            become=True,
            roles=['ceph-rgw'],
        ),
        dict(
            hosts='restapis',
            become=True,
            roles=['ceph-restapi'],
        ),
    ]

    __doc__ = """
    A task to setup ceph cluster using ceph-ansible

    - ceph-ansible:
        repo: {git_base}ceph-ansible.git
        branch: mybranch # defaults to master
        ansible-version: 2.2 # defaults to 2.1
        vars:
          ceph_dev: True ( default)
          ceph_conf_overrides:
             global:
                mon pg warn min per osd: 2

    It always uses a dynamic inventory.

    It will optionally do the following automatically based on ``vars`` that
    are passed in:
        * Set ``devices`` for each host if ``osd_auto_discovery`` is not True
        * Set ``monitor_interface`` for each host if ``monitor_interface`` is
          unset
        * Set ``public_network`` for each host if ``public_network`` is unset
    """.format(
        git_base=teuth_config.ceph_git_base_url,
        playbook=_default_playbook,
    )

    def __init__(self, ctx, config):
        super(CephAnsible, self).__init__(ctx, config)
        config = self.config or dict()
        if 'playbook' not in config:
            self.playbook = self._default_playbook
        else:
            self.playbook = self.config['playbook']
        if 'repo' not in config:
            self.config['repo'] = os.path.join(teuth_config.ceph_git_base_url,
                                               'ceph-ansible.git')
        # default vars to dev builds
        if 'vars' not in config:
            vars = dict()
            config['vars'] = vars
        vars = config['vars']
        if 'ceph_dev' not in vars:
            vars['ceph_dev'] = True
        if 'ceph_dev_key' not in vars:
            vars['ceph_dev_key'] = 'https://download.ceph.com/keys/autobuild.asc'
        if 'ceph_dev_branch' not in vars:
            vars['ceph_dev_branch'] = ctx.config.get('branch', 'master')

    def get_inventory(self):
        """
        Stub this method so we always generate the hosts file
        """
        pass

    def generate_hosts_file(self):
        groups_to_roles = dict(
            mons='mon',
            mdss='mds',
            osds='osd',
        )
        hosts_dict = dict()
        for group in sorted(groups_to_roles.keys()):
            role_prefix = groups_to_roles[group]
            want = lambda role: role.startswith(role_prefix)
            for (remote, roles) in self.cluster.only(want).remotes.iteritems():
                hostname = remote.hostname
                host_vars = self.get_host_vars(remote)
                if group not in hosts_dict:
                    hosts_dict[group] = {hostname: host_vars}
                elif hostname not in hosts_dict[group]:
                    hosts_dict[group][hostname] = host_vars

        hosts_stringio = StringIO()
        for group in sorted(hosts_dict.keys()):
            hosts_stringio.write('[%s]\n' % group)
            for hostname in sorted(hosts_dict[group].keys()):
                vars = hosts_dict[group][hostname]
                if vars:
                    vars_list = []
                    for key in sorted(vars.keys()):
                        vars_list.append(
                            "%s='%s'" % (key, json.dumps(vars[key]).strip('"'))
                        )
                    host_line = "{hostname} {vars}".format(
                        hostname=hostname,
                        vars=' '.join(vars_list),
                    )
                else:
                    host_line = hostname
                hosts_stringio.write('%s\n' % host_line)
            hosts_stringio.write('\n')
        hosts_stringio.seek(0)
        self.inventory = self._write_hosts_file(hosts_stringio.read().strip())
        self.generated_inventory = True

    def get_host_vars(self, remote):
        extra_vars = self.config.get('vars', dict())
        host_vars = dict()
        if not extra_vars.get('osd_auto_discovery', False):
            host_vars['devices'] = get_scratch_devices(remote)
        if 'monitor_interface' not in extra_vars:
            host_vars['monitor_interface'] = remote.interface
        if 'public_network' not in extra_vars:
            host_vars['public_network'] = remote.cidr
        return host_vars

task = CephAnsible
<|endoftext|>


<commit_before>import os
import subprocess
from .simulator import Simulator

class Modelsim(Simulator):

    def __init__(self, system):
        super(Modelsim, self).__init__(system)
        self.sim_root = os.path.join(self.build_root, 'sim-modelsim')
        if system.config.has_option('modelsim', 'vlog_options'):
            self.vlog_options = system.config.get('modelsim','vlog_options').split()
        else:
            self.vlog_options = []

        if system.config.has_option('modelsim', 'vsim_options'):
            self.vsim_options = system.config.get('modelsim','vsim_options').split()
        else:
            self.vsim_options = []

    def configure(self):
        super(Modelsim, self).configure()
        self._write_config_files()

    def _write_config_files(self):
        self.cfg_file = self.system.name+'.scr'

        f = open(os.path.join(self.sim_root,self.cfg_file),'w')

        for include_dir in self.include_dirs:
            f.write("+incdir+" + include_dir + '\n')
        for src_file in self.src_files:
            f.write(src_file + '\n')

        f.close()

    def build(self):
        super(Modelsim, self).build()
        self.model_tech = os.getenv('MODEL_TECH')
        if not self.model_tech:
            print("Environment variable MODEL_TECH was not found. It should be set to <modelsim install path>/bin")
            exit(1)

        #FIXME: Handle failures. Save stdout/stderr. Build vmem file from elf file argument
        try:
            subprocess.check_call([self.model_tech+'/vlib', 'work'],
                                  cwd = self.sim_root)

        except OSError:
            print("Error: Command vlib not found. Make sure it is in $PATH")
            exit(1)
        except CalledProcessError:
            print("Error: Failed to create library work")
            exit(1)

        try:
            logfile = os.path.join(self.sim_root, 'vlog.log')
            subprocess.check_call([self.model_tech+'/vlog', '-f', self.cfg_file, '-quiet', '-l', logfile] +
                                  self.vlog_options,
                            cwd = self.sim_root)
        except OSError:
            print("Error: Command vlog not found. Make sure it is in $PATH")
            exit(1)
        except subprocess.CalledProcessError:
            print("Error: Failed to compile simulation model. Compile log is available in " + logfile)
            exit(1)

        for vpi_module in self.vpi_modules:
            objs = []
            for src_file in vpi_module['src_files']:
                try:
                    subprocess.check_call(['gcc', '-c', '-fPIC', '-g','-m32','-DMODELSIM_VPI'] +
                                          ['-I'+self.model_tech+'/../include'] +
                                          ['-I'+s for s in vpi_module['include_dirs']] +
                                          [src_file],
                                          cwd = self.sim_root,
                                          stdin=subprocess.PIPE)
                except OSError:
                    print("Error: Command gcc not found. Make sure it is in $PATH")
                    exit(1)
                except subprocess.CalledProcessError:
                    print("Error: Compilation of "+src_file + "failed")
                    exit(1)

	        object_files = [os.path.splitext(os.path.basename(s))[0]+'.o' for s in vpi_module['src_files']]
            try:
                subprocess.check_call(['ld', '-shared','-E','-melf_i386','-o',vpi_module['name']] +
                                      object_files,
                                      cwd = self.sim_root,
                                      stdin=subprocess.PIPE)
            except OSError:
                print("Error: Command ld not found. Make sure it is in $PATH")
                exit(1)
            except subprocess.CalledProcessError:
                print("Error: Linking of "+vpi_module['name'] + " failed")
                exit(1)

    def run(self, args):
        super(Modelsim, self).run(args)

        #FIXME: Handle failures. Save stdout/stderr
        vpi_options = []
        for vpi_module in self.vpi_modules:
            vpi_options += ['-pli', vpi_module['name']]
        try:
            logfile = os.path.join(self.sim_root, 'vsim.log')
            subprocess.check_call([self.model_tech+'/vsim', '-c', '-do', 'run -all'] +
                                  ['-l', logfile] +
                                  self.vsim_options +
                                  vpi_options +
                                  ['work.orpsoc_tb'] +
                                  ['+'+s for s in self.plusargs],
                                  cwd = self.sim_root,
                                  stdin=subprocess.PIPE)
        except OSError:
            print("Error: Command vsim not found. Make sure it is in $PATH")
            exit(1)
        except subprocess.CalledProcessError:
            print("Error: Simulation failed. Simulation log is available in " + logfile)
            exit(1)
<commit_msg>Whitespace fix in modelsim.py<commit_after>import os
import subprocess
from .simulator import Simulator

class Modelsim(Simulator):

    def __init__(self, system):
        super(Modelsim, self).__init__(system)
        self.sim_root = os.path.join(self.build_root, 'sim-modelsim')
        if system.config.has_option('modelsim', 'vlog_options'):
            self.vlog_options = system.config.get('modelsim','vlog_options').split()
        else:
            self.vlog_options = []

        if system.config.has_option('modelsim', 'vsim_options'):
            self.vsim_options = system.config.get('modelsim','vsim_options').split()
        else:
            self.vsim_options = []

    def configure(self):
        super(Modelsim, self).configure()
        self._write_config_files()

    def _write_config_files(self):
        self.cfg_file = self.system.name+'.scr'

        f = open(os.path.join(self.sim_root,self.cfg_file),'w')

        for include_dir in self.include_dirs:
            f.write("+incdir+" + include_dir + '\n')
        for src_file in self.src_files:
            f.write(src_file + '\n')

        f.close()

    def build(self):
        super(Modelsim, self).build()
        self.model_tech = os.getenv('MODEL_TECH')
        if not self.model_tech:
            print("Environment variable MODEL_TECH was not found. It should be set to <modelsim install path>/bin")
            exit(1)

        #FIXME: Handle failures. Save stdout/stderr. Build vmem file from elf file argument
        try:
            subprocess.check_call([self.model_tech+'/vlib', 'work'],
                                  cwd = self.sim_root)

        except OSError:
            print("Error: Command vlib not found. Make sure it is in $PATH")
            exit(1)
        except CalledProcessError:
            print("Error: Failed to create library work")
            exit(1)

        try:
            logfile = os.path.join(self.sim_root, 'vlog.log')
            subprocess.check_call([self.model_tech+'/vlog', '-f', self.cfg_file, '-quiet', '-l', logfile] +
                                  self.vlog_options,
                            cwd = self.sim_root)
        except OSError:
            print("Error: Command vlog not found. Make sure it is in $PATH")
            exit(1)
        except subprocess.CalledProcessError:
            print("Error: Failed to compile simulation model. Compile log is available in " + logfile)
            exit(1)

        for vpi_module in self.vpi_modules:
            objs = []
            for src_file in vpi_module['src_files']:
                try:
                    subprocess.check_call(['gcc', '-c', '-fPIC', '-g','-m32','-DMODELSIM_VPI'] +
                                          ['-I'+self.model_tech+'/../include'] +
                                          ['-I'+s for s in vpi_module['include_dirs']] +
                                          [src_file],
                                          cwd = self.sim_root,
                                          stdin=subprocess.PIPE)
                except OSError:
                    print("Error: Command gcc not found. Make sure it is in $PATH")
                    exit(1)
                except subprocess.CalledProcessError:
                    print("Error: Compilation of "+src_file + "failed")
                    exit(1)

                object_files = [os.path.splitext(os.path.basename(s))[0]+'.o' for s in vpi_module['src_files']]
            try:
                subprocess.check_call(['ld', '-shared','-E','-melf_i386','-o',vpi_module['name']] +
                                      object_files,
                                      cwd = self.sim_root,
                                      stdin=subprocess.PIPE)
            except OSError:
                print("Error: Command ld not found. Make sure it is in $PATH")
                exit(1)
            except subprocess.CalledProcessError:
                print("Error: Linking of "+vpi_module['name'] + " failed")
                exit(1)

    def run(self, args):
        super(Modelsim, self).run(args)

        #FIXME: Handle failures. Save stdout/stderr
        vpi_options = []
        for vpi_module in self.vpi_modules:
            vpi_options += ['-pli', vpi_module['name']]
        try:
            logfile = os.path.join(self.sim_root, 'vsim.log')
            subprocess.check_call([self.model_tech+'/vsim', '-c', '-do', 'run -all'] +
                                  ['-l', logfile] +
                                  self.vsim_options +
                                  vpi_options +
                                  ['work.orpsoc_tb'] +
                                  ['+'+s for s in self.plusargs],
                                  cwd = self.sim_root,
                                  stdin=subprocess.PIPE)
        except OSError:
            print("Error: Command vsim not found. Make sure it is in $PATH")
            exit(1)
        except subprocess.CalledProcessError:
            print("Error: Simulation failed. Simulation log is available in " + logfile)
            exit(1)
<|endoftext|>


<commit_before>from cStringIO import StringIO
import random
import time
import re
import gevent
import json
from ..orchestra import run

class Thrasher(gevent.Greenlet):
    def __init__(self, manager, config, logger=None):
        self.ceph_manager = manager
        self.ceph_manager.wait_till_clean()
        osd_status = self.ceph_manager.get_osd_status()
        self.in_osds = osd_status['in']
        self.live_osds = osd_status['live']
        self.out_osds = osd_status['out']
        self.dead_osds = osd_status['dead']
        self.stopping = False
        self.logger = logger
        self.config = config
        if self.logger is not None:
            self.log = lambda x: self.logger.info(x)
        else:
            def tmp(x):
                print x
            self.log = tmp
        if self.config is None:
            self.config = dict()
        gevent.Greenlet.__init__(self, self.do_thrash)
        self.start()

    def kill_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.live_osds)
        self.log("Killing osd %s, live_osds are %s"%(str(osd),str(self.live_osds)))
        self.live_osds.remove(osd)
        self.dead_osds.append(osd)
        if osd in self.in_osds:
            self.in_osds.remove(osd)
            self.out_osds.append(osd)
        self.ceph_manager.kill_osd(osd)

    def revive_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.dead_osds)
        self.log("Reviving osd %s"%(str(osd),))
        self.live_osds.append(osd)
        self.dead_osds.remove(osd)
        if osd in self.out_osds:
            self.in_osds.append(osd)
            self.out_osds.remove(osd)
        self.ceph_manager.revive_osd(osd)

    def out_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.in_osds)
        self.log("Removing osd %s, in_osds are: %s"%(str(osd),str(self.in_osds)))
        self.ceph_manager.mark_out_osd(osd)
        self.in_osds.remove(osd)
        self.out_osds.append(osd)

    def in_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.out_osds)
        if osd in self.dead_osds:
            return self.revive_osd(osd)
        self.log("Adding osd %s"%(str(osd),))
        self.out_osds.remove(osd)
        self.in_osds.append(osd)
        self.ceph_manager.mark_in_osd(osd)

    def all_up(self):
        while len(self.dead_osds) > 0:
            self.revive_osd()
        while len(self.out_osds) > 0:
            self.in_osd()

    def do_join(self):
        self.stopping = True
        self.get()

    def choose_action(self):
        chance_down = self.config.get("chance_down", 0)
        if isinstance(chance_down, int):
            chance_down = float(chance_down) / 100
        minin = self.config.get("min_in", 2)
        minout = self.config.get("min_out", 0)
        minlive = self.config.get("min_live", 2)
        mindead = self.config.get("min_dead", 2)

        actions = []
        if len(self.in_osds) > minin:
            actions.append((self.out_osd, 1.0,))
            if len(self.live_osds) > minlive and chance_down > 0:
                actions.append((self.kill_osd, chance_down))
        if len(self.out_osds) > minout:
            actions.append((self.in_osd, 1.0,))
            if len(self.dead_osds) > mindead:
                actions.append((self.revive_osd, 1.0))

        total = sum([y for (x,y) in actions])
        rev_cum = reduce(lambda l,(y1,y2): l+[(y1, (l[-1][1]+y2)/total)], actions, [(0, 0)])[1:]
        rev_cum.reverse()
        final_rev_cum = [(x, y-rev_cum[-1][1]) for (x,y) in rev_cum]
        val = random.uniform(0, 1)
        for (action, prob) in final_rev_cum:
            if (prob < val):
                return action
        return None

    def do_thrash(self):
        cleanint = self.config.get("clean_interval", 60)
        delay = self.config.get("op_delay", 5)
        self.log("starting do_thrash")
        while not self.stopping:
            self.log(" ".join([str(x) for x in ["in_osds: ", self.in_osds, " out_osds: ", self.out_osds,
                                                "dead_osds: ", self.dead_osds, "live_osds: ",
                                                self.live_osds]]))
            if random.uniform(0,1) < (float(delay) / cleanint):
                self.ceph_manager.wait_till_clean()
            self.choose_action()()
            time.sleep(delay)
        self.all_up()

class CephManager:
    def __init__(self, controller, ctx=None, logger=None):
        self.ctx = ctx
        self.controller = controller
        if (logger):
            self.log = lambda x: logger.info(x)
        else:
            def tmp(x):
                print x
            self.log = tmp

    def raw_cluster_cmd(self, *args):
        ceph_args = [
                'LD_LIBRARY_PRELOAD=/tmp/cephtest/binary/usr/local/lib',
                '/tmp/cephtest/enable-coredump',
                '/tmp/cephtest/binary/usr/local/bin/ceph-coverage',
                '/tmp/cephtest/archive/coverage',
                '/tmp/cephtest/binary/usr/local/bin/ceph',
                '-k', '/tmp/cephtest/ceph.keyring',
                '-c', '/tmp/cephtest/ceph.conf',
                '--concise',
                ]
        ceph_args.extend(args)
        proc = self.controller.run(
            args=ceph_args,
            stdout=StringIO(),
            )
        return proc.stdout.getvalue()

    def raw_cluster_status(self):
        return self.raw_cluster_cmd('-s')

    def raw_osd_status(self):
        return self.raw_cluster_cmd('osd', 'dump')

    def get_osd_status(self):
        osd_lines = filter(
            lambda x: x.startswith('osd.') and (("up" in x) or ("down" in x)),
            self.raw_osd_status().split('\n'))
        self.log(osd_lines)
        in_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " in " in x,
                osd_lines)]
        out_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " out " in x,
                osd_lines)]
        up_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " up " in x,
                osd_lines)]
        down_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " down " in x,
                osd_lines)]
        dead_osds = [int(x.id_) for x in
                     filter(lambda x: not x.running(), self.ctx.daemons.iter_daemons_of_role('osd'))]
        live_osds = [int(x.id_) for x in
                     filter(lambda x: x.running(), self.ctx.daemons.iter_daemons_of_role('osd'))]
        return { 'in' : in_osds, 'out' : out_osds, 'up' : up_osds,
                 'down' : down_osds, 'dead' : dead_osds, 'live' : live_osds, 'raw' : osd_lines }

    def get_num_pgs(self):
        status = self.raw_cluster_status()
        return int(re.search(
                "\d* pgs:",
                status).group(0).split()[0])

    def get_pg_stats(self):
        out = self.raw_cluster_cmd('--', 'pg','dump','--format=json')
        j = json.loads('\n'.join(out.split('\n')[1:]))
        return j['pg_stats']

    def get_osd_dump(self):
        out = self.raw_cluster_cmd('--', 'osd','dump','--format=json')
        j = json.loads('\n'.join(out.split('\n')[1:]))
        return j['osds']

    def get_num_unfound_objects(self):
        status = self.raw_cluster_status()
        self.log(status)
        match = re.search(
            "\d+/\d+ unfound",
            status)
        if match == None:
            return 0
        else:
            return int(match.group(0).split('/')[0])

    def get_num_active_clean(self):
        status = self.raw_cluster_status()
        self.log(status)
        match = re.search(
            "\d* active.clean",
            status)
        if match == None:
            return 0
        else:
            return int(match.group(0).split()[0])

    def get_num_active(self):
        pgs = self.get_pg_stats()
        num = 0
        for pg in pgs:
            if pg['state'].startswith('active'):
                num += 1
        return num

    def is_clean(self):
        return self.get_num_active_clean() == self.get_num_pgs()

    def wait_till_clean(self, timeout=None):
        self.log("waiting till clean")
        start = time.time()
        while not self.is_clean():
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'failed to become clean before timeout expired'
            time.sleep(3)
        self.log("clean!")

    def osd_is_up(self, osd):
        osds = self.get_osd_dump()
        return osds[osd]['up'] > 0

    def wait_till_osd_is_up(self, osd, timeout=None):
        self.log('waiting for osd.%d to be up' % osd);
        start = time.time()
        while not self.osd_is_up(osd):
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'osd.%d failed to come up before timeout expired' % osd
            time.sleep(3)
        self.log('osd.%d is up' % osd)

    def is_active(self):
        return self.get_num_active() == self.get_num_pgs()

    def wait_till_active(self, timeout=None):
        self.log("waiting till active")
        start = time.time()
        while not self.is_active():
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'failed to become active before timeout expired'
            time.sleep(3)
        self.log("active!")

    def mark_out_osd(self, osd):
        self.raw_cluster_cmd('osd', 'out', str(osd))

    def kill_osd(self, osd):
        self.ctx.daemons.get_daemon('osd', osd).stop()

    def revive_osd(self, osd):
        self.ctx.daemons.get_daemon('osd', osd).restart()

    def mark_down_osd(self, osd):
        self.raw_cluster_cmd('osd', 'down', str(osd))

    def mark_in_osd(self, osd):
        self.raw_cluster_cmd('osd', 'in', str(osd))
<commit_msg>ceph_manager: count active+clean+<somjething else> as active+clean<commit_after>from cStringIO import StringIO
import random
import time
import re
import gevent
import json
from ..orchestra import run

class Thrasher(gevent.Greenlet):
    def __init__(self, manager, config, logger=None):
        self.ceph_manager = manager
        self.ceph_manager.wait_till_clean()
        osd_status = self.ceph_manager.get_osd_status()
        self.in_osds = osd_status['in']
        self.live_osds = osd_status['live']
        self.out_osds = osd_status['out']
        self.dead_osds = osd_status['dead']
        self.stopping = False
        self.logger = logger
        self.config = config
        if self.logger is not None:
            self.log = lambda x: self.logger.info(x)
        else:
            def tmp(x):
                print x
            self.log = tmp
        if self.config is None:
            self.config = dict()
        gevent.Greenlet.__init__(self, self.do_thrash)
        self.start()

    def kill_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.live_osds)
        self.log("Killing osd %s, live_osds are %s"%(str(osd),str(self.live_osds)))
        self.live_osds.remove(osd)
        self.dead_osds.append(osd)
        if osd in self.in_osds:
            self.in_osds.remove(osd)
            self.out_osds.append(osd)
        self.ceph_manager.kill_osd(osd)

    def revive_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.dead_osds)
        self.log("Reviving osd %s"%(str(osd),))
        self.live_osds.append(osd)
        self.dead_osds.remove(osd)
        if osd in self.out_osds:
            self.in_osds.append(osd)
            self.out_osds.remove(osd)
        self.ceph_manager.revive_osd(osd)

    def out_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.in_osds)
        self.log("Removing osd %s, in_osds are: %s"%(str(osd),str(self.in_osds)))
        self.ceph_manager.mark_out_osd(osd)
        self.in_osds.remove(osd)
        self.out_osds.append(osd)

    def in_osd(self, osd=None):
        if osd is None:
            osd = random.choice(self.out_osds)
        if osd in self.dead_osds:
            return self.revive_osd(osd)
        self.log("Adding osd %s"%(str(osd),))
        self.out_osds.remove(osd)
        self.in_osds.append(osd)
        self.ceph_manager.mark_in_osd(osd)

    def all_up(self):
        while len(self.dead_osds) > 0:
            self.revive_osd()
        while len(self.out_osds) > 0:
            self.in_osd()

    def do_join(self):
        self.stopping = True
        self.get()

    def choose_action(self):
        chance_down = self.config.get("chance_down", 0)
        if isinstance(chance_down, int):
            chance_down = float(chance_down) / 100
        minin = self.config.get("min_in", 2)
        minout = self.config.get("min_out", 0)
        minlive = self.config.get("min_live", 2)
        mindead = self.config.get("min_dead", 2)

        actions = []
        if len(self.in_osds) > minin:
            actions.append((self.out_osd, 1.0,))
            if len(self.live_osds) > minlive and chance_down > 0:
                actions.append((self.kill_osd, chance_down))
        if len(self.out_osds) > minout:
            actions.append((self.in_osd, 1.0,))
            if len(self.dead_osds) > mindead:
                actions.append((self.revive_osd, 1.0))

        total = sum([y for (x,y) in actions])
        rev_cum = reduce(lambda l,(y1,y2): l+[(y1, (l[-1][1]+y2)/total)], actions, [(0, 0)])[1:]
        rev_cum.reverse()
        final_rev_cum = [(x, y-rev_cum[-1][1]) for (x,y) in rev_cum]
        val = random.uniform(0, 1)
        for (action, prob) in final_rev_cum:
            if (prob < val):
                return action
        return None

    def do_thrash(self):
        cleanint = self.config.get("clean_interval", 60)
        delay = self.config.get("op_delay", 5)
        self.log("starting do_thrash")
        while not self.stopping:
            self.log(" ".join([str(x) for x in ["in_osds: ", self.in_osds, " out_osds: ", self.out_osds,
                                                "dead_osds: ", self.dead_osds, "live_osds: ",
                                                self.live_osds]]))
            if random.uniform(0,1) < (float(delay) / cleanint):
                self.ceph_manager.wait_till_clean()
            self.choose_action()()
            time.sleep(delay)
        self.all_up()

class CephManager:
    def __init__(self, controller, ctx=None, logger=None):
        self.ctx = ctx
        self.controller = controller
        if (logger):
            self.log = lambda x: logger.info(x)
        else:
            def tmp(x):
                print x
            self.log = tmp

    def raw_cluster_cmd(self, *args):
        ceph_args = [
                'LD_LIBRARY_PRELOAD=/tmp/cephtest/binary/usr/local/lib',
                '/tmp/cephtest/enable-coredump',
                '/tmp/cephtest/binary/usr/local/bin/ceph-coverage',
                '/tmp/cephtest/archive/coverage',
                '/tmp/cephtest/binary/usr/local/bin/ceph',
                '-k', '/tmp/cephtest/ceph.keyring',
                '-c', '/tmp/cephtest/ceph.conf',
                '--concise',
                ]
        ceph_args.extend(args)
        proc = self.controller.run(
            args=ceph_args,
            stdout=StringIO(),
            )
        return proc.stdout.getvalue()

    def raw_cluster_status(self):
        return self.raw_cluster_cmd('-s')

    def raw_osd_status(self):
        return self.raw_cluster_cmd('osd', 'dump')

    def get_osd_status(self):
        osd_lines = filter(
            lambda x: x.startswith('osd.') and (("up" in x) or ("down" in x)),
            self.raw_osd_status().split('\n'))
        self.log(osd_lines)
        in_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " in " in x,
                osd_lines)]
        out_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " out " in x,
                osd_lines)]
        up_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " up " in x,
                osd_lines)]
        down_osds = [int(i[4:].split()[0]) for i in filter(
                lambda x: " down " in x,
                osd_lines)]
        dead_osds = [int(x.id_) for x in
                     filter(lambda x: not x.running(), self.ctx.daemons.iter_daemons_of_role('osd'))]
        live_osds = [int(x.id_) for x in
                     filter(lambda x: x.running(), self.ctx.daemons.iter_daemons_of_role('osd'))]
        return { 'in' : in_osds, 'out' : out_osds, 'up' : up_osds,
                 'down' : down_osds, 'dead' : dead_osds, 'live' : live_osds, 'raw' : osd_lines }

    def get_num_pgs(self):
        status = self.raw_cluster_status()
        return int(re.search(
                "\d* pgs:",
                status).group(0).split()[0])

    def get_pg_stats(self):
        out = self.raw_cluster_cmd('--', 'pg','dump','--format=json')
        j = json.loads('\n'.join(out.split('\n')[1:]))
        return j['pg_stats']

    def get_osd_dump(self):
        out = self.raw_cluster_cmd('--', 'osd','dump','--format=json')
        j = json.loads('\n'.join(out.split('\n')[1:]))
        return j['osds']

    def get_num_unfound_objects(self):
        status = self.raw_cluster_status()
        self.log(status)
        match = re.search(
            "\d+/\d+ unfound",
            status)
        if match == None:
            return 0
        else:
            return int(match.group(0).split('/')[0])

    def get_num_active_clean(self):
        pgs = self.get_pg_stats()
        num = 0
        for pg in pgs:
            if pg['state'].startswith('active+clean'):
                num += 1
        return num

    def get_num_active(self):
        pgs = self.get_pg_stats()
        num = 0
        for pg in pgs:
            if pg['state'].startswith('active'):
                num += 1
        return num

    def is_clean(self):
        return self.get_num_active_clean() == self.get_num_pgs()

    def wait_till_clean(self, timeout=None):
        self.log("waiting till clean")
        start = time.time()
        while not self.is_clean():
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'failed to become clean before timeout expired'
            time.sleep(3)
        self.log("clean!")

    def osd_is_up(self, osd):
        osds = self.get_osd_dump()
        return osds[osd]['up'] > 0

    def wait_till_osd_is_up(self, osd, timeout=None):
        self.log('waiting for osd.%d to be up' % osd);
        start = time.time()
        while not self.osd_is_up(osd):
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'osd.%d failed to come up before timeout expired' % osd
            time.sleep(3)
        self.log('osd.%d is up' % osd)

    def is_active(self):
        return self.get_num_active() == self.get_num_pgs()

    def wait_till_active(self, timeout=None):
        self.log("waiting till active")
        start = time.time()
        while not self.is_active():
            if timeout is not None:
                assert time.time() - start < timeout, \
                    'failed to become active before timeout expired'
            time.sleep(3)
        self.log("active!")

    def mark_out_osd(self, osd):
        self.raw_cluster_cmd('osd', 'out', str(osd))

    def kill_osd(self, osd):
        self.ctx.daemons.get_daemon('osd', osd).stop()

    def revive_osd(self, osd):
        self.ctx.daemons.get_daemon('osd', osd).restart()

    def mark_down_osd(self, osd):
        self.raw_cluster_cmd('osd', 'down', str(osd))

    def mark_in_osd(self, osd):
        self.raw_cluster_cmd('osd', 'in', str(osd))
<|endoftext|>


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

"""Utilities for interacting with Identity Servers"""

import logging

from canonicaljson import json

from twisted.internet import defer

from synapse.api.errors import (
    CodeMessageException,
    Codes,
    HttpResponseException,
    SynapseError,
)

from ._base import BaseHandler

logger = logging.getLogger(__name__)

class IdentityHandler(BaseHandler):

    def __init__(self, hs):
        super(IdentityHandler, self).__init__(hs)

        self.http_client = hs.get_simple_http_client()
        self.federation_http_client = hs.get_http_client()

        self.trusted_id_servers = set(hs.config.trusted_third_party_id_servers)
        self.trust_any_id_server_just_for_testing_do_not_use = (
            hs.config.use_insecure_ssl_client_just_for_testing_do_not_use
        )

    def _should_trust_id_server(self, id_server):
        if id_server not in self.trusted_id_servers:
            if self.trust_any_id_server_just_for_testing_do_not_use:
                logger.warn(
                    "Trusting untrustworthy ID server %r even though it isn't"
                    " in the trusted id list for testing because"
                    " 'use_insecure_ssl_client_just_for_testing_do_not_use'"
                    " is set in the config",
                    id_server,
                )
            else:
                return False
        return True

    @defer.inlineCallbacks
    def threepid_from_creds(self, creds):
        if 'id_server' in creds:
            id_server = creds['id_server']
        elif 'idServer' in creds:
            id_server = creds['idServer']
        else:
            raise SynapseError(400, "No id_server in creds")

        if 'client_secret' in creds:
            client_secret = creds['client_secret']
        elif 'clientSecret' in creds:
            client_secret = creds['clientSecret']
        else:
            raise SynapseError(400, "No client_secret in creds")

        if not self._should_trust_id_server(id_server):
            logger.warn(
                '%s is not a trusted ID server: rejecting 3pid ' +
                'credentials', id_server
            )
            defer.returnValue(None)

        try:
            data = yield self.http_client.get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/3pid/getValidated3pid"
                ),
                {'sid': creds['sid'], 'client_secret': client_secret}
            )
        except HttpResponseException as e:
            logger.info("getValidated3pid failed with Matrix error: %r", e)
            raise e.to_synapse_error()

        if 'medium' in data:
            defer.returnValue(data)
        defer.returnValue(None)

    @defer.inlineCallbacks
    def bind_threepid(self, creds, mxid):
        logger.debug("binding threepid %r to %s", creds, mxid)
        data = None

        if 'id_server' in creds:
            id_server = creds['id_server']
        elif 'idServer' in creds:
            id_server = creds['idServer']
        else:
            raise SynapseError(400, "No id_server in creds")

        if 'client_secret' in creds:
            client_secret = creds['client_secret']
        elif 'clientSecret' in creds:
            client_secret = creds['clientSecret']
        else:
            raise SynapseError(400, "No client_secret in creds")

        try:
            data = yield self.http_client.post_urlencoded_get_json(
                "https://%s%s" % (
                    id_server, "/_matrix/identity/api/v1/3pid/bind"
                ),
                {
                    'sid': creds['sid'],
                    'client_secret': client_secret,
                    'mxid': mxid,
                }
            )
            logger.debug("bound threepid %r to %s", creds, mxid)

            # Remember where we bound the threepid
            yield self.store.add_user_bound_threepid(
                user_id=mxid,
                medium=data["medium"],
                address=data["address"],
                id_server=id_server,
            )
        except CodeMessageException as e:
            data = json.loads(e.msg)  # XXX WAT?
        defer.returnValue(data)

    @defer.inlineCallbacks
    def try_unbind_threepid(self, mxid, threepid):
        """Removes a binding from an identity server

        Args:
            mxid (str): Matrix user ID of binding to be removed
            threepid (dict): Dict with medium & address of binding to be
                removed, and an optional id_server.

        Raises:
            SynapseError: If we failed to contact the identity server

        Returns:
            Deferred[bool]: True on success, otherwise False if the identity
            server doesn't support unbinding (or no identity server found to
            contact).
        """
        if threepid.get("id_server"):
            id_servers = [threepid["id_server"]]
        else:
            id_servers = yield self.store.get_id_servers_user_bound(
                user_id=mxid,
                medium=threepid["medium"],
                address=threepid["address"],
            )

        # We don't know where to unbind, so we don't have a choice but to return
        if not id_servers:
            defer.returnValue(False)

        changed = True
        for id_server in id_servers:
            changed &= yield self.try_unbind_threepid_with_id_server(
                mxid, threepid, id_server,
            )

        defer.returnValue(changed)

    @defer.inlineCallbacks
    def try_unbind_threepid_with_id_server(self, mxid, threepid, id_server):
        """Removes a binding from an identity server

        Args:
            mxid (str): Matrix user ID of binding to be removed
            threepid (dict): Dict with medium & address of binding to be removed
            id_server (str): Identity server to unbind from

        Raises:
            SynapseError: If we failed to contact the identity server

        Returns:
            Deferred[bool]: True on success, otherwise False if the identity
            server doesn't support unbinding
        """
        url = "https://%s/_matrix/identity/api/v1/3pid/unbind" % (id_server,)
        content = {
            "mxid": mxid,
            "threepid": {
                "medium": threepid["medium"],
                "address": threepid["address"],
            },
        }

        # we abuse the federation http client to sign the request, but we have to send it
        # using the normal http client since we don't want the SRV lookup and want normal
        # 'browser-like' HTTPS.
        auth_headers = self.federation_http_client.build_auth_headers(
            destination=None,
            method='POST',
            url_bytes='/_matrix/identity/api/v1/3pid/unbind'.encode('ascii'),
            content=content,
            destination_is=id_server,
        )
        headers = {
            b"Authorization": auth_headers,
        }

        try:
            yield self.http_client.post_json_get_json(
                url,
                content,
                headers,
            )

            yield self.store.remove_user_bound_threepid(
                user_id=mxid,
                medium=threepid["medium"],
                address=threepid["address"],
                id_server=id_server,
            )
        except HttpResponseException as e:
            if e.code in (400, 404, 501,):
                # The remote server probably doesn't support unbinding (yet)
                logger.warn("Received %d response while unbinding threepid", e.code)
                defer.returnValue(False)
            else:
                logger.error("Failed to unbind threepid on identity server: %s", e)
                raise SynapseError(502, "Failed to contact identity server")

        defer.returnValue(True)

    @defer.inlineCallbacks
    def requestEmailToken(self, id_server, email, client_secret, send_attempt, **kwargs):
        if not self._should_trust_id_server(id_server):
            raise SynapseError(
                400, "Untrusted ID server '%s'" % id_server,
                Codes.SERVER_NOT_TRUSTED
            )

        params = {
            'email': email,
            'client_secret': client_secret,
            'send_attempt': send_attempt,
        }
        params.update(kwargs)

        try:
            data = yield self.http_client.post_json_get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/validate/email/requestToken"
                ),
                params
            )
            defer.returnValue(data)
        except HttpResponseException as e:
            logger.info("Proxied requestToken failed: %r", e)
            raise e.to_synapse_error()

    @defer.inlineCallbacks
    def requestMsisdnToken(
            self, id_server, country, phone_number,
            client_secret, send_attempt, **kwargs
    ):
        if not self._should_trust_id_server(id_server):
            raise SynapseError(
                400, "Untrusted ID server '%s'" % id_server,
                Codes.SERVER_NOT_TRUSTED
            )

        params = {
            'country': country,
            'phone_number': phone_number,
            'client_secret': client_secret,
            'send_attempt': send_attempt,
        }
        params.update(kwargs)

        try:
            data = yield self.http_client.post_json_get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/validate/msisdn/requestToken"
                ),
                params
            )
            defer.returnValue(data)
        except HttpResponseException as e:
            logger.info("Proxied requestToken failed: %r", e)
            raise e.to_synapse_error()
<commit_msg>Remove threepid binding if id server returns 400/404/501<commit_after># -*- coding: utf-8 -*-
# Copyright 2015, 2016 OpenMarket Ltd
# Copyright 2017 Vector Creations Ltd
# Copyright 2018 New Vector Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for interacting with Identity Servers"""

import logging

from canonicaljson import json

from twisted.internet import defer

from synapse.api.errors import (
    CodeMessageException,
    Codes,
    HttpResponseException,
    SynapseError,
)

from ._base import BaseHandler

logger = logging.getLogger(__name__)

class IdentityHandler(BaseHandler):

    def __init__(self, hs):
        super(IdentityHandler, self).__init__(hs)

        self.http_client = hs.get_simple_http_client()
        self.federation_http_client = hs.get_http_client()

        self.trusted_id_servers = set(hs.config.trusted_third_party_id_servers)
        self.trust_any_id_server_just_for_testing_do_not_use = (
            hs.config.use_insecure_ssl_client_just_for_testing_do_not_use
        )

    def _should_trust_id_server(self, id_server):
        if id_server not in self.trusted_id_servers:
            if self.trust_any_id_server_just_for_testing_do_not_use:
                logger.warn(
                    "Trusting untrustworthy ID server %r even though it isn't"
                    " in the trusted id list for testing because"
                    " 'use_insecure_ssl_client_just_for_testing_do_not_use'"
                    " is set in the config",
                    id_server,
                )
            else:
                return False
        return True

    @defer.inlineCallbacks
    def threepid_from_creds(self, creds):
        if 'id_server' in creds:
            id_server = creds['id_server']
        elif 'idServer' in creds:
            id_server = creds['idServer']
        else:
            raise SynapseError(400, "No id_server in creds")

        if 'client_secret' in creds:
            client_secret = creds['client_secret']
        elif 'clientSecret' in creds:
            client_secret = creds['clientSecret']
        else:
            raise SynapseError(400, "No client_secret in creds")

        if not self._should_trust_id_server(id_server):
            logger.warn(
                '%s is not a trusted ID server: rejecting 3pid ' +
                'credentials', id_server
            )
            defer.returnValue(None)

        try:
            data = yield self.http_client.get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/3pid/getValidated3pid"
                ),
                {'sid': creds['sid'], 'client_secret': client_secret}
            )
        except HttpResponseException as e:
            logger.info("getValidated3pid failed with Matrix error: %r", e)
            raise e.to_synapse_error()

        if 'medium' in data:
            defer.returnValue(data)
        defer.returnValue(None)

    @defer.inlineCallbacks
    def bind_threepid(self, creds, mxid):
        logger.debug("binding threepid %r to %s", creds, mxid)
        data = None

        if 'id_server' in creds:
            id_server = creds['id_server']
        elif 'idServer' in creds:
            id_server = creds['idServer']
        else:
            raise SynapseError(400, "No id_server in creds")

        if 'client_secret' in creds:
            client_secret = creds['client_secret']
        elif 'clientSecret' in creds:
            client_secret = creds['clientSecret']
        else:
            raise SynapseError(400, "No client_secret in creds")

        try:
            data = yield self.http_client.post_urlencoded_get_json(
                "https://%s%s" % (
                    id_server, "/_matrix/identity/api/v1/3pid/bind"
                ),
                {
                    'sid': creds['sid'],
                    'client_secret': client_secret,
                    'mxid': mxid,
                }
            )
            logger.debug("bound threepid %r to %s", creds, mxid)

            # Remember where we bound the threepid
            yield self.store.add_user_bound_threepid(
                user_id=mxid,
                medium=data["medium"],
                address=data["address"],
                id_server=id_server,
            )
        except CodeMessageException as e:
            data = json.loads(e.msg)  # XXX WAT?
        defer.returnValue(data)

    @defer.inlineCallbacks
    def try_unbind_threepid(self, mxid, threepid):
        """Removes a binding from an identity server

        Args:
            mxid (str): Matrix user ID of binding to be removed
            threepid (dict): Dict with medium & address of binding to be
                removed, and an optional id_server.

        Raises:
            SynapseError: If we failed to contact the identity server

        Returns:
            Deferred[bool]: True on success, otherwise False if the identity
            server doesn't support unbinding (or no identity server found to
            contact).
        """
        if threepid.get("id_server"):
            id_servers = [threepid["id_server"]]
        else:
            id_servers = yield self.store.get_id_servers_user_bound(
                user_id=mxid,
                medium=threepid["medium"],
                address=threepid["address"],
            )

        # We don't know where to unbind, so we don't have a choice but to return
        if not id_servers:
            defer.returnValue(False)

        changed = True
        for id_server in id_servers:
            changed &= yield self.try_unbind_threepid_with_id_server(
                mxid, threepid, id_server,
            )

        defer.returnValue(changed)

    @defer.inlineCallbacks
    def try_unbind_threepid_with_id_server(self, mxid, threepid, id_server):
        """Removes a binding from an identity server

        Args:
            mxid (str): Matrix user ID of binding to be removed
            threepid (dict): Dict with medium & address of binding to be removed
            id_server (str): Identity server to unbind from

        Raises:
            SynapseError: If we failed to contact the identity server

        Returns:
            Deferred[bool]: True on success, otherwise False if the identity
            server doesn't support unbinding
        """
        url = "https://%s/_matrix/identity/api/v1/3pid/unbind" % (id_server,)
        content = {
            "mxid": mxid,
            "threepid": {
                "medium": threepid["medium"],
                "address": threepid["address"],
            },
        }

        # we abuse the federation http client to sign the request, but we have to send it
        # using the normal http client since we don't want the SRV lookup and want normal
        # 'browser-like' HTTPS.
        auth_headers = self.federation_http_client.build_auth_headers(
            destination=None,
            method='POST',
            url_bytes='/_matrix/identity/api/v1/3pid/unbind'.encode('ascii'),
            content=content,
            destination_is=id_server,
        )
        headers = {
            b"Authorization": auth_headers,
        }

        try:
            yield self.http_client.post_json_get_json(
                url,
                content,
                headers,
            )
            changed = True
        except HttpResponseException as e:
            changed = False
            if e.code in (400, 404, 501,):
                # The remote server probably doesn't support unbinding (yet)
                logger.warn("Received %d response while unbinding threepid", e.code)
            else:
                logger.error("Failed to unbind threepid on identity server: %s", e)
                raise SynapseError(502, "Failed to contact identity server")

        yield self.store.remove_user_bound_threepid(
            user_id=mxid,
            medium=threepid["medium"],
            address=threepid["address"],
            id_server=id_server,
        )

        defer.returnValue(changed)

    @defer.inlineCallbacks
    def requestEmailToken(self, id_server, email, client_secret, send_attempt, **kwargs):
        if not self._should_trust_id_server(id_server):
            raise SynapseError(
                400, "Untrusted ID server '%s'" % id_server,
                Codes.SERVER_NOT_TRUSTED
            )

        params = {
            'email': email,
            'client_secret': client_secret,
            'send_attempt': send_attempt,
        }
        params.update(kwargs)

        try:
            data = yield self.http_client.post_json_get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/validate/email/requestToken"
                ),
                params
            )
            defer.returnValue(data)
        except HttpResponseException as e:
            logger.info("Proxied requestToken failed: %r", e)
            raise e.to_synapse_error()

    @defer.inlineCallbacks
    def requestMsisdnToken(
            self, id_server, country, phone_number,
            client_secret, send_attempt, **kwargs
    ):
        if not self._should_trust_id_server(id_server):
            raise SynapseError(
                400, "Untrusted ID server '%s'" % id_server,
                Codes.SERVER_NOT_TRUSTED
            )

        params = {
            'country': country,
            'phone_number': phone_number,
            'client_secret': client_secret,
            'send_attempt': send_attempt,
        }
        params.update(kwargs)

        try:
            data = yield self.http_client.post_json_get_json(
                "https://%s%s" % (
                    id_server,
                    "/_matrix/identity/api/v1/validate/msisdn/requestToken"
                ),
                params
            )
            defer.returnValue(data)
        except HttpResponseException as e:
            logger.info("Proxied requestToken failed: %r", e)
            raise e.to_synapse_error()
<|endoftext|>


<commit_before>import gzip
import sys
import os
from django.core.management.base import BaseCommand
from xbrowse_server.base.models import Project
from xbrowse.parsers import vcf_stuff, fam_stuff
from xbrowse_server.phenotips.utilities import add_individuals_to_phenotips

class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('args', nargs='*')
        parser.add_argument('--vcf',
                            dest='vcf',
                            help='A VCF file to gather patient information from.'
                            )
        parser.add_argument('--ped',
                            dest='ped',
                            help='A PED file to gather patient information from (PREFERRED due to richer information set).'
                            )

    def handle(self, *args, **options):
        """
        Handles populating Phenotips.
        Currently favors PED file for richer set of information (vs VCF) to create patients
        """

        if len(args) == 0 or options['vcf'] is None and options['ped'] is None:
            print '\n\nAdds a series of patients to Phenotips.\n'
            print 'Please enter a VCF file (--vcf), OR IDEALLY A PED file (--ped), and a project ID (first positional argument).'
            print 'For example: python manage.py add_individuals_to_phenotips  myProjectId --ped myPed.ped\n'
            sys.exit()

        project_id = args[0]
        project = Project.objects.get(project_id=project_id)

        if options.get('vcf') and not os.path.exists(options.get('vcf')):
            print '\n\nError: the VCF file you entered does not exist or is invalid\n\n'
            sys.exit()
        if options.get('ped') and not os.path.exists(options.get('ped')):
            print '\n\nError: the PED file you entered does not exist or is invalid\n\n'
            sys.exit()

        indiv_id_list = None
        if options.get('vcf'):
            vcf_path = options.get('vcf')
            if vcf_path.endswith('.gz'):
                vcf = gzip.open(vcf_path)
            else:
                vcf = open(vcf_path)
            indiv_id_list = vcf_stuff.get_ids_from_vcf(vcf)
            add_individuals_to_phenotips(project_id, indiv_id_list)
        elif options.get('ped'):
            fam_file = open(options.get('ped'))
            individuals = fam_stuff.get_individuals_from_fam_file(fam_file)
            indiv_id_list = []
            for ind in individuals:
                indiv_id_list.append(ind.indiv_id)

        # if no vcf or ped file was specified, add all individuals in this projects
        print("Creating phenotips records for new individuals.")
        add_individuals_to_phenotips(project_id, indiv_id_list)
        print("Done.")
<commit_msg>Added --all option<commit_after>import gzip
import sys
import os
from django.core.management.base import BaseCommand
from xbrowse_server.base.models import Project, Individual
from xbrowse.parsers import vcf_stuff, fam_stuff
from xbrowse_server.phenotips.utilities import add_individuals_to_phenotips

class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('args', nargs='*')
        parser.add_argument('--vcf', dest='vcf',
                            help='A VCF file to gather patient information from.' )
        parser.add_argument('--ped', dest='ped',
                            help='A PED file to gather patient information from (PREFERRED due to richer information set).' )
        parser.add_argument('--all', dest='all', action="store_true",
                            help='A PED file to gather patient information from (PREFERRED due to richer information set).' )

    def handle(self, *args, **options):
        """
        Handles populating Phenotips.
        Currently favors PED file for richer set of information (vs VCF) to create patients
        """

        if len(args) == 0 or options['vcf'] is None and options['ped'] is None and options['all'] is None:
            print('\n\nAdds a series of patients to Phenotips.\n')
            print('Please enter a project id (first positional argument) followed by either a PED file (--ped), a VCF file (--vcf), or --all ('
                  'if you just want to create PhenoTips patients for the individuals that already exist in the project)')
            print('For example: python manage.py add_individuals_to_phenotips  myProjectId --ped myPed.ped\n')
            sys.exit()

        project_id = args[0]
        project = Project.objects.get(project_id=project_id)

        if options.get('vcf') and not os.path.exists(options.get('vcf')):
            print '\n\nError: the VCF file you entered does not exist or is invalid\n\n'
            sys.exit()
        if options.get('ped') and not os.path.exists(options.get('ped')):
            print '\n\nError: the PED file you entered does not exist or is invalid\n\n'
            sys.exit()

        indiv_id_list = None
        if options.get('vcf'):
            vcf_path = options.get('vcf')
            if vcf_path.endswith('.gz'):
                vcf = gzip.open(vcf_path)
            else:
                vcf = open(vcf_path)
            indiv_id_list = vcf_stuff.get_ids_from_vcf(vcf)
            add_individuals_to_phenotips(project_id, indiv_id_list)
        elif options.get('ped'):
            fam_file = open(options.get('ped'))
            individuals = fam_stuff.get_individuals_from_fam_file(fam_file)
            indiv_id_list = [ind.indiv_id for ind in individuals]

        elif options.get("all"):
            indiv_id_list = [ind.indiv_id for ind in Individual.objects.filter(project=project)]

        # if no vcf or ped file was specified, add all individuals in this projects
        print("Creating phenotips records for new individuals.")
        add_individuals_to_phenotips(project_id, indiv_id_list)
        print("Done.")
<|endoftext|>


<commit_before>import json
from os import path

PROJECT_DIR = path.dirname(path.dirname(path.abspath(__file__)))
BASE_DIR = path.dirname(PROJECT_DIR)

def load_config_json(name):
    json_path = path.abspath(path.join(BASE_DIR,
                                       'microservice',
                                       '{name}.json'))
    with open(json_path) as data_file:
        info = json.load(data_file)
    return info
<commit_msg>fixed string formatting function<commit_after>import json
from os import path

PROJECT_DIR = path.dirname(path.dirname(path.abspath(__file__)))
BASE_DIR = path.dirname(PROJECT_DIR)

def load_config_json(name):
    json_path = path.abspath(path.join(BASE_DIR,
                                       'microservice',
                                       f'{name}.json'))
    with open(json_path) as data_file:
        info = json.load(data_file)
    return info
<|endoftext|>


<commit_before>from __future__ import absolute_import, print_function

import logging
import pprint

from googleapiclient.discovery import build as google_build
from oauth2client.client import GoogleCredentials
from six import iteritems

from . import utils
from .s3 import events_s3
from .schema import mk_bigquery_schema

__license__ = """
Copyright 2016 Parsely, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

log = logging.getLogger(__name__)

def streaming_insert_bigquery(jsonlines,
                              bq_conn=None,
                              project_id=None,
                              dataset_id=None,
                              table_id=None):
    """Write a stream of events to BigQuery

    :param bq_conn: The BigQuery connection to write to
    :type bq_conn: googleapiclient.discovery.Resource
    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    """
    insert_body = {
        "kind": "bigquery#tableDataInsertAllRequest",
        "skipInvalidRows": False,
        "ignoreUnknownValues": False,
        "rows": [jsonline for jsonline in jsonlines]
    }
    if bq_conn is None:
        pprint.PrettyPrinter(indent=0).pprint(insert_body)
        return False
    query = bq_conn.tabledata().insertAll(projectId=project_id,
                                          datasetId=dataset_id,
                                          tableId=table_id,
                                          body=insert_body)
    response = query.execute(num_retries=5)
    if 'insertErrors' in response:
        for error_set in response['insertErrors']:
            for error in error_set['errors']:
                log.error(error)
        return False
    return True

def copy_from_s3(network,
                 s3_prefix="",
                 access_key_id="",
                 secret_access_key="",
                 region_name="us-east-1",
                 project_id=None,
                 dataset_id=None,
                 table_id=None,
                 dry_run=False):
    """Load events from S3 to BigQuery using the BQ streaming insert API.

    :param network: The Parse.ly network for which to perform writes (eg
        "parsely-blog")
    :type network: str
    :param s3_prefix: The S3 timestamp directory prefix from which to fetch data
        batches, formatted as YYYY/MM/DD
    :type s3_prefix: str
    :param access_key_id: The AWS access key to use when fetching data batches
    :type access_key_id: str
    :param secret_access_key: The AWS secret key to use when fetching data batches
    :type secret_access_key: str
    :param region_name: The AWS region in which to perform fetches
    :type region_name: str
    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    :param dry_run: If True, don't perform BigQuery writes
    :type dry_run: bool
    """
    bq_conn = None
    if not dry_run:
        bq_conn = google_build(
            'bigquery', 'v2',
            credentials=GoogleCredentials.get_application_default())
    s3_stream = events_s3(network, prefix=s3_prefix, access_key_id=access_key_id,
                          secret_access_key=secret_access_key,
                          region_name=region_name)

    def chunked(seq, chunk_size):
        chunk = []
        for item in seq:
            chunk.append(item)
            if len(chunk) >= chunk_size:
                yield chunk
                chunk = []
    for events in chunked(s3_stream, 500):
        write_events_bigquery(events, bq_conn=bq_conn, project_id=project_id,
                              dataset_id=dataset_id, table_id=table_id)

def create_table(project_id, table_id, dataset_id, debug=False):
    """Create a BigQuery table using a schema compatible with Parse.ly events

    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    """
    fields = mk_bigquery_schema()
    schema = {
        "description": "Parse.ly Data Pipeline",
        "schema": {"fields": fields},
        "tableReference": {
            "projectId": project_id,
            "tableId": table_id,
            "datasetId": dataset_id
        }
    }
    if debug:
        print("Running the following BigQuery JSON table insert:")
        print(json.dumps(schema, indent=4, sort_keys=True))
    credentials = GoogleCredentials.get_application_default()
    bigquery = google_build('bigquery', 'v2', credentials=credentials)
    bigquery.tables().insert(projectId=project_id,
                             datasetId=dataset_id,
                             tableId=table_id,
                             body=schema).execute()

def main():
    commands = ["copy_from_s3", "create_table"]
    parser = utils.get_default_parser("Google BigQuery utilities for Parse.ly",
                                      commands=commands)
    parser.add_argument('--dry_run', action="store_true",
                        help="If true, don't perform writes to Bigquery"
                             'connect, ending in "redshift.amazonaws.com"')
    parser.add_argument('--bigquery_project_id', type=str,
                        help='The ID of the BigQuery project to which to connect')
    parser.add_argument('--bigquery_dataset_id', type=str,
                        help='The ID of the BigQuery dataset to which to connect')
    parser.add_argument('--bigquery_table_id', type=str,
                        help='The ID of the BigQuery table to which to connect')
    args = parser.parse_args()

    if args.command == "copy_from_s3":
        copy_from_s3(
            args.network,
            s3_prefix=args.s3_prefix,
            access_key_id=args.aws_access_key_id,
            secret_access_key=args.aws_secret_access_key,
            region_name=args.aws_region_name,
            project_id=args.bigquery_project_id,
            dataset_id=args.bigquery_dataset_id,
            table_id=args.bigquery_table_id,
            dry_run=args.dry_run
        )
    elif args.command == "create_table":
        create_table(
            project_id=args.project_id,
            dataset_id=args.dataset_id,
            table_id=args.table_id
        )

if __name__ == "__main__":
    main()
<commit_msg>Add json import, homogenize args names<commit_after>from __future__ import absolute_import, print_function

import logging
import pprint
import json

from googleapiclient.discovery import build as google_build
from oauth2client.client import GoogleCredentials
from six import iteritems

from . import utils
from .s3 import events_s3
from .schema import mk_bigquery_schema

__license__ = """
Copyright 2016 Parsely, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

log = logging.getLogger(__name__)

def streaming_insert_bigquery(jsonlines,
                              bq_conn=None,
                              project_id=None,
                              dataset_id=None,
                              table_id=None):
    """Write a stream of events to BigQuery

    :param bq_conn: The BigQuery connection to write to
    :type bq_conn: googleapiclient.discovery.Resource
    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    """
    insert_body = {
        "kind": "bigquery#tableDataInsertAllRequest",
        "skipInvalidRows": False,
        "ignoreUnknownValues": False,
        "rows": [jsonline for jsonline in jsonlines]
    }
    if bq_conn is None:
        pprint.PrettyPrinter(indent=0).pprint(insert_body)
        return False
    query = bq_conn.tabledata().insertAll(projectId=project_id,
                                          datasetId=dataset_id,
                                          tableId=table_id,
                                          body=insert_body)
    response = query.execute(num_retries=5)
    if 'insertErrors' in response:
        for error_set in response['insertErrors']:
            for error in error_set['errors']:
                log.error(error)
        return False
    return True

def copy_from_s3(network,
                 s3_prefix="",
                 access_key_id="",
                 secret_access_key="",
                 region_name="us-east-1",
                 project_id=None,
                 dataset_id=None,
                 table_id=None,
                 dry_run=False):
    """Load events from S3 to BigQuery using the BQ streaming insert API.

    :param network: The Parse.ly network for which to perform writes (eg
        "parsely-blog")
    :type network: str
    :param s3_prefix: The S3 timestamp directory prefix from which to fetch data
        batches, formatted as YYYY/MM/DD
    :type s3_prefix: str
    :param access_key_id: The AWS access key to use when fetching data batches
    :type access_key_id: str
    :param secret_access_key: The AWS secret key to use when fetching data batches
    :type secret_access_key: str
    :param region_name: The AWS region in which to perform fetches
    :type region_name: str
    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    :param dry_run: If True, don't perform BigQuery writes
    :type dry_run: bool
    """
    bq_conn = None
    if not dry_run:
        bq_conn = google_build(
            'bigquery', 'v2',
            credentials=GoogleCredentials.get_application_default())
    s3_stream = events_s3(network, prefix=s3_prefix, access_key_id=access_key_id,
                          secret_access_key=secret_access_key,
                          region_name=region_name)

    def chunked(seq, chunk_size):
        chunk = []
        for item in seq:
            chunk.append(item)
            if len(chunk) >= chunk_size:
                yield chunk
                chunk = []
    for events in chunked(s3_stream, 500):
        write_events_bigquery(events, bq_conn=bq_conn, project_id=project_id,
                              dataset_id=dataset_id, table_id=table_id)

def create_table(project_id, table_id, dataset_id, debug=False):
    """Create a BigQuery table using a schema compatible with Parse.ly events

    :param project_id: The BigQuery project ID to write to
    :type project_id: str
    :param table_id: The BigQuery table ID to write to
    :type table_id: str
    :param dataset_id: The BigQuery dataset ID to write to
    :type dataset_id: str
    """
    fields = mk_bigquery_schema()
    schema = {
        "description": "Parse.ly Data Pipeline",
        "schema": {"fields": fields},
        "tableReference": {
            "projectId": project_id,
            "tableId": table_id,
            "datasetId": dataset_id
        }
    }
    if debug:
        print("Running the following BigQuery JSON table insert:")
        print(json.dumps(schema, indent=4, sort_keys=True))
    credentials = GoogleCredentials.get_application_default()
    bigquery = google_build('bigquery', 'v2', credentials=credentials)
    bigquery.tables().insert(projectId=project_id,
                             datasetId=dataset_id,
                             tableId=table_id,
                             body=schema).execute()

def main():
    commands = ["copy_from_s3", "create_table"]
    parser = utils.get_default_parser("Google BigQuery utilities for Parse.ly",
                                      commands=commands)
    parser.add_argument('--dry_run', action="store_true",
                        help="If true, don't perform writes to Bigquery"
                             'connect, ending in "redshift.amazonaws.com"')
    parser.add_argument('--bigquery_project_id', type=str,
                        help='The ID of the BigQuery project to which to connect')
    parser.add_argument('--bigquery_dataset_id', type=str,
                        help='The ID of the BigQuery dataset to which to connect')
    parser.add_argument('--bigquery_table_id', type=str,
                        help='The ID of the BigQuery table to which to connect')
    args = parser.parse_args()

    if args.command == "copy_from_s3":
        copy_from_s3(
            args.network,
            s3_prefix=args.s3_prefix,
            access_key_id=args.aws_access_key_id,
            secret_access_key=args.aws_secret_access_key,
            region_name=args.aws_region_name,
            project_id=args.bigquery_project_id,
            dataset_id=args.bigquery_dataset_id,
            table_id=args.bigquery_table_id,
            dry_run=args.dry_run
        )
    elif args.command == "create_table":
        create_table(
            project_id=args.bigquery_project_id,
            dataset_id=args.bigquery_dataset_id,
            table_id=args.bigquery_table_id,
            debug=args.debug
        )

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


<commit_before># SMOP compiler runtime support library
# Copyright 2014 Victor Leikehman

# MIT license

import numpy

def ceil_(t):
    '''
    >>> ceil_(-1.1)
    -1.0
    '''
    return numpy.ceil(t)

def false_(*args):
    return numpy.zeros(args,dtype=bool)

def find_(a,nargout=2):
    i,j = a.nonzero()
    return (i+1),(j+1)

def floor_(t):
    '''
    >>> floor_(-0.1)
    -1.0
    '''
    return numpy.floor(t)

def inf_(*args):
    t = numpy.empty(numpy.prod(args))
    t[:] = numpy.inf
    return t.reshape(args)

Inf_ = inf_

def isequal_(a,b):
    return numpy.array_equal(a,b)

def length_(t):
    '''
    >>> length_(zeros_(3,4,5))
    5
    >>> length_(12345)
    1
    >>> length_([3,4,5])
    3
    '''
    return max(numpy.shape(t) or (1,))

def max_(t, d=0, nargout=0):
    '''
    >>> max_(123)
    123
    >>> max_(range(10))
    9
    >>> max_(numpy.arange(10))
    9
    '''
    if d or nargout:
        raise ErrorNotImplemented
    return numpy.amax(t)

def min_(t, d=0, nargout=0):
    '''
    >>> min_(123)
    123
    >>> min_(range(10))
    0
    >>> min_(numpy.arange(10))
    0
    '''
    if d or nargout:
        raise ErrorNotImplemented
    return numpy.amin(t)

def ones_(*args):
        return numpy.ones(args)

def rand_(*args):
    '''
    >>> size_(rand_())
    (1, 1)
    >>> size_(rand_(2))
    (2, 2)
    >>> size_(rand_(2,3))
    (2, 3)
    '''
    if len(args) == 1:
        args += args
    return numpy.random.rand(*args)

def ravel_(t):
    '''
    >>> size_(ravel_(3))
    (1, 1)
    >>> size_(ravel_(zeros_(3,3)))
    (9, 1)
    '''
    return numpy.reshape(t,[-1,1])

def round_(t):
    '''
    >>> round_(0.5)
    0.0
    >>> round_(0.6)
    1.0
    '''
    return numpy.round(t)

def size_(a, b=0, nargout=2):
    '''
    >>> size_(123, 100)
    1
    >>> size_(123)
    (1, 1)
    >>> (p,q,r) = size_(123,nargout=3)
    >>> (p,q,r)
    (1, 1, 1)
    >>> z = zeros_(2,3,4)
    >>> a = size_(z)
    >>> a
    (2, 3, 4)
    >>> b = zeros_(2,3)
    >>> (m,n) = size_(b,nargout=2)
    >>> (m,n)
    (2, 3)
    '''
    s = numpy.shape(a)
    if s is ():
        return 1 if b else (1,)*nargout
    # a is not a scalar
    return s[b-1] if b else s

def true_(*args):
    return numpy.ones(args,dtype=bool)

def zeros_(*args,**kwargs):
    '''
    >>> size_(zeros_(2,3))
    (2, 3)
    '''
    return numpy.zeros(args,**kwargs)

if __name__ == "__main__":
    import doctest
    doctest.testmod()

# vim:et:sw=4:si:
<commit_msg>work in progress<commit_after># SMOP compiler runtime support library
# Copyright 2014 Victor Leikehman

# MIT license
import numpy as np
import os
from scipy.io import loadmat

# abs sort find1 numel rand('seed') sum(a,d) log dot cumsum any(a,d) sqrt proda isinf sign mod

class matlabarray(np.ndarray):
    def __new__(cls,input_array=[],**kwargs):
        #obj = np.array(input_array,ndmin=2,order="F",copy=None).view(cls)
        obj = np.asarray(input_array,order="F").view(cls)
        return obj

    def __array_finalize__(self,obj):
        self.delegate = obj

    def compute_indices(self,index):
        if not isinstance(index,tuple):
           index = index,
        if len(index) != 1 and len(index) != self.delegate.ndim:
            raise IndexError()
        indices = []
        for i,ix in enumerate(index):
            if ix.__class__ is slice:
                if len(index) == 1:
                    n = self.delegate.size
                else:
                    n = self.delegate.shape[i]
                start,stop,step = ix.indices(n)
                indices.append(slice(start-1 if start else 0,
                                     stop,
                                     step))
            else:
                try:
                    indices.append(ix-1)
                except TypeError:
                    indices.append(np.asarray(ix)-1)
        return tuple(indices)

    def __getslice__(self,i,j):
        return self.__getitem__(slice(i,j,None))

    def __getitem__(self,index):
        if index  == slice(None):
            return self.delegate.reshape(-1,1,order="F")
        indices = self.compute_indices(index)
        if len(indices) == 1:
            return self.delegate.reshape(-1,order="F").__getitem__(indices)
        else:
            return self.delegate.__getitem__(indices)

    def __setitem__(self,index,value):
        indices = self.compute_indices(index)
        for j in (1,2):
            try:
                if isinstance(index,tuple):
                    np.ndarray.__setitem__(self.delegate,indices,value)
                else:
                    np.ndarray.__setitem__(self.delegate.reshape(-1,order="F"),indices,value)
                break
            except IndexError:
                if isinstance(index,tuple):
                    new_shape = [max(self.delegate.shape[d], i+1) for d,i in enumerate(indices)]
                    # new_shape = [self.new_size(i) for i in indices]
                else:
                    raise NotImplementedError("Expanding arrays not implemented yet")
                #self.delegate = np.resize(self.delegate,new_shape)
                self.delegate = self.delegate.copy()
                self.delegate.resize(new_shape)

    def reshape_(self,*args):
        return self.delegate.reshape(*args,order="F")

#    def __copy__(self):
#        obj = matlabarray()
#        obj.delegate = self.delegate.copy()
#        return obj

    def __repr__(self):
        return self.delegate.__repr__()

    def __str__(self):
        return self.delegate.__str__()

    def __add__(self,other):
        return np.ndarray.__add__(self,other)

    def __neg__(self):
        return matlabarray(self.delegate.__neg__())

def asarray(a):
    return a.delegate if isinstance(a,matlabarray) else np.asarray(a,order="F")

def abs_(a):
    return np.abs(asarray(a))

def arange_(*args):
#    if args is ():
#        return slice(None,None,None)
#    if len(args) == 1:
#        return np.arange(1,args[0]+1).reshape(-1,1)
#   if len(args) == 2:
        return np.arange(args[0],args[1]+1).reshape(-1,1)
#   if len(args) == 3:
#       return np.arange(args[0],args[1]+1,args[2]).reshape(-1,1)
#   assert 0

def ceil_(a):
    return np.ceil(asarray(a))

def disp_(*args):
    print (args)

def false_(*args):
    if len(args) == 1:
        args += args
    return np.zeros(args,dtype=bool,order="F")

def find_(a,nargout=2):
    i,j = asarray(a).nonzero()
    return (i+1),(j+1)

def floor_(a):
    return np.floor_(asarray(a))

def fullfile_(*args):
    return os.path.join(*args)

def intersect_(a,b,nargout=1):
    if nargout == 1:
        c = sorted(set(a) & set(b))
        if isinstance(a,str):
            return "".join(c)
        elif isinstance(a,list):
            return c
        else:
            # FIXME: the result is a column vector if
            # both args are column vectors; otherwise row vector
            return np.array(c)
    raise ErrorNotImplemented

#def inf_(*args):
#    t = np.empty(np.prod(args))
#    t[:] = np.inf
#    return t.reshape(args)
#
#Inf_ = inf_
#
#def int_(t):
#    return np.int(t)
#
#def int32_(t):
#    return np.int(t)
#
def isempty_(a):
    try:
        return 0 in asarray(a).shape
    except AttributeError:
        return False

def isequal_(a,b):
    return np.array_equal(asarray(a),
                          asarray(b))

def length_(a):
    try:
        return max(asarray(a).shape)
    except ValueError:
        return 1

def load_(a):
    return loadmat(a) # FIXME

def max_(a, d=0, nargout=0):
    if d or nargout:
        raise ErrorNotImplemented()
    return asarray(a).amax()

def min_(a, d=0, nargout=0):
    if d or nargout:
        raise ErrorNotImplemented()
    return asarray(a).amin()

def numel_(a):
    return asarray(a).size

def ones_(*args,**kwargs):
        return matlabarray(np.ones(args,dtype="int",**kwargs))

def rand_(*args):
    if len(args) == 1:
        args += args
    return np.rand(args,order="F")

def ravel_(a):
    return asarray(a).reshape(-1,1)

def round_(a):
    return np.round(asarray(a))

def size_(a, b=0, nargout=2):
    s = asarray(a).shape
    if s is ():
        return 1 if b else (1,)*nargout
    # a is not a scalar
    try:
        return s[b-1] if b else s
    except IndexError:
        return 1

def sum_(a):
    return asarray(a).sum()

def true_(*args):
    if len(args) == 1:
        args += args
    return np.ones(args,dtype=bool,order="F")

def zeros_(*args):
    if not args:
        return 0.0
    if len(args) == 1:
        args += args
    return np.zeros(args,order="F")

# vim:et:sw=4:si:
<|endoftext|>


<commit_before># nhlib: A New Hazard Library
# Copyright (C) 2012 GEM Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
Package :mod:`nhlib.calc` contains hazard calculator modules
and utilities for them, such as :mod:`~nhlib.calc.filters`.
"""
from nhlib.calc.hazard_curve import hazard_curves_poissonian
from nhlib.calc import filters
<commit_msg>calc/__init__: import stochastic_event_set()<commit_after># nhlib: A New Hazard Library
# Copyright (C) 2012 GEM Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
Package :mod:`nhlib.calc` contains hazard calculator modules
and utilities for them, such as :mod:`~nhlib.calc.filters`.
"""
from nhlib.calc.hazard_curve import hazard_curves_poissonian
from nhlib.calc.stochastic import stochastic_event_set
from nhlib.calc import filters
<|endoftext|>