<commit_before># Copyright (c) 2015-2017 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

import numpy as np
import healpy as hp

from .. import qarray as qa
from .tod import TOD
from ..op import Operator
from ..map import DistRings, PySMSky, LibSharpSmooth, DistPixels
from ..mpi import MPI
from ..ctoast import sim_map_scan_map

class OpSimGradient(Operator):
    """
    Generate a fake sky signal as a gradient between the poles.

    This passes through each observation and creates a fake signal timestream
    based on the cartesian Z coordinate of the HEALPix pixel containing the
    detector pointing.

    Args:
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
        nside (int): the HEALPix NSIDE value to use.
        min (float): the minimum value to use at the South Pole.
        max (float): the maximum value to use at the North Pole.
        nest (bool): whether to use NESTED ordering.
    """

    def __init__(self, out='grad', nside=512, min=-100.0, max=100.0, nest=False,
                 flag_mask=255, common_flag_mask=255):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._nside = nside
        self._out = out
        self._min = min
        self._max = max
        self._nest = nest
        self._flag_mask = flag_mask
        self._common_flag_mask = common_flag_mask

    def exec(self, data):
        """
        Create the gradient timestreams.

        This pixelizes each detector's pointing and then assigns a
        timestream value based on the cartesian Z coordinate of the pixel
        center.

        Args:
            data (toast.Data): The distributed data.
        """
        comm = data.comm

        zaxis = np.array([0,0,1], dtype=np.float64)
        nullquat = np.array([0,0,0,1], dtype=np.float64)

        range = self._max - self._min

        for obs in data.obs:
            tod = obs['tod']
            base = obs['baselines']
            nse = obs['noise']
            intrvl = obs['intervals']

            for det in tod.local_dets:
                pdata = np.copy(tod.read_pntg(detector=det, local_start=0,
                                              n=tod.local_samples[1]))
                flags, common = tod.read_flags(detector=det, local_start=0,
                                               n=tod.local_samples[1])
                totflags = flags & self._flag_mask
                totflags |= (common & self._common_flag_mask)

                del flags
                del common

                pdata[totflags != 0,:] = nullquat

                dir = qa.rotate(pdata, zaxis)
                pixels = hp.vec2pix(self._nside, dir[:,0], dir[:,1], dir[:,2],
                                    nest=self._nest)
                x, y, z = hp.pix2vec(self._nside, pixels, nest=self._nest)
                z += 1.0
                z *= 0.5
                z *= range
                z += self._min
                z[totflags != 0] = 0.0

                cachename = "{}_{}".format(self._out, det)
                if not tod.cache.exists(cachename):
                    tod.cache.create(cachename, np.float64,
                                     (tod.local_samples[1],))
                ref = tod.cache.reference(cachename)
                ref[:] += z
                del ref
                #print('Grad timestream:', ref, np.sum(ref!=0),' non-zeros', flush=True) # DEBUG

        return

    def sigmap(self):
        """
        (array): Return the underlying signal map (full map on all processes).
        """
        range = self._max - self._min
        pix = np.arange(0, 12*self._nside*self._nside, dtype=np.int64)
        x, y, z = hp.pix2vec(self._nside, pix, nest=self._nest)
        z += 1.0
        z *= 0.5
        z *= range
        z += self._min
        return z

class OpSimScan(Operator):
    """
    Operator which generates sky signal by scanning from a map.

    The signal to use should already be in a distributed pixel structure,
    and local pointing should already exist.

    Args:
        distmap (DistPixels): the distributed map domain data.
        pixels (str): the name of the cache object (<pixels>_<detector>)
            containing the pixel indices to use.
        weights (str): the name of the cache object (<weights>_<detector>)
            containing the pointing weights to use.
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
    """
    def __init__(self, distmap=None, pixels='pixels', weights='weights',
                 out='scan'):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._map = distmap
        self._pixels = pixels
        self._weights = weights
        self._out = out

    def exec(self, data):
        """
        Create the timestreams by scanning from the map.

        This loops over all observations and detectors and uses the pointing
        matrix to project the distributed map into a timestream.

        Args:
            data (toast.Data): The distributed data.
        """
        comm = data.comm
        # the global communicator
        cworld = comm.comm_world
        # the communicator within the group
        cgroup = comm.comm_group
        # the communicator with all processes with
        # the same rank within their group
        crank = comm.comm_rank

        for obs in data.obs:
            tod = obs['tod']

            for det in tod.local_dets:

                # get the pixels and weights from the cache

                pixelsname = "{}_{}".format(self._pixels, det)
                weightsname = "{}_{}".format(self._weights, det)
                pixels = tod.cache.reference(pixelsname)
                weights = tod.cache.reference(weightsname)

                nsamp, nnz = weights.shape

                sm, lpix = self._map.global_to_local(pixels)

                #f = (np.dot(weights[x], self._map.data[sm[x], lpix[x]])
                #     if (lpix[x] >= 0) else 0
                #     for x in range(tod.local_samples[1]))
                #maptod = np.fromiter(f, np.float64, count=tod.local_samples[1])
                maptod = np.zeros(nsamp)
                sim_map_scan_map(sm, weights, lpix, self._map.data, maptod)

                cachename = "{}_{}".format(self._out, det)
                if not tod.cache.exists(cachename):
                    tod.cache.create(cachename, np.float64,
                                     (tod.local_samples[1],))
                ref = tod.cache.reference(cachename)
                ref[:] += maptod

                del ref
                del pixels
                del weights

        return

def extract_local_dets(data):
    """Extracts the local detectors from the TOD objects

    Some detectors could only appear in some observations, so we need
    to loop through all observations and accumulate all detectors in
    a set
    """
    local_dets = set()
    for obs in data.obs:
        tod = obs['tod']
        local_dets.update(tod.local_dets)
    return local_dets

def assemble_map_on_rank0(comm, local_map, n_components, npix):
    full_maps_rank0 = np.zeros((n_components, npix), dtype=np.float64) if comm.comm_world.rank == 0 else None
    comm.comm_world.Reduce(local_map, full_maps_rank0, root=0, op=MPI.SUM)
    return full_maps_rank0

def extract_detector_parameters(det, focalplanes):
    for fp in focalplanes:
        if det in fp:
            return fp[det]["bandcenter_ghz"], fp[det]["bandwidth_ghz"], fp[det]["fwhm"]
    raise RuntimeError("Cannot find detector {} in any focalplane")

class OpSimPySM(Operator):
    """
    Operator which generates sky signal by scanning from a map.

    The signal to use should already be in a distributed pixel structure,
    and local pointing should already exist.

    Args:
        distmap (DistPixels): the distributed map domain data.
        pixels (str): the name of the cache object (<pixels>_<detector>)
            containing the pixel indices to use.
        weights (str): the name of the cache object (<weights>_<detector>)
            containing the pointing weights to use.
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
    """
    def __init__(self, comm=None,
                 out='signal', pysm_model='', focalplanes=None, nside=None,
                 subnpix=None, localsm=None):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._out = out
        self.comm = comm
        self.dist_rings = DistRings(comm.comm_rank,
                            nside = nside,
                            nnz = 3)

        pysm_sky_components = [
            'synchrotron',
            'dust',
            'freefree',
            'cmb',
            'ame',
        ]
        pysm_sky_config = dict()
        for component_model in pysm_model.split(','):
            full_component_name = \
                    [each for each in pysm_sky_components
                     if each.startswith(component_model[0])][0]
            pysm_sky_config[full_component_name] = component_model
        self.pysm_sky = PySMSky(local_pixels=self.dist_rings.local_pixels,
                              nside=nside,
                              pysm_sky_config=pysm_sky_config)
        self.focalplanes = focalplanes
        self.npix = hp.nside2npix(nside)
        self.distmap = DistPixels(
            comm=comm.comm_world, size=self.npix, nnz=3,
            dtype=np.float32, submap=subnpix, local=localsm)

    def exec(self, data):
        local_dets = extract_local_dets(data)

        bandpasses = {}
        fwhm_deg = {}
        N_POINTS_BANDPASS = 10  # possibly take as parameter
        for det in local_dets:
            bandcenter, bandwidth, fwhm_deg[det] = \
                    extract_detector_parameters(det, self.focalplanes)
            bandpasses[det] = \
                (np.linspace(bandcenter-bandwidth/2,
                             bandcenter+bandwidth/2,
                             N_POINTS_BANDPASS),
                 np.ones(N_POINTS_BANDPASS))

        lmax = 512

        if self.comm.comm_world.rank == 0:
            print('Collecting, Broadcasting map', flush=True)
        start = MPI.Wtime()
        local_maps = dict()  # FIXME use Cache instead
        for det in local_dets:
            self.pysm_sky.exec(local_maps, out="sky",
                               bandpasses={"": bandpasses[det]})
            # LibSharp also supports transforming multiple channels together each with own beam
            smooth = LibSharpSmooth(self.comm.comm_rank, signal_map="sky",
                                       lmax=lmax, grid=self.dist_rings.libsharp_grid,
                                       fwhm_deg=fwhm_deg[det], beam=None)
            smooth.exec(local_maps)

            n_components = 3

            full_map_rank0 = assemble_map_on_rank0(self.comm,
                                     local_maps["sky"], n_components, self.npix)
            # full_map_rank0 dict contains on rank 0 the smoothed PySM map

            if self.comm.comm_world.rank == 0:
                print(full_map_rank0[:10])

            self.distmap.broadcast_healpix_map(full_map_rank0)
            scansim = OpSimScan(distmap=self.distmap, out=self._out+'_'+det)
            scansim.exec(data)

        stop = MPI.Wtime()
        if self.comm.comm_world.rank == 0:
            print('PySM Operator completed:  {:.2f} seconds'
                  ''.format(stop-start), flush=True)
<commit_msg>use rank communicator in assembling<commit_after># Copyright (c) 2015-2017 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

import numpy as np
import healpy as hp

from .. import qarray as qa
from .tod import TOD
from ..op import Operator
from ..map import DistRings, PySMSky, LibSharpSmooth, DistPixels
from ..mpi import MPI
from ..ctoast import sim_map_scan_map

class OpSimGradient(Operator):
    """
    Generate a fake sky signal as a gradient between the poles.

    This passes through each observation and creates a fake signal timestream
    based on the cartesian Z coordinate of the HEALPix pixel containing the
    detector pointing.

    Args:
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
        nside (int): the HEALPix NSIDE value to use.
        min (float): the minimum value to use at the South Pole.
        max (float): the maximum value to use at the North Pole.
        nest (bool): whether to use NESTED ordering.
    """

    def __init__(self, out='grad', nside=512, min=-100.0, max=100.0, nest=False,
                 flag_mask=255, common_flag_mask=255):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._nside = nside
        self._out = out
        self._min = min
        self._max = max
        self._nest = nest
        self._flag_mask = flag_mask
        self._common_flag_mask = common_flag_mask

    def exec(self, data):
        """
        Create the gradient timestreams.

        This pixelizes each detector's pointing and then assigns a
        timestream value based on the cartesian Z coordinate of the pixel
        center.

        Args:
            data (toast.Data): The distributed data.
        """
        comm = data.comm

        zaxis = np.array([0,0,1], dtype=np.float64)
        nullquat = np.array([0,0,0,1], dtype=np.float64)

        range = self._max - self._min

        for obs in data.obs:
            tod = obs['tod']
            base = obs['baselines']
            nse = obs['noise']
            intrvl = obs['intervals']

            for det in tod.local_dets:
                pdata = np.copy(tod.read_pntg(detector=det, local_start=0,
                                              n=tod.local_samples[1]))
                flags, common = tod.read_flags(detector=det, local_start=0,
                                               n=tod.local_samples[1])
                totflags = flags & self._flag_mask
                totflags |= (common & self._common_flag_mask)

                del flags
                del common

                pdata[totflags != 0,:] = nullquat

                dir = qa.rotate(pdata, zaxis)
                pixels = hp.vec2pix(self._nside, dir[:,0], dir[:,1], dir[:,2],
                                    nest=self._nest)
                x, y, z = hp.pix2vec(self._nside, pixels, nest=self._nest)
                z += 1.0
                z *= 0.5
                z *= range
                z += self._min
                z[totflags != 0] = 0.0

                cachename = "{}_{}".format(self._out, det)
                if not tod.cache.exists(cachename):
                    tod.cache.create(cachename, np.float64,
                                     (tod.local_samples[1],))
                ref = tod.cache.reference(cachename)
                ref[:] += z
                del ref
                #print('Grad timestream:', ref, np.sum(ref!=0),' non-zeros', flush=True) # DEBUG

        return

    def sigmap(self):
        """
        (array): Return the underlying signal map (full map on all processes).
        """
        range = self._max - self._min
        pix = np.arange(0, 12*self._nside*self._nside, dtype=np.int64)
        x, y, z = hp.pix2vec(self._nside, pix, nest=self._nest)
        z += 1.0
        z *= 0.5
        z *= range
        z += self._min
        return z

class OpSimScan(Operator):
    """
    Operator which generates sky signal by scanning from a map.

    The signal to use should already be in a distributed pixel structure,
    and local pointing should already exist.

    Args:
        distmap (DistPixels): the distributed map domain data.
        pixels (str): the name of the cache object (<pixels>_<detector>)
            containing the pixel indices to use.
        weights (str): the name of the cache object (<weights>_<detector>)
            containing the pointing weights to use.
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
    """
    def __init__(self, distmap=None, pixels='pixels', weights='weights',
                 out='scan'):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._map = distmap
        self._pixels = pixels
        self._weights = weights
        self._out = out

    def exec(self, data):
        """
        Create the timestreams by scanning from the map.

        This loops over all observations and detectors and uses the pointing
        matrix to project the distributed map into a timestream.

        Args:
            data (toast.Data): The distributed data.
        """
        comm = data.comm
        # the global communicator
        cworld = comm.comm_world
        # the communicator within the group
        cgroup = comm.comm_group
        # the communicator with all processes with
        # the same rank within their group
        crank = comm.comm_rank

        for obs in data.obs:
            tod = obs['tod']

            for det in tod.local_dets:

                # get the pixels and weights from the cache

                pixelsname = "{}_{}".format(self._pixels, det)
                weightsname = "{}_{}".format(self._weights, det)
                pixels = tod.cache.reference(pixelsname)
                weights = tod.cache.reference(weightsname)

                nsamp, nnz = weights.shape

                sm, lpix = self._map.global_to_local(pixels)

                #f = (np.dot(weights[x], self._map.data[sm[x], lpix[x]])
                #     if (lpix[x] >= 0) else 0
                #     for x in range(tod.local_samples[1]))
                #maptod = np.fromiter(f, np.float64, count=tod.local_samples[1])
                maptod = np.zeros(nsamp)
                sim_map_scan_map(sm, weights, lpix, self._map.data, maptod)

                cachename = "{}_{}".format(self._out, det)
                if not tod.cache.exists(cachename):
                    tod.cache.create(cachename, np.float64,
                                     (tod.local_samples[1],))
                ref = tod.cache.reference(cachename)
                ref[:] += maptod

                del ref
                del pixels
                del weights

        return

def extract_local_dets(data):
    """Extracts the local detectors from the TOD objects

    Some detectors could only appear in some observations, so we need
    to loop through all observations and accumulate all detectors in
    a set
    """
    local_dets = set()
    for obs in data.obs:
        tod = obs['tod']
        local_dets.update(tod.local_dets)
    return local_dets

def assemble_map_on_rank0(comm, local_map, n_components, npix):
    full_maps_rank0 = np.zeros((n_components, npix), dtype=np.float64) if comm.rank == 0 else None
    comm.Reduce(local_map, full_maps_rank0, root=0, op=MPI.SUM)
    return full_maps_rank0

def extract_detector_parameters(det, focalplanes):
    for fp in focalplanes:
        if det in fp:
            return fp[det]["bandcenter_ghz"], fp[det]["bandwidth_ghz"], fp[det]["fwhm"]
    raise RuntimeError("Cannot find detector {} in any focalplane")

class OpSimPySM(Operator):
    """
    Operator which generates sky signal by scanning from a map.

    The signal to use should already be in a distributed pixel structure,
    and local pointing should already exist.

    Args:
        distmap (DistPixels): the distributed map domain data.
        pixels (str): the name of the cache object (<pixels>_<detector>)
            containing the pixel indices to use.
        weights (str): the name of the cache object (<weights>_<detector>)
            containing the pointing weights to use.
        out (str): accumulate data to the cache with name <out>_<detector>.
            If the named cache objects do not exist, then they are created.
    """
    def __init__(self, comm=None,
                 out='signal', pysm_model='', focalplanes=None, nside=None,
                 subnpix=None, localsm=None):
        # We call the parent class constructor, which currently does nothing
        super().__init__()
        self._out = out
        self.comm = comm
        self.dist_rings = DistRings(comm.comm_rank,
                            nside = nside,
                            nnz = 3)

        pysm_sky_components = [
            'synchrotron',
            'dust',
            'freefree',
            'cmb',
            'ame',
        ]
        pysm_sky_config = dict()
        for component_model in pysm_model.split(','):
            full_component_name = \
                    [each for each in pysm_sky_components
                     if each.startswith(component_model[0])][0]
            pysm_sky_config[full_component_name] = component_model
        self.pysm_sky = PySMSky(local_pixels=self.dist_rings.local_pixels,
                              nside=nside,
                              pysm_sky_config=pysm_sky_config)
        self.focalplanes = focalplanes
        self.npix = hp.nside2npix(nside)
        self.distmap = DistPixels(
            comm=comm.comm_world, size=self.npix, nnz=3,
            dtype=np.float32, submap=subnpix, local=localsm)

    def exec(self, data):
        local_dets = extract_local_dets(data)

        bandpasses = {}
        fwhm_deg = {}
        N_POINTS_BANDPASS = 10  # possibly take as parameter
        for det in local_dets:
            bandcenter, bandwidth, fwhm_deg[det] = \
                    extract_detector_parameters(det, self.focalplanes)
            bandpasses[det] = \
                (np.linspace(bandcenter-bandwidth/2,
                             bandcenter+bandwidth/2,
                             N_POINTS_BANDPASS),
                 np.ones(N_POINTS_BANDPASS))

        lmax = 512

        if self.comm.comm_world.rank == 0:
            print('Collecting, Broadcasting map', flush=True)
        start = MPI.Wtime()
        local_maps = dict()  # FIXME use Cache instead
        for det in local_dets:
            self.pysm_sky.exec(local_maps, out="sky",
                               bandpasses={"": bandpasses[det]})
            # LibSharp also supports transforming multiple channels together each with own beam
            smooth = LibSharpSmooth(self.comm.comm_rank, signal_map="sky",
                                       lmax=lmax, grid=self.dist_rings.libsharp_grid,
                                       fwhm_deg=fwhm_deg[det], beam=None)
            smooth.exec(local_maps)

            n_components = 3

            full_map_rank0 = assemble_map_on_rank0(self.comm.comm_rank,
                                     local_maps["sky"], n_components, self.npix)
            # full_map_rank0 dict contains on rank 0 the smoothed PySM map

            if self.comm.comm_world.rank == 0:
                print(full_map_rank0[:10])

            self.distmap.broadcast_healpix_map(full_map_rank0)
            scansim = OpSimScan(distmap=self.distmap, out=self._out+'_'+det)
            scansim.exec(data)

        stop = MPI.Wtime()
        if self.comm.comm_world.rank == 0:
            print('PySM Operator completed:  {:.2f} seconds'
                  ''.format(stop-start), flush=True)
<|endoftext|>


<commit_before>from django.contrib import admin

# Register your models here.
<commit_msg>Admin poll model<commit_after>from django.contrib import admin
from polls.models import Poll

admin.site.register(Poll)
<|endoftext|>


<commit_before>from __future__ import unicode_literals
import base64
import datetime
import json
import logging
from urllib2 import Request, urlopen

from django.db.models import Max
from django.utils import safestring
import numpy as np
import re
from django.contrib.auth.models import User
from django.db import models
from django.utils import timezone
from pyMSpec.pyisocalc import pyisocalc

# Create your models here.
class Adduct(models.Model):
    nM = models.IntegerField(default=1)
    delta_formula = models.TextField(default="")  # addition/loss groups per M (sum of +X-Y)
    delta_atoms = models.TextField(default="")  # net atom addition/loss
    charge = models.IntegerField(default=1)  # overall charge after atom change

    def charge_str(self):
        if np.sign(self.charge) == 1:
            c_sign = "+"
        else:
            c_sign = "-"
        n_charges = np.abs(self.charge)
        if n_charges == 1:
            n_charges=""
        return "{}{}".format(n_charges,c_sign)

    def nice_str(self):
        _nM = self.nM
        _charge = self.charge
        if _nM == 1:
            _nM = ""
        if np.abs(_charge) == 1:
            _charge = self.charge_str()
        return "[{}M{}]{}".format(_nM, self.delta_formula, _charge)

    def html_str(self):
        return safestring.mark_safe("[{}M{}]{}".format(self.nM, self.delta_formula, self.charge_str()).replace("1", ""))

    def __unicode__(self):
        #!! don't edit this - I'm an idiot so it's used as a key in Molecule!!#
        return "[{}M{}]{}".format(self.nM, self.delta_formula, self.charge)

    def get_delta_atoms(self):
        def addElement(elDict, element, number):
            elDict.setdefault(element, []).append(number)

        self.delta_formula = self.delta_formula.strip()
        if all([self.delta_formula.startswith("+"), self.delta_formula.startswith("-")]):
            self.delta_formula = "+" + self.delta_formula
        formula_split = re.split(u'([+-])', self.delta_formula)
        logging.debug(formula_split)
        el_dict = {}
        for sign, el in zip(formula_split[1::2], formula_split[2::2]):
            this_el_dict = dict([(segment.element().name(), int("{}1".format(sign)) * segment.amount()) for segment in
                                 pyisocalc.parseSumFormula(el).get_segments()])
            for this_el in this_el_dict:
                logging.debug(el_dict)
                addElement(el_dict, this_el, this_el_dict[this_el])
        sign_dict = {1: "+", -1: "-"}
        for this_el in el_dict:
            el_dict[this_el] = sum(el_dict[this_el])
        logging.debug(el_dict)
        el_string = "".join(["{}{}{}".format(sign_dict[np.sign(el_dict[el])], el, abs(el_dict[el])) for el in el_dict if
                             el_dict[el] != 0])
        logging.debug(el_string)
        return el_string

    def save(self, *args, **kwargs):
        self.delta_atoms = self.get_delta_atoms()
        super(Adduct, self).save(*args, **kwargs)

class Molecule(models.Model):
    _adduct_mzs = models.TextField(default="")
    name = models.TextField(default="")
    sum_formula = models.TextField(null=True)
    inchi_code = models.TextField(default="")
    exact_mass = models.FloatField(default=0.0)
    solubility = models.TextField(null=True, blank=True)
    # External reference numbers
    hmdb_id = models.TextField(null=True, blank=True)
    chebi_id = models.TextField(null=True, blank=True)
    lipidmaps_id = models.TextField(null=True, blank=True)
    cas_id = models.TextField(null=True, blank=True)
    pubchem_id = models.TextField(null=True, blank=True)
    tags = models.ManyToManyField('MoleculeTag', blank=True)
    natural_product = models.BooleanField(default=True)

    def get_adduct_mzs(self):
        return json.loads(self._adduct_mzs)

    def set_adduct_mzs(self):
        adduct_dict = {}
        for adduct in Adduct.objects.all():
            adduct_dict[str(adduct)] = self.get_mz(adduct)
        self._adduct_mzs = json.dumps(adduct_dict)

    def get_adduct_mzs_by_pk(self):
        by_str = self.get_adduct_mzs()
        return {add.pk: by_str[str(add)] for add in Adduct.objects.all()}

    adduct_mzs = property(get_adduct_mzs, set_adduct_mzs)
    adduct_mzs_by_pk = property(get_adduct_mzs_by_pk)

    def get_mass(self):
        logging.info(self.sum_formula)
        logging.info(pyisocalc.parseSumFormula(self.sum_formula))
        spec = pyisocalc.perfect_pattern(pyisocalc.parseSumFormula(self.sum_formula), charge=0)
        logging.info(spec)
        mass = spec.get_spectrum(source='centroids')[0][np.argmax(spec.get_spectrum(source='centroids')[1])]
        logging.info(mass)
        return mass

    def __unicode__(self):
        return u"".join([i for i in self.name if ord(i) < 128])

    def html_str(self):
        return safestring.mark_safe("{}".format(self.name))

    def save(self, *args, **kwargs):
        logging.info('starting save')
        self.sum_formula = self.sum_formula.strip()
        logging.info('starting sf')
        self.exact_mass = self.get_mass()
        logging.info('starting adduct')
        self.set_adduct_mzs()
        logging.info('ready to save')
        super(Molecule, self).save(*args, **kwargs)

    def make_ion_formula(self, adduct):
        formula = "({}){}{}".format(self.sum_formula, adduct.nM, adduct.delta_atoms)
        return formula

    def get_mz(self, adduct):
        """
        Calculate the precursor mass for this molecule with a given adduct
        :param adduct: object of class Adduct
        :return: float
        """
        try:
            formula = self.make_ion_formula(adduct)
            spec = pyisocalc.perfect_pattern(pyisocalc.parseSumFormula(formula), charge=adduct.charge)
            mass = spec.get_spectrum(source='centroids')[0][np.argmax(spec.get_spectrum(source='centroids')[1])]
            return mass
        except:
            logging.debug(self.name, adduct)
            return -1.

    def spectra_count(self):
        return FragmentationSpectrum.objects.all().filter(standard__molecule=self).count()

    @property
    def smiles(self):
        import pybel
        return pybel.readstring('inchi', self.inchi_code).writestring('smi')

class Standard(models.Model):
    inventory_id = models.IntegerField(db_column='MCFID', unique=True)
    molecule = models.ForeignKey(Molecule)
    vendor = models.TextField(null=True, blank=True)
    vendor_cat = models.TextField(null=True, blank=True)
    lot_num = models.TextField(null=True, blank=True)
    location = models.TextField(null=True, blank=True)
    purchase_date = models.DateField(null=True, blank=True)

    def save(self, *args, **kwargs):
        if self.inventory_id is None:
            standards = Standard.objects.all()
            max_ = standards.aggregate(Max('inventory_id'))['inventory_id__max']
            if max_ is None:  # if there are no standards
                max_ = 0
            self.inventory_id = max_ + 1
        super(Standard, self).save(*args, **kwargs)

    def __unicode__(self):
        return "{}: {}".format(self.inventory_id, self.molecule.name)

class LcInfo(models.Model):
    content = models.TextField()

class MsInfo(models.Model):
    content = models.TextField()

class InstrumentInfo(models.Model):
    content = models.TextField()

class Dataset(models.Model):
    processing_finished = models.BooleanField(default=False)
    name = models.TextField(default="")
    path = models.TextField(default="")
    adducts_present = models.ManyToManyField(Adduct, blank=True)
    standards_present = models.ManyToManyField(Standard, blank=True)
    mass_accuracy_ppm = models.FloatField(default=10.0)
    quad_window_mz = models.FloatField(default=1.0)
    lc_info = models.ManyToManyField(to=LcInfo)
    ms_info = models.ManyToManyField(to=MsInfo)
    instrument_info = models.ManyToManyField(to=InstrumentInfo)
    ionization_method = models.TextField(default="")
    ion_analyzer = models.TextField(default="")

    # (for xic search)
    def __unicode__(self):
        return self.name

class Xic(models.Model):
    mz = models.FloatField(default=0.0)
    dataset = models.ForeignKey(Dataset)
    _xic = models.TextField(
        db_column='data',
        blank=True)
    _rt = models.TextField(
        db_column='rt_data',
        blank=True)

    standard = models.ForeignKey(Standard, blank=True, null=True)
    adduct = models.ForeignKey(Adduct, blank=True, null=True)
    collision = models.TextField(default='')

    def set_xic(self, xic):
        xic = np.asarray(xic, dtype=np.float64)
        self._xic = base64.b64encode(xic)

    def set_rt(self, rt):
        rt = np.asarray(rt, dtype=np.float64)
        self._rt = base64.b64encode(rt)

    def get_xic(self):
        r = base64.decodestring(self._xic)
        return np.frombuffer(r, dtype=np.float64)

    def get_rt(self):
        r = base64.decodestring(self._rt)
        return np.frombuffer(r, dtype=np.float64)

    xic = property(get_xic, set_xic)
    rt = property(get_rt, set_rt)

    def check_mass(self, tol_ppm=100):
        tol_mz = self.mz * tol_ppm * 1e-6
        theor_mz = self.standard.molecule.get_mz(self.adduct)
        if np.abs(theor_mz - self.mz) > tol_mz:
            raise ValueError('Mass tolerance not satisfied {} {}'.format(theor_mz, self.mz))
        return True
        # todo(An)
        # extend save to check that standard+adduct mass == precursor

class FragmentationSpectrum(models.Model):
    precursor_mz = models.FloatField(null=True)
    ms1_intensity = models.FloatField(default=0.0)
    _centroid_mzs = models.TextField()
    _centroid_ints = models.TextField()
    collision_energy = models.TextField(default="")
    dataset = models.ForeignKey(Dataset)
    standard = models.ForeignKey(Standard, blank=True, null=True)
    adduct = models.ForeignKey(Adduct, blank=True, null=True)
    spec_num = models.IntegerField(blank=True, null=True)
    rt = models.FloatField(blank=True, null=True)
    precursor_quad_fraction = models.FloatField(blank=True, null=True)
    reviewed = models.BooleanField(default=0)
    date_added = models.DateField(default=timezone.now)
    date_edited = models.DateField(default=timezone.now)
    last_editor = models.ForeignKey(User, blank=True, null=True)

    def __unicode__(self):
        return "{} {:3.2f}".format(self.spec_num, self.precursor_mz)

    def set_centroid_mzs(self, mzs):
        mzs = np.asarray(mzs, dtype=np.float64)
        self._centroid_mzs = base64.b64encode(mzs)

    def get_centroid_mzs(self):
        r = base64.decodestring(self._centroid_mzs)
        return np.frombuffer(r, dtype=np.float64)

    centroid_mzs = property(get_centroid_mzs, set_centroid_mzs)

    def set_centroid_ints(self, values):
        values = np.asarray(values, dtype=np.float64)
        self._centroid_ints = base64.b64encode(values)

    def get_centroid_ints(self):
        r = base64.decodestring(self._centroid_ints)
        return np.frombuffer(r, dtype=np.float64)

    centroid_ints = property(get_centroid_ints, set_centroid_ints)

    def get_centroids(self):
        return self.centroid_mzs, self.centroid_ints

    def save(self, *args, **kwargs):
        logging.debug(self.pk)
        if not self.pk:
            self.date_added = datetime.datetime.now()
        self.date_edited = datetime.datetime.now()
        super(FragmentationSpectrum, self).save(*args, **kwargs)

    @property
    def base_peak(self):
        spec = self.get_centroids()
        return spec[0][np.argmax(spec[1])]

    @property
    def splash(self):
        splash_payload = json.dumps({
            "ions": [{"mass": mz, "intensity": int_} for mz, int_ in zip(self.centroid_mzs, self.centroid_ints)],
            "type": "MS"})
        url = "http://splash.fiehnlab.ucdavis.edu/splash/it"
        request = Request(url, data=splash_payload, headers={'Content-Type': "application/json"})
        response = urlopen(request).read().decode()
        return response

class MoleculeSpectraCount(models.Model):
    molecule = models.ForeignKey(Molecule, primary_key=True, on_delete=models.DO_NOTHING)
    spectra_count = models.IntegerField()

    class Meta:
        managed = False

class ProcessingError(models.Model):
    dataset = models.ForeignKey(Dataset)
    message = models.TextField()

class MoleculeTag(models.Model):
    name = models.TextField()

    def __unicode__(self):
        return self.name
<commit_msg>Fix SMILES conversion<commit_after>from __future__ import unicode_literals
import base64
import datetime
import json
import logging
from urllib2 import Request, urlopen

from django.db.models import Max
from django.utils import safestring
import numpy as np
import re
from django.contrib.auth.models import User
from django.db import models
from django.utils import timezone
from pyMSpec.pyisocalc import pyisocalc

# Create your models here.
class Adduct(models.Model):
    nM = models.IntegerField(default=1)
    delta_formula = models.TextField(default="")  # addition/loss groups per M (sum of +X-Y)
    delta_atoms = models.TextField(default="")  # net atom addition/loss
    charge = models.IntegerField(default=1)  # overall charge after atom change

    def charge_str(self):
        if np.sign(self.charge) == 1:
            c_sign = "+"
        else:
            c_sign = "-"
        n_charges = np.abs(self.charge)
        if n_charges == 1:
            n_charges=""
        return "{}{}".format(n_charges,c_sign)

    def nice_str(self):
        _nM = self.nM
        _charge = self.charge
        if _nM == 1:
            _nM = ""
        if np.abs(_charge) == 1:
            _charge = self.charge_str()
        return "[{}M{}]{}".format(_nM, self.delta_formula, _charge)

    def html_str(self):
        return safestring.mark_safe("[{}M{}]{}".format(self.nM, self.delta_formula, self.charge_str()).replace("1", ""))

    def __unicode__(self):
        #!! don't edit this - I'm an idiot so it's used as a key in Molecule!!#
        return "[{}M{}]{}".format(self.nM, self.delta_formula, self.charge)

    def get_delta_atoms(self):
        def addElement(elDict, element, number):
            elDict.setdefault(element, []).append(number)

        self.delta_formula = self.delta_formula.strip()
        if all([self.delta_formula.startswith("+"), self.delta_formula.startswith("-")]):
            self.delta_formula = "+" + self.delta_formula
        formula_split = re.split(u'([+-])', self.delta_formula)
        logging.debug(formula_split)
        el_dict = {}
        for sign, el in zip(formula_split[1::2], formula_split[2::2]):
            this_el_dict = dict([(segment.element().name(), int("{}1".format(sign)) * segment.amount()) for segment in
                                 pyisocalc.parseSumFormula(el).get_segments()])
            for this_el in this_el_dict:
                logging.debug(el_dict)
                addElement(el_dict, this_el, this_el_dict[this_el])
        sign_dict = {1: "+", -1: "-"}
        for this_el in el_dict:
            el_dict[this_el] = sum(el_dict[this_el])
        logging.debug(el_dict)
        el_string = "".join(["{}{}{}".format(sign_dict[np.sign(el_dict[el])], el, abs(el_dict[el])) for el in el_dict if
                             el_dict[el] != 0])
        logging.debug(el_string)
        return el_string

    def save(self, *args, **kwargs):
        self.delta_atoms = self.get_delta_atoms()
        super(Adduct, self).save(*args, **kwargs)

class Molecule(models.Model):
    _adduct_mzs = models.TextField(default="")
    name = models.TextField(default="")
    sum_formula = models.TextField(null=True)
    inchi_code = models.TextField(default="")
    exact_mass = models.FloatField(default=0.0)
    solubility = models.TextField(null=True, blank=True)
    # External reference numbers
    hmdb_id = models.TextField(null=True, blank=True)
    chebi_id = models.TextField(null=True, blank=True)
    lipidmaps_id = models.TextField(null=True, blank=True)
    cas_id = models.TextField(null=True, blank=True)
    pubchem_id = models.TextField(null=True, blank=True)
    tags = models.ManyToManyField('MoleculeTag', blank=True)
    natural_product = models.BooleanField(default=True)

    def get_adduct_mzs(self):
        return json.loads(self._adduct_mzs)

    def set_adduct_mzs(self):
        adduct_dict = {}
        for adduct in Adduct.objects.all():
            adduct_dict[str(adduct)] = self.get_mz(adduct)
        self._adduct_mzs = json.dumps(adduct_dict)

    def get_adduct_mzs_by_pk(self):
        by_str = self.get_adduct_mzs()
        return {add.pk: by_str[str(add)] for add in Adduct.objects.all()}

    adduct_mzs = property(get_adduct_mzs, set_adduct_mzs)
    adduct_mzs_by_pk = property(get_adduct_mzs_by_pk)

    def get_mass(self):
        logging.info(self.sum_formula)
        logging.info(pyisocalc.parseSumFormula(self.sum_formula))
        spec = pyisocalc.perfect_pattern(pyisocalc.parseSumFormula(self.sum_formula), charge=0)
        logging.info(spec)
        mass = spec.get_spectrum(source='centroids')[0][np.argmax(spec.get_spectrum(source='centroids')[1])]
        logging.info(mass)
        return mass

    def __unicode__(self):
        return u"".join([i for i in self.name if ord(i) < 128])

    def html_str(self):
        return safestring.mark_safe("{}".format(self.name))

    def save(self, *args, **kwargs):
        logging.info('starting save')
        self.sum_formula = self.sum_formula.strip()
        logging.info('starting sf')
        self.exact_mass = self.get_mass()
        logging.info('starting adduct')
        self.set_adduct_mzs()
        logging.info('ready to save')
        super(Molecule, self).save(*args, **kwargs)

    def make_ion_formula(self, adduct):
        formula = "({}){}{}".format(self.sum_formula, adduct.nM, adduct.delta_atoms)
        return formula

    def get_mz(self, adduct):
        """
        Calculate the precursor mass for this molecule with a given adduct
        :param adduct: object of class Adduct
        :return: float
        """
        try:
            formula = self.make_ion_formula(adduct)
            spec = pyisocalc.perfect_pattern(pyisocalc.parseSumFormula(formula), charge=adduct.charge)
            mass = spec.get_spectrum(source='centroids')[0][np.argmax(spec.get_spectrum(source='centroids')[1])]
            return mass
        except:
            logging.debug(self.name, adduct)
            return -1.

    def spectra_count(self):
        return FragmentationSpectrum.objects.all().filter(standard__molecule=self).count()

    @property
    def smiles(self):
        import pybel
        try:
            return pybel.readstring(b'inchi', self.inchi_code.encode('ascii')).write(b'smi').strip()
        except IOError:
            logging.error('Could not read InChI code: {}'.format(self.inchi_code))
            return '??'

class Standard(models.Model):
    inventory_id = models.IntegerField(db_column='MCFID', unique=True)
    molecule = models.ForeignKey(Molecule)
    vendor = models.TextField(null=True, blank=True)
    vendor_cat = models.TextField(null=True, blank=True)
    lot_num = models.TextField(null=True, blank=True)
    location = models.TextField(null=True, blank=True)
    purchase_date = models.DateField(null=True, blank=True)

    def save(self, *args, **kwargs):
        if self.inventory_id is None:
            standards = Standard.objects.all()
            max_ = standards.aggregate(Max('inventory_id'))['inventory_id__max']
            if max_ is None:  # if there are no standards
                max_ = 0
            self.inventory_id = max_ + 1
        super(Standard, self).save(*args, **kwargs)

    def __unicode__(self):
        return "{}: {}".format(self.inventory_id, self.molecule.name)

class LcInfo(models.Model):
    content = models.TextField()

class MsInfo(models.Model):
    content = models.TextField()

class InstrumentInfo(models.Model):
    content = models.TextField()

class Dataset(models.Model):
    processing_finished = models.BooleanField(default=False)
    name = models.TextField(default="")
    path = models.TextField(default="")
    adducts_present = models.ManyToManyField(Adduct, blank=True)
    standards_present = models.ManyToManyField(Standard, blank=True)
    mass_accuracy_ppm = models.FloatField(default=10.0)
    quad_window_mz = models.FloatField(default=1.0)
    lc_info = models.ManyToManyField(to=LcInfo)
    ms_info = models.ManyToManyField(to=MsInfo)
    instrument_info = models.ManyToManyField(to=InstrumentInfo)
    ionization_method = models.TextField(default="")
    ion_analyzer = models.TextField(default="")

    # (for xic search)
    def __unicode__(self):
        return self.name

class Xic(models.Model):
    mz = models.FloatField(default=0.0)
    dataset = models.ForeignKey(Dataset)
    _xic = models.TextField(
        db_column='data',
        blank=True)
    _rt = models.TextField(
        db_column='rt_data',
        blank=True)

    standard = models.ForeignKey(Standard, blank=True, null=True)
    adduct = models.ForeignKey(Adduct, blank=True, null=True)
    collision = models.TextField(default='')

    def set_xic(self, xic):
        xic = np.asarray(xic, dtype=np.float64)
        self._xic = base64.b64encode(xic)

    def set_rt(self, rt):
        rt = np.asarray(rt, dtype=np.float64)
        self._rt = base64.b64encode(rt)

    def get_xic(self):
        r = base64.decodestring(self._xic)
        return np.frombuffer(r, dtype=np.float64)

    def get_rt(self):
        r = base64.decodestring(self._rt)
        return np.frombuffer(r, dtype=np.float64)

    xic = property(get_xic, set_xic)
    rt = property(get_rt, set_rt)

    def check_mass(self, tol_ppm=100):
        tol_mz = self.mz * tol_ppm * 1e-6
        theor_mz = self.standard.molecule.get_mz(self.adduct)
        if np.abs(theor_mz - self.mz) > tol_mz:
            raise ValueError('Mass tolerance not satisfied {} {}'.format(theor_mz, self.mz))
        return True
        # todo(An)
        # extend save to check that standard+adduct mass == precursor

class FragmentationSpectrum(models.Model):
    precursor_mz = models.FloatField(null=True)
    ms1_intensity = models.FloatField(default=0.0)
    _centroid_mzs = models.TextField()
    _centroid_ints = models.TextField()
    collision_energy = models.TextField(default="")
    dataset = models.ForeignKey(Dataset)
    standard = models.ForeignKey(Standard, blank=True, null=True)
    adduct = models.ForeignKey(Adduct, blank=True, null=True)
    spec_num = models.IntegerField(blank=True, null=True)
    rt = models.FloatField(blank=True, null=True)
    precursor_quad_fraction = models.FloatField(blank=True, null=True)
    reviewed = models.BooleanField(default=0)
    date_added = models.DateField(default=timezone.now)
    date_edited = models.DateField(default=timezone.now)
    last_editor = models.ForeignKey(User, blank=True, null=True)

    def __unicode__(self):
        return "{} {:3.2f}".format(self.spec_num, self.precursor_mz)

    def set_centroid_mzs(self, mzs):
        mzs = np.asarray(mzs, dtype=np.float64)
        self._centroid_mzs = base64.b64encode(mzs)

    def get_centroid_mzs(self):
        r = base64.decodestring(self._centroid_mzs)
        return np.frombuffer(r, dtype=np.float64)

    centroid_mzs = property(get_centroid_mzs, set_centroid_mzs)

    def set_centroid_ints(self, values):
        values = np.asarray(values, dtype=np.float64)
        self._centroid_ints = base64.b64encode(values)

    def get_centroid_ints(self):
        r = base64.decodestring(self._centroid_ints)
        return np.frombuffer(r, dtype=np.float64)

    centroid_ints = property(get_centroid_ints, set_centroid_ints)

    def get_centroids(self):
        return self.centroid_mzs, self.centroid_ints

    def save(self, *args, **kwargs):
        logging.debug(self.pk)
        if not self.pk:
            self.date_added = datetime.datetime.now()
        self.date_edited = datetime.datetime.now()
        super(FragmentationSpectrum, self).save(*args, **kwargs)

    @property
    def base_peak(self):
        spec = self.get_centroids()
        return spec[0][np.argmax(spec[1])]

    @property
    def splash(self):
        splash_payload = json.dumps({
            "ions": [{"mass": mz, "intensity": int_} for mz, int_ in zip(self.centroid_mzs, self.centroid_ints)],
            "type": "MS"})
        url = "http://splash.fiehnlab.ucdavis.edu/splash/it"
        request = Request(url, data=splash_payload, headers={'Content-Type': "application/json"})
        response = urlopen(request).read().decode()
        return response

class MoleculeSpectraCount(models.Model):
    molecule = models.ForeignKey(Molecule, primary_key=True, on_delete=models.DO_NOTHING)
    spectra_count = models.IntegerField()

    class Meta:
        managed = False

class ProcessingError(models.Model):
    dataset = models.ForeignKey(Dataset)
    message = models.TextField()

class MoleculeTag(models.Model):
    name = models.TextField()

    def __unicode__(self):
        return self.name
<|endoftext|>


<commit_before>import requests
import json
import re
import fileinput

import argparse
import subprocess
import signal
from threading import Timer

import os
import sys
import time
import traceback

from couchbase.bucket import Bucket
import couchbase
from couchbase.n1ql import N1QLQuery

"""

# An evolving thing - takes as input:
- a file which is the output from perfrunner - this file will contain some json which describes the perf results
- the perf keys and expected values

This program parses out the results from the files and compares them against the expected values

"""

test_workload_output = '''
[20/Oct/2015 15:01:26] INFO - Creating new database: iostatperfregression_410-4859-enterprise_27b10170106
[20/Oct/2015 15:01:29] INFO - Creating new database: ns_serverperfregression_410-4859-enterprise_27bbucket-110170107
[20/Oct/2015 15:02:08] INFO - Adding snapshot: perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:02:15] INFO - http://cbmonitor.sc.couchbase.com/reports/html/?snapshot=perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:03:04] INFO - http://cbmonitor.sc.couchbase.com/reports/get_corr_matrix/?snapshot=perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:03:31] INFO - Dry run stats: {
    "build": "4.1.0-4859-enterprise",
    "build_url": null,
    "metric": "perf_sanity_kv_latency_mixed_2M_short_get_95th_perf_sanity_base_test",
    "snapshots": [
        "perfregression_410-4859-enterprise_27b_access"
    ],
    "value": 0.56
}
[20/Oct/2015 15:03:31] INFO - Dry run stats: {
    "build": "4.1.0-4859-enterprise",
    "build_url": null,
    "metric": "perf_sanity_kv_latency_mixed_2M_short_set_95th_perf_sanity_base_test",
    "snapshots": [
        "perfregression_410-4859-enterprise_27b_access"
    ],
    "value": 0.95
}
[20/Oct/2015 15:03:31] INFO - Terminating local Celery workers
'''

# subprocess timeout code from here http://stackoverflow.com/questions/1191374/using-module-subprocess-with-timeout

def kill_proc(proc, timeout):
  timeout["value"] = True
  proc.kill()

def run_with_timeout(cmd, env, timeout_sec):
  print 'the timeout value is', timeout_sec
  proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  timeout = {"value": False}
  timer = Timer(timeout_sec, kill_proc, [proc, timeout])
  timer.start()

  # real time output is handy but it didn't seem to work with the timeout stuff
  #for line in iter(proc.stdout.readline, ''):
        #print 'workload line', line
        #workload_output += line

  print 'before communicate'
  stdout, stderr = proc.communicate()
  print 'after communicate'
  timer.cancel()
  return proc.returncode, stdout.decode("utf-8"), stderr.decode("utf-8"), timeout["value"]

def checkResults( results, testDescriptor):
            #print '\n\nthe results are', results
            p = re.compile(r'Dry run stats: {(.*?)}', re.MULTILINE)
            matches = p.findall(results.replace('\n', ''))
            results = []
            actual_values = {}
            for m in matches:
                #print '\n\nhave a match', m
                actual = json.loads('{' + m + '}')
                actual_values[actual['metric']] = actual  # ( json.loads('{' + m + '}') )

            expected_keys = testDescriptor['KPIs']
            for k in expected_keys.keys():
                haveAMatch = False
                for i in actual_values.keys():
                    if k in i:
                        haveAMatch = True
                        actualIndex = i
                        break

                if haveAMatch:
                    passResult = True
                    if actual_values[actualIndex]['value'] > 1.1 * expected_keys[k]:
                        passResult = False
                        print '  ', k, ' is greater than expected. Expected', expected_keys[k], 'Actual', actual_values[actualIndex][
                            'value']

                    elif actual_values[actualIndex]['value'] < 0.9 * expected_keys[k]:
                        passResult = False
                        # sort of want to yellow flag this but for now all we have is a red flag so use that
                        print '  ', k, ' is less than expected. Expected', expected_keys[k], 'Actual', actual_values[actualIndex][
                            'value']

                    results.append({'testMetric':k, 'expectedValue':expected_keys[k], 'actualValue':actual_values[actualIndex]['value'], 'pass':passResult})
                    del actual_values[actualIndex]
                else:
                    print '  Expected key', k, ' is not found'

            if len(actual_values) > 0:
                print '  The following key(s) were present but not expected:'
                for i in actual_values:
                    print '    ', i
            return results

def runPerfRunner( testDescriptor, options):
    print testDescriptor['testType']
    testName = testDescriptor['testName']

    test = testDescriptor['testFile'] + '.test'
    if options.specFile is None:
        # backup and restore needs a list of spec files to do the installs
        if type(testDescriptor['specFile']) is unicode:
            spec = [ testDescriptor['specFile']]
        else:
            print 'have a list of spec files'
            spec = []
            for i in testDescriptor['specFile']:
                spec.append( i )
    else:
        spec = [options.specFile]
    print 'specfile', spec
    KPIs = testDescriptor['KPIs']

    my_env = os.environ
    my_env['test_config'] = 'perfSanity/tests/' + test
    if options.url is None:
        my_env['version'] = options.version
    else:
        my_env['url'] = options.url

    for i in spec:
        my_env['cluster'] = 'perfSanity/clusters/' + i + '.spec'
        if re.search('n1ql.*Q[2].*',testName):
            print '-'*100
            print 'Skipping Setup for N1QL Q2 queries ... '
            print '-'*100
        else:
            proc = subprocess.Popen('./scripts/setup.sh', env=my_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                        shell=True)

            for line in iter(proc.stdout.readline, ''):
                print 'Setup output', line
                sys.stdout.flush()

            (stdoutdata, stderrdata) = proc.communicate()

            if proc.returncode == 1:
                print '\n\nHave an error during setup'
                print stderrdata
                print stdoutdata
                return  [{'pass':False, 'reason':'Have an error during setup'}]

            print 'Setup complete, starting workload'

    rc, stdout, stderr, timeout = run_with_timeout( './perfSanity/scripts/workload_dev.sh', my_env, 7200)  # 2 hours
    print 'rc is', rc
    print 'timeout is', timeout
    print 'stderr', stderr
    print 'stdout is', stdout

    if proc.returncode == 1:
        print '  Have an error during workload generation'
        return [{'pass':False, 'reason':'Have an error during workload generation'}]
    elif timeout:
        print '  test timeout'
        return [{'pass':False, 'reason':'test timed out'}]
    else:
        print '\n\nWorkload complete, analyzing results'
        return checkResults( stdout, testDescriptor)

def runForestDBTest( testDescriptor, options):

    if options.url is not None:
        print 'runForestDBTest and url option is not supported'
        return []

    testName = testDescriptor['testName']

    command = testDescriptor['command'] + ' --version=' + options.version
    print 'the command is', command

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,shell=True)
    commandOutput = ''
    for line in iter(proc.stdout.readline, ''):
        print line
        commandOutput += line

    (stdoutdata, stderrdata) = proc.communicate()

    print 'stderrdata', stderrdata

    if proc.returncode == 1:
        print '  Have an error during forest DB'
        return [{'pass':False, 'reason':'Check logs'}]
    else:
        return checkResults( commandOutput, testDescriptor)

def runTest( testDescriptor, options, bucket ):
    print testDescriptor['testType']
    testName = testDescriptor['testName']

    testStartTime = time.strftime("%y-%m-%d-%H:%M:%S", time.strptime(time.ctime() ))
    startTime = time.time()   # in seconds to get the elapsed time
    print '\n\n', time.asctime( time.localtime(time.time()) ), 'Now running', testName

    baseResult = {'runStartTime':options.runStartTime, 'testStartTime':testStartTime, 'build':options.version, 'testName':testName}
    if testDescriptor['testType'] == 'perfRunner':
        res = runPerfRunner(testDescriptor, options)
    elif testDescriptor['testType'] == 'perfRunnerForestDB':
        print 'have the forest DB test', testDescriptor['command']
        res = runForestDBTest(testDescriptor, options)
    elif testDescriptor['testType'] == 'SampleDBs':
        return True
    else:
        print 'Unknown test type', testDescriptor['testType']
        return True

    for i in res:
        combinedResults = dict(baseResult.items() + i.items()  + {'elapsedTime': round(time.time() - startTime,0)}.items() )
        print 'the result is ', combinedResults
        if bucket is not None:
            if 'testMetric' in combinedResults:
                testKey = combinedResults['testName'] + '-' + combinedResults['testMetric']
            else:
                testKey = combinedResults['testName']
            bucket.upsert( testStartTime + '-' + options.version + '-' + testKey, combinedResults, format=couchbase.FMT_JSON)

def main():
    print 'Starting the perf regression runner'

    usage = '%prog -f conf-file'
    #parser = OptionParser(usage)

    ##parser.add_option('-f', '--filename', dest='filename')
    #parser.add_option('-v', '--version', dest='version')
    #parser.add_option('-u', '--url', dest='url')
    #parser.add_option('-q', '--query', dest='query')
    #parser.add_option('-s', '--specFile', dest='specFile')
    #parser.add_option('-r', '--runStartTime', dest='runStartTime')
    #parser.add_option('-b', '--betaTests', dest='betaTests', default=False, action='store_true')
    #parser.add_option('-n', '--nop', dest='nop',default=False, action='store_true')

    #options, args = parser.parse_args()

    # the option parsing way
    parser = argparse.ArgumentParser(description=usage)
    parser.add_argument('-q', '--query', nargs='+')
    parser.add_argument('-v', '--version', dest='version')
    parser.add_argument('-u', '--url', dest='url')
    parser.add_argument('-s', '--specFile', dest='specFile')
    parser.add_argument('-r', '--runStartTime', dest='runStartTime')
    parser.add_argument('-b', '--betaTests', dest='betaTests', default=False, action='store_true')
    parser.add_argument('-a', '--allTests', dest='allTests', default=False, action='store_true')
    parser.add_argument('-n', '--nop', dest='nop',default=False, action='store_true')

    options = parser.parse_args()

    print 'query', options.query
    print 'specfile', options.specFile

    runStartTime = options.runStartTime
    summary = []

    print 'version', options.version
    print 'url', options.url

    # open the bucket
    if options.nop:
        bucket = None
    else:
        bucket = Bucket('couchbase://'+ '172.23.105.177:8091/Daily-Performance')

    testBucket = Bucket('couchbase://'+ '172.23.105.177:8091/Daily-Performance-Tests')
    queryString = "select `Daily-Performance-Tests`.* from `Daily-Performance-Tests`"
    wherePredicates = []
    if options.query is not None:
        wherePredicates.append( ' '.join(options.query[0] ) )
    if not options.allTests:
        if options.betaTests:
            wherePredicates.append( "status='beta'")
        else:
            wherePredicates.append( "status!='beta'")

    if len(wherePredicates) > 0:
       for i in range(len(wherePredicates)):
          if i == 0:
              queryString += ' where '
          else:
              queryString += ' and '
          queryString += wherePredicates[i]

    print 'the query string is', queryString
    query = N1QLQuery(queryString )
    testsToRun = testBucket.n1ql_query( queryString )
    tests = [row for row in testsToRun]
    print 'the tests are', tests
    testsToRerun = []

    for row in tests:
        try:
            if row['status'].lower() == 'disabled':
                print row['testName'], ' is disabled.'
            else:
                if not runTest( row, options, bucket ):
                    testsToRerun.append(row)
        except:
            print 'Exception in ', row['testName']
            traceback.print_exc()

        #time.sleep(10)
    # end the for loop - print the results
    print 'done'

if __name__ == "__main__":
    if not main():
        sys.exit(1)

<commit_msg>CBQE-0: Allow parameters with spaces<commit_after>import requests
import json
import re
import fileinput

import argparse
import subprocess
import signal
from threading import Timer

import os
import sys
import time
import traceback

from couchbase.bucket import Bucket
import couchbase
from couchbase.n1ql import N1QLQuery

"""

# An evolving thing - takes as input:
- a file which is the output from perfrunner - this file will contain some json which describes the perf results
- the perf keys and expected values

This program parses out the results from the files and compares them against the expected values

"""

test_workload_output = '''
[20/Oct/2015 15:01:26] INFO - Creating new database: iostatperfregression_410-4859-enterprise_27b10170106
[20/Oct/2015 15:01:29] INFO - Creating new database: ns_serverperfregression_410-4859-enterprise_27bbucket-110170107
[20/Oct/2015 15:02:08] INFO - Adding snapshot: perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:02:15] INFO - http://cbmonitor.sc.couchbase.com/reports/html/?snapshot=perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:03:04] INFO - http://cbmonitor.sc.couchbase.com/reports/get_corr_matrix/?snapshot=perfregression_410-4859-enterprise_27b_access
[20/Oct/2015 15:03:31] INFO - Dry run stats: {
    "build": "4.1.0-4859-enterprise",
    "build_url": null,
    "metric": "perf_sanity_kv_latency_mixed_2M_short_get_95th_perf_sanity_base_test",
    "snapshots": [
        "perfregression_410-4859-enterprise_27b_access"
    ],
    "value": 0.56
}
[20/Oct/2015 15:03:31] INFO - Dry run stats: {
    "build": "4.1.0-4859-enterprise",
    "build_url": null,
    "metric": "perf_sanity_kv_latency_mixed_2M_short_set_95th_perf_sanity_base_test",
    "snapshots": [
        "perfregression_410-4859-enterprise_27b_access"
    ],
    "value": 0.95
}
[20/Oct/2015 15:03:31] INFO - Terminating local Celery workers
'''

# subprocess timeout code from here http://stackoverflow.com/questions/1191374/using-module-subprocess-with-timeout

def kill_proc(proc, timeout):
  timeout["value"] = True
  proc.kill()

def run_with_timeout(cmd, env, timeout_sec):
  print 'the timeout value is', timeout_sec
  proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  timeout = {"value": False}
  timer = Timer(timeout_sec, kill_proc, [proc, timeout])
  timer.start()

  # real time output is handy but it didn't seem to work with the timeout stuff
  #for line in iter(proc.stdout.readline, ''):
        #print 'workload line', line
        #workload_output += line

  print 'before communicate'
  stdout, stderr = proc.communicate()
  print 'after communicate'
  timer.cancel()
  return proc.returncode, stdout.decode("utf-8"), stderr.decode("utf-8"), timeout["value"]

def checkResults( results, testDescriptor):
            #print '\n\nthe results are', results
            p = re.compile(r'Dry run stats: {(.*?)}', re.MULTILINE)
            matches = p.findall(results.replace('\n', ''))
            results = []
            actual_values = {}
            for m in matches:
                #print '\n\nhave a match', m
                actual = json.loads('{' + m + '}')
                actual_values[actual['metric']] = actual  # ( json.loads('{' + m + '}') )

            expected_keys = testDescriptor['KPIs']
            for k in expected_keys.keys():
                haveAMatch = False
                for i in actual_values.keys():
                    if k in i:
                        haveAMatch = True
                        actualIndex = i
                        break

                if haveAMatch:
                    passResult = True
                    if actual_values[actualIndex]['value'] > 1.1 * expected_keys[k]:
                        passResult = False
                        print '  ', k, ' is greater than expected. Expected', expected_keys[k], 'Actual', actual_values[actualIndex][
                            'value']

                    elif actual_values[actualIndex]['value'] < 0.9 * expected_keys[k]:
                        passResult = False
                        # sort of want to yellow flag this but for now all we have is a red flag so use that
                        print '  ', k, ' is less than expected. Expected', expected_keys[k], 'Actual', actual_values[actualIndex][
                            'value']

                    results.append({'testMetric':k, 'expectedValue':expected_keys[k], 'actualValue':actual_values[actualIndex]['value'], 'pass':passResult})
                    del actual_values[actualIndex]
                else:
                    print '  Expected key', k, ' is not found'

            if len(actual_values) > 0:
                print '  The following key(s) were present but not expected:'
                for i in actual_values:
                    print '    ', i
            return results

def runPerfRunner( testDescriptor, options):
    print testDescriptor['testType']
    testName = testDescriptor['testName']

    test = testDescriptor['testFile'] + '.test'
    if options.specFile is None:
        # backup and restore needs a list of spec files to do the installs
        if type(testDescriptor['specFile']) is unicode:
            spec = [ testDescriptor['specFile']]
        else:
            print 'have a list of spec files'
            spec = []
            for i in testDescriptor['specFile']:
                spec.append( i )
    else:
        spec = [options.specFile]
    print 'specfile', spec
    KPIs = testDescriptor['KPIs']

    my_env = os.environ
    my_env['test_config'] = 'perfSanity/tests/' + test
    if options.url is None:
        my_env['version'] = options.version
    else:
        my_env['url'] = options.url

    for i in spec:
        my_env['cluster'] = 'perfSanity/clusters/' + i + '.spec'
        if re.search('n1ql.*Q[2].*',testName):
            print '-'*100
            print 'Skipping Setup for N1QL Q2 queries ... '
            print '-'*100
        else:
            proc = subprocess.Popen('./scripts/setup.sh', env=my_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                        shell=True)

            for line in iter(proc.stdout.readline, ''):
                print 'Setup output', line
                sys.stdout.flush()

            (stdoutdata, stderrdata) = proc.communicate()

            if proc.returncode == 1:
                print '\n\nHave an error during setup'
                print stderrdata
                print stdoutdata
                return  [{'pass':False, 'reason':'Have an error during setup'}]

            print 'Setup complete, starting workload'

    rc, stdout, stderr, timeout = run_with_timeout( './perfSanity/scripts/workload_dev.sh', my_env, 7200)  # 2 hours
    print 'rc is', rc
    print 'timeout is', timeout
    print 'stderr', stderr
    print 'stdout is', stdout

    if proc.returncode == 1:
        print '  Have an error during workload generation'
        return [{'pass':False, 'reason':'Have an error during workload generation'}]
    elif timeout:
        print '  test timeout'
        return [{'pass':False, 'reason':'test timed out'}]
    else:
        print '\n\nWorkload complete, analyzing results'
        return checkResults( stdout, testDescriptor)

def runForestDBTest( testDescriptor, options):

    if options.url is not None:
        print 'runForestDBTest and url option is not supported'
        return []

    testName = testDescriptor['testName']

    command = testDescriptor['command'] + ' --version=' + options.version
    print 'the command is', command

    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,shell=True)
    commandOutput = ''
    for line in iter(proc.stdout.readline, ''):
        print line
        commandOutput += line

    (stdoutdata, stderrdata) = proc.communicate()

    print 'stderrdata', stderrdata

    if proc.returncode == 1:
        print '  Have an error during forest DB'
        return [{'pass':False, 'reason':'Check logs'}]
    else:
        return checkResults( commandOutput, testDescriptor)

def runTest( testDescriptor, options, bucket ):
    print testDescriptor['testType']
    testName = testDescriptor['testName']

    testStartTime = time.strftime("%y-%m-%d-%H:%M:%S", time.strptime(time.ctime() ))
    startTime = time.time()   # in seconds to get the elapsed time
    print '\n\n', time.asctime( time.localtime(time.time()) ), 'Now running', testName

    baseResult = {'runStartTime':options.runStartTime, 'testStartTime':testStartTime, 'build':options.version, 'testName':testName}
    if testDescriptor['testType'] == 'perfRunner':
        res = runPerfRunner(testDescriptor, options)
    elif testDescriptor['testType'] == 'perfRunnerForestDB':
        print 'have the forest DB test', testDescriptor['command']
        res = runForestDBTest(testDescriptor, options)
    elif testDescriptor['testType'] == 'SampleDBs':
        return True
    else:
        print 'Unknown test type', testDescriptor['testType']
        return True

    for i in res:
        combinedResults = dict(baseResult.items() + i.items()  + {'elapsedTime': round(time.time() - startTime,0)}.items() )
        print 'the result is ', combinedResults
        if bucket is not None:
            if 'testMetric' in combinedResults:
                testKey = combinedResults['testName'] + '-' + combinedResults['testMetric']
            else:
                testKey = combinedResults['testName']
            bucket.upsert( testStartTime + '-' + options.version + '-' + testKey, combinedResults, format=couchbase.FMT_JSON)

def main():
    print 'Starting the perf regression runner'

    usage = '%prog -f conf-file'
    #parser = OptionParser(usage)

    ##parser.add_option('-f', '--filename', dest='filename')
    #parser.add_option('-v', '--version', dest='version')
    #parser.add_option('-u', '--url', dest='url')
    #parser.add_option('-q', '--query', dest='query')
    #parser.add_option('-s', '--specFile', dest='specFile')
    #parser.add_option('-r', '--runStartTime', dest='runStartTime')
    #parser.add_option('-b', '--betaTests', dest='betaTests', default=False, action='store_true')
    #parser.add_option('-n', '--nop', dest='nop',default=False, action='store_true')

    #options, args = parser.parse_args()

    # the option parsing way
    parser = argparse.ArgumentParser(description=usage)
    parser.add_argument('-q', '--query', nargs='+')
    parser.add_argument('-v', '--version', dest='version')
    parser.add_argument('-u', '--url', dest='url')
    parser.add_argument('-s', '--specFile', dest='specFile')
    parser.add_argument('-r', '--runStartTime', dest='runStartTime')
    parser.add_argument('-b', '--betaTests', dest='betaTests', default=False, action='store_true')
    parser.add_argument('-a', '--allTests', dest='allTests', default=False, action='store_true')
    parser.add_argument('-n', '--nop', dest='nop',default=False, action='store_true')

    options = parser.parse_args()

    print 'query', options.query
    print 'specfile', options.specFile

    runStartTime = options.runStartTime
    summary = []

    print 'version', options.version
    print 'url', options.url

    # open the bucket
    if options.nop:
        bucket = None
    else:
        bucket = Bucket('couchbase://'+ '172.23.105.177:8091/Daily-Performance')

    testBucket = Bucket('couchbase://'+ '172.23.105.177:8091/Daily-Performance-Tests')
    queryString = "select `Daily-Performance-Tests`.* from `Daily-Performance-Tests`"
    wherePredicates = []
    if options.query is not None:
        wherePredicates.append( ' '.join(options.query) )
    if not options.allTests:
        if options.betaTests:
            wherePredicates.append( "status='beta'")
        else:
            wherePredicates.append( "status!='beta'")

    if len(wherePredicates) > 0:
       for i in range(len(wherePredicates)):
          if i == 0:
              queryString += ' where '
          else:
              queryString += ' and '
          queryString += wherePredicates[i]

    print 'the query string is', queryString
    query = N1QLQuery(queryString )
    testsToRun = testBucket.n1ql_query( queryString )
    tests = [row for row in testsToRun]
    print 'the tests are', tests
    testsToRerun = []

    for row in tests:
        try:
            if row['status'].lower() == 'disabled':
                print row['testName'], ' is disabled.'
            else:
                if not runTest( row, options, bucket ):
                    testsToRerun.append(row)
        except:
            print 'Exception in ', row['testName']
            traceback.print_exc()

        #time.sleep(10)
    # end the for loop - print the results
    print 'done'

if __name__ == "__main__":
    if not main():
        sys.exit(1)

<|endoftext|>


<commit_before>from .helpers import *
import os
import discord
from collections import OrderedDict

class BotDataItem:
	def __init__(self, botdata, list_key, primary_keys, defaults):
		self.__dict__.update({
			"_botdata": botdata,
			"_list_key": list_key,
			"_primary_keys": primary_keys,
			"defaults": defaults
		})

	@property
	def json_data(self):
		for item in self._botdata.json_data[self._list_key]:
			if all(item.get(key) == self._primary_keys[key] for key in self._primary_keys):
				return item
		return None

	def __getattr__(self, key):
		if key in self._primary_keys:
			return self._primary_keys[key]
		if key not in self.defaults:
			raise ValueError(f"Tried to get invalid '{key}' in {self._list_key}")
		if self.json_data:
			return self.json_data.get(key, self.defaults.get(key))
		return self.defaults.get(key)

	def __setattr__(self, key, val):
		if key in self._primary_keys:
			raise ValueError("You can't set a primary key")
		if key not in self.defaults:
			raise ValueError(f"Tried to set invalid '{key}' in {self._list_key}")

		# recreate to order correctly
		newdict = OrderedDict(self._primary_keys)
		for k in self.defaults:
			if k == key:
				if val != self.defaults[key]:
					newdict[k] = val
			elif self.json_data and k in self.json_data:
				newdict[k] = self.json_data[k]
		# now save to json
		if self.json_data:
			index = self._botdata.json_data[self._list_key].index(self.json_data)
			self._botdata.json_data[self._list_key][index] = newdict
		else:
			self._botdata.json_data[self._list_key].append(newdict)
		self._botdata.save_data()

class UserInfo(BotDataItem):
	def __init__(self, botdata, discord):
		BotDataItem.__init__(self, botdata, "userinfo", { "discord": discord }, OrderedDict([
			("steam32", None),
			("intro", ""),
			("outro", "")
		]))

class GuildInfo:
	def __init__(self, botdata, guildid):
		self.botdata = botdata
		self.id = guildid
		if(self.json_data is None):
			self.botdata.json_data['guildinfo'].append(OrderedDict([
				("id", self.id),
				("voicechannel", None),
				("reactions", False),
				("invalidcommands", False),
				("banned_users", []),
				("ttschannel", [])
			]))
			self.botdata.save_data()

	@property
	def json_data(self):
		for guild in self.botdata.json_data['guildinfo']:
			if (guild['id'] == self.id):
				return guild
		# Should only happen when loading a guildinfo for the first time
		return None

	@property
	def voicechannel(self):
		return self.json_data.get("voicechannel", None)

	@voicechannel.setter
	def voicechannel(self, value):
		self.json_data["voicechannel"] = value
		self.botdata.save_data()

	@property
	def ttschannel(self):
		return self.json_data.get("ttschannel", None)

	@ttschannel.setter
	def ttschannel(self, value):
		self.json_data["ttschannel"] = value
		self.botdata.save_data()

	@property
	def reactions(self):
		return self.json_data.get("reactions", False)

	@property
	def invalidcommands(self):
		return self.json_data.get("invalidcommands", False)

	def is_banned(self, user):
		return user.id in self.json_data.get("banned_users", [])

	def botban(self, user):
		if self.json_data.get("banned_users") is None:
			self.json_data["banned_users"] = []
		if user.id not in self.json_data["banned_users"]:
			self.json_data["banned_users"].append(user.id)
			self.botdata.save_data()

	def botunban(self, user):
		if self.json_data.get("banned_users") is None:
			self.json_data["banned_users"] = []
		if user.id in self.json_data["banned_users"]:
			self.json_data["banned_users"].remove(user.id)
			self.botdata.save_data()

class BotData:
	def __init__(self):
		self.path = "botdata.json"
		self.defaults = OrderedDict([ ("userinfo" , []), ("guildinfo" , []) ])
		if not os.path.exists(self.path):
			self.json_data = self.defaults
			self.save_data()
		else:
			current = read_json(self.path)
			if current.keys() != self.defaults.keys():
				for key in self.defaults.keys():
					if key not in current.keys():
						current[key] = self.defaults[key]
						print("Adding " + str(key) + " field to botdata.json")
				write_json(self.path, current)
			self.json_data = read_json(self.path)

	def save_data(self):
		write_json(self.path, self.json_data)

	def userinfo(self, userid):
		if isinstance(userid, discord.User):
			userid = userid.id
		return UserInfo(self, userid)

	def guildinfo(self, guildid):
		if isinstance(guildid, discord.Guild):
			guildid = guildid.id
		return GuildInfo(self, guildid)

	def guildinfo_list(self):
		guildinfos = []
		for data in self.json_data['guildinfo']:
			guildinfos.append(GuildInfo(self, data['id']))
		return guildinfos

<commit_msg>implemented new botdata system for guildinfo<commit_after>from .helpers import *
import os
import discord
from collections import OrderedDict

class BotDataItem:
	def __init__(self, botdata, list_key, primary_keys, defaults):
		self.__dict__.update({
			"_botdata": botdata,
			"_list_key": list_key,
			"_primary_keys": primary_keys,
			"defaults": defaults
		})

	@property
	def json_data(self):
		for item in self._botdata.json_data[self._list_key]:
			if all(item.get(key) == self._primary_keys[key] for key in self._primary_keys):
				return item
		return None

	def __getattr__(self, key):
		if key in self._primary_keys:
			return self._primary_keys[key]
		if key not in self.defaults:
			raise ValueError(f"Tried to get invalid '{key}' in {self._list_key}")
		if self.json_data:
			return self.json_data.get(key, self.defaults.get(key))
		return self.defaults.get(key)

	def __setattr__(self, key, val):
		if key in self._primary_keys:
			raise ValueError("You can't set a primary key")
		if key not in self.defaults:
			raise ValueError(f"Tried to set invalid '{key}' in {self._list_key}")

		# recreate to order correctly
		newdict = OrderedDict(self._primary_keys)
		for k in self.defaults:
			if k == key:
				if val != self.defaults[key]:
					newdict[k] = val
			elif self.json_data and k in self.json_data:
				newdict[k] = self.json_data[k]
		# now save to json
		if self.json_data:
			index = self._botdata.json_data[self._list_key].index(self.json_data)
			self._botdata.json_data[self._list_key][index] = newdict
		else:
			self._botdata.json_data[self._list_key].append(newdict)
		self._botdata.save_data()

	__getitem__ = __getattr__
	__setitem__ = __setattr__

	# adds an item to a list variable, like banned_users
	def add_list_item(self, key, item):
		if item not in self[key]:
			new_list = list(self[key])
			new_list.append(item)
			self[key] = new_list

	# removes an item from a list variable, like banned_users
	def remove_list_item(self, key, item):
		if item in self[key]:
			new_list = list(self[key])
			new_list.remove(item)
			self[key] = new_list

class UserInfo(BotDataItem):
	def __init__(self, botdata, discord):
		BotDataItem.__init__(self, botdata, "userinfo", { "discord": discord }, OrderedDict([
			("steam32", None),
			("intro", ""),
			("outro", "")
		]))

class GuildInfo(BotDataItem):
	def __init__(self, botdata, guildid):
		BotDataItem.__init__(self, botdata, "guildinfo", { "id": guildid }, OrderedDict([
			("voicechannel", None),
			("reactions", False),
			("invalidcommands", False),
			("banned_users", []),
			("ttschannel", None)
		]))

	def is_banned(self, user):
		return user.id in self.banned_users

	def botban(self, user):
		self.add_list_item("banned_users", user.id)

	def botunban(self, user):
		self.remove_list_item("banned_users", user.id)

class BotData:
	def __init__(self):
		self.path = "botdata.json"
		self.defaults = OrderedDict([ ("userinfo" , []), ("guildinfo" , []) ])
		if not os.path.exists(self.path):
			self.json_data = self.defaults
			self.save_data()
		else:
			current = read_json(self.path)
			if current.keys() != self.defaults.keys():
				for key in self.defaults.keys():
					if key not in current.keys():
						current[key] = self.defaults[key]
						print("Adding " + str(key) + " field to botdata.json")
				write_json(self.path, current)
			self.json_data = read_json(self.path)

	def save_data(self):
		write_json(self.path, self.json_data)

	def userinfo(self, userid):
		if isinstance(userid, discord.User):
			userid = userid.id
		return UserInfo(self, userid)

	def guildinfo(self, guildid):
		if isinstance(guildid, discord.Guild):
			guildid = guildid.id
		return GuildInfo(self, guildid)

	def guildinfo_list(self):
		guildinfos = []
		for data in self.json_data['guildinfo']:
			guildinfos.append(GuildInfo(self, data['id']))
		return guildinfos

<|endoftext|>


<commit_before>"""
    Credit goes to https://stackoverflow.com/users/684592/kotaro
    https://stackoverflow.com/questions/14716497/how-can-i-find-a-list-of-street-intersections-from-openstreetmap-data?rq=1
"""

try:
    from xml.etree import cElementTree as ET
except ImportError, e:
    from xml.etree import ElementTree as ET

def extract_intersections(osm):
    """
    This method reads the passed osm file (xml) and finds intersections (nodes that are shared by two or more roads)
    :param osm: An osm file or a string from get_osm()
    """
    ret_val = []

    # step 1: parse either XML string or file
    if '<' in osm and '>' in osm:
        tree = ET.fromstring(osm)
        children = tree.getchildren()
    else:
        tree = ET.parse(osm)
        root = tree.getroot()
        children = root.getchildren()

    counter = {}
    for child in children:
        if child.tag == 'way':
            is_road = False

            # TODO: make road types configurable
            road_types = ('primary', 'secondary', 'residential', 'tertiary', 'service')
            for item in child:
                if item.tag == 'tag' and item.attrib['k'] == 'highway' and item.attrib['v'] in road_types:
                    is_road = True
                    break

            if is_road:
                for item in child:
                    if item.tag == 'nd':
                        nd_ref = item.attrib['ref']
                        if nd_ref not in counter:
                            counter[nd_ref] = 0
                        counter[nd_ref] += 1

    # Find nodes that are shared with more than one way, which might correspond to intersections
    intersections = filter(lambda x: counter[x] > 1, counter)

    for child in children:
        if child.tag == 'node' and child.attrib['id'] in intersections:
            coordinate = child.attrib['lat'] + ',' + child.attrib['lon']
            ret_val.append(coordinate)

    return ret_val

def main():
    intersections = extract_intersections('./test/portland.osm')
    for entry in intersections:
        print entry

if __name__ == '__main__':
    main()
<commit_msg>initial add of intersections converter<commit_after>"""
    Credit goes to https://stackoverflow.com/users/684592/kotaro
    https://stackoverflow.com/questions/14716497/how-can-i-find-a-list-of-street-intersections-from-openstreetmap-data?rq=1
"""

try:
    from xml.etree import cElementTree as ET
except ImportError, e:
    from xml.etree import ElementTree as ET

def extract_intersections(osm):
    """
    This method reads the passed osm file (xml) and finds intersections (nodes that are shared by two or more roads)
    :param osm: An osm file or a string from get_osm()
    """
    ret_val = []

    # step 1: parse either XML string or file
    if '<' in osm and '>' in osm:
        tree = ET.fromstring(osm)
        children = tree.getchildren()
    else:
        tree = ET.parse(osm)
        root = tree.getroot()
        children = root.getchildren()

    counter = {}
    for child in children:
        if child.tag == 'way':
            is_road = False

            # TODO: make road types configurable
            road_types = ('primary', 'secondary', 'residential', 'tertiary', 'service')
            for item in child:
                if item.tag == 'tag' and item.attrib['k'] == 'highway' and item.attrib['v'] in road_types:
                    is_road = True
                    break

            if is_road:
                for item in child:
                    if item.tag == 'nd':
                        nd_ref = item.attrib['ref']
                        if nd_ref not in counter:
                            counter[nd_ref] = 0
                        counter[nd_ref] += 1

    # Find nodes that are shared with more than one way, which might correspond to intersections
    intersections = filter(lambda x: counter[x] > 1, counter)

    for child in children:
        if child.tag == 'node' and child.attrib['id'] in intersections:
            coordinate = child.attrib['lat'] + ',' + child.attrib['lon']
            ret_val.append(coordinate)

    return ret_val

def main():
    intersections = extract_intersections('./test/portland.osm')
    for entry in intersections:
        print entry

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


<commit_before>import os
import tarfile
import time
import io
import stat

class SpoutManager:
    """
    A convenience context manager for creating spouts, allowing you to create spout code like:

    ```
    while True:
        with SpoutManager() as spout:
            spout.add_from_bytes("foo", b"#")
        time.sleep(1.0)
    ```
    """

    def __init__(self, max_attempts=10, sleep_time=1):
        """
        Creates a new spout manager.

        Params:

        * `max_attempts`: An optional int specifying the maximum number of
        times the manager should attempt to open the spout tarpipe.
        * `sleep_time`: An optional int specifying the sleep time between
        attempting to open the spout tarpipe.
        """

        self.f = None
        self.max_attempts = max_attempts
        self.sleep_time = 1

    def __enter__(self):
        attempts = 0
        umask_original = os.umask(0o777 ^ stat.S_IWUSR)

        try:
            while True:
                try:
                    f1 = os.open("/pfs/out", os.O_WRONLY, stat.S_IWUSR)
                    f2 = os.fdopen(f1, "wb")
                    self.f = tarfile.open(fileobj=f2, mode="w|", encoding="utf-8")
                    return self
                except OSError:
                    attempts += 1
                    if attempts == self.max_attempts:
                        raise
                    time.sleep(self.sleep_time)
        finally:
            os.umask(umask_original)

    def __exit__(self, type, value, traceback):
        self.f.close()

    def add_from_fileobj(self, path, size, fileobj):
        """
        Adds a file to the spout from a file-like object.

        Params:

        * `path`: The path to the file in the spout.
        * `size`: The size of the file.
        * `fileobj`: The file-like object to add.
        """

        tar_info = tarfile.TarInfo(path)
        tar_info.size = size
        tar_info.mode = 0o600
        self.f.addfile(tarinfo=tar_info, fileobj=fileobj)

    def add_from_bytes(self, path, bytes):
        """
        Adds a file to the spout from a bytestring.

        Params:

        * `path`: The path to the file in the spout.
        * `bytes`: The bytestring representing the file contents.
        """

        self.add_from_fileobj(path, len(bytes), io.BytesIO(bytes))
<commit_msg>Attempt retries for file addition too<commit_after>import os
import tarfile
import time
import io
import stat

class SpoutManager:
    """
    A convenience context manager for creating spouts, allowing you to create spout code like:

    ```
    while True:
        with SpoutManager() as spout:
            spout.add_from_bytes("foo", b"#")
        time.sleep(1.0)
    ```
    """

    def __init__(self, max_attempts=10, sleep_time=1):
        """
        Creates a new spout manager.

        Params:

        * `max_attempts`: An optional int specifying the maximum number of
        times the manager should attempt to open the spout tarpipe.
        * `sleep_time`: An optional int specifying the sleep time between
        attempting to open the spout tarpipe.
        """

        self.f = None
        self.max_attempts = max_attempts
        self.sleep_time = 1

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        if self.f is None:
            self.f.close()

    def add_from_fileobj(self, path, size, fileobj):
        """
        Adds a file to the spout from a file-like object.

        Params:

        * `path`: The path to the file in the spout.
        * `size`: The size of the file.
        * `fileobj`: The file-like object to add.
        """

        attempts = 0
        umask_original = os.umask(0o777 ^ stat.S_IWUSR)

        try:
            while True:
                try:
                    if self.f is None:
                        f1 = os.open("/pfs/out", os.O_WRONLY, stat.S_IWUSR)
                        f2 = os.fdopen(f1, "wb")
                        self.f = tarfile.open(fileobj=f2, mode="w|", encoding="utf-8")

                    tar_info = tarfile.TarInfo(path)
                    tar_info.size = size
                    tar_info.mode = 0o600
                    self.f.addfile(tarinfo=tar_info, fileobj=fileobj)
                except:
                    attempts += 1
                    if attempts == self.max_attempts:
                        raise
                    time.sleep(self.sleep_time)
        finally:
            os.umask(umask_original)

    def add_from_bytes(self, path, bytes):
        """
        Adds a file to the spout from a bytestring.

        Params:

        * `path`: The path to the file in the spout.
        * `bytes`: The bytestring representing the file contents.
        """

        self.add_from_fileobj(path, len(bytes), io.BytesIO(bytes))
<|endoftext|>