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

from ._base import SQLBaseStore
from twisted.internet import defer
from synapse.util.caches.descriptors import cachedInlineCallbacks

import logging
import ujson as json

logger = logging.getLogger(__name__)

class EventPushActionsStore(SQLBaseStore):
    def _set_push_actions_for_event_and_users_txn(self, txn, event, tuples):
        """
        :param event: the event set actions for
        :param tuples: list of tuples of (user_id, actions)
        """
        values = []
        for uid, actions in tuples:
            values.append({
                'room_id': event.room_id,
                'event_id': event.event_id,
                'user_id': uid,
                'actions': json.dumps(actions),
                'stream_ordering': event.internal_metadata.stream_ordering,
                'topological_ordering': event.depth,
                'notif': 1,
                'highlight': 1 if _action_has_highlight(actions) else 0,
            })

        for uid, __ in tuples:
            txn.call_after(
                self.get_unread_event_push_actions_by_room_for_user.invalidate_many,
                (event.room_id, uid)
            )
        self._simple_insert_many_txn(txn, "event_push_actions", values)

    @cachedInlineCallbacks(num_args=3, lru=True, tree=True)
    def get_unread_event_push_actions_by_room_for_user(
            self, room_id, user_id, last_read_event_id
    ):
        def _get_unread_event_push_actions_by_room(txn):
            sql = (
                "SELECT stream_ordering, topological_ordering"
                " FROM events"
                " WHERE room_id = ? AND event_id = ?"
            )
            txn.execute(
                sql, (room_id, last_read_event_id)
            )
            results = txn.fetchall()
            if len(results) == 0:
                return {"notify_count": 0, "highlight_count": 0}

            stream_ordering = results[0][0]
            topological_ordering = results[0][1]

            sql = (
                "SELECT sum(notif), sum(highlight)"
                " FROM event_push_actions ea"
                " WHERE"
                " user_id = ?"
                " AND room_id = ?"
                " AND ("
                "       topological_ordering > ?"
                "       OR (topological_ordering = ? AND stream_ordering > ?)"
                ")"
            )
            txn.execute(sql, (
                user_id, room_id,
                topological_ordering, topological_ordering, stream_ordering
            ))
            row = txn.fetchone()
            if row:
                return {
                    "notify_count": row[0] or 0,
                    "highlight_count": row[1] or 0,
                }
            else:
                return {"notify_count": 0, "highlight_count": 0}

        ret = yield self.runInteraction(
            "get_unread_event_push_actions_by_room",
            _get_unread_event_push_actions_by_room
        )
        defer.returnValue(ret)

    def _remove_push_actions_for_event_id_txn(self, txn, room_id, event_id):
        # Sad that we have to blow away the cache for the whole room here
        txn.call_after(
            self.get_unread_event_push_actions_by_room_for_user.invalidate_many,
            (room_id,)
        )
        txn.execute(
            "DELETE FROM event_push_actions WHERE room_id = ? AND event_id = ?",
            (room_id, event_id)
        )

def _action_has_highlight(actions):
    for action in actions:
        try:
            if action.get("set_tweak", None) == "highlight":
                return action.get("value", True)
        except AttributeError:
            pass

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

from ._base import SQLBaseStore
from twisted.internet import defer
from synapse.util.caches.descriptors import cachedInlineCallbacks

import logging
import ujson as json

logger = logging.getLogger(__name__)

class EventPushActionsStore(SQLBaseStore):
    def _set_push_actions_for_event_and_users_txn(self, txn, event, tuples):
        """
        :param event: the event set actions for
        :param tuples: list of tuples of (user_id, actions)
        """
        values = []
        for uid, actions in tuples:
            values.append({
                'room_id': event.room_id,
                'event_id': event.event_id,
                'user_id': uid,
                'actions': json.dumps(actions),
                'stream_ordering': event.internal_metadata.stream_ordering,
                'topological_ordering': event.depth,
                'notif': 1,
                'highlight': 1 if _action_has_highlight(actions) else 0,
            })

        for uid, __ in tuples:
            txn.call_after(
                self.get_unread_event_push_actions_by_room_for_user.invalidate_many,
                (event.room_id, uid)
            )
        self._simple_insert_many_txn(txn, "event_push_actions", values)

    @cachedInlineCallbacks(num_args=3, lru=True, tree=True, max_entries=5000)
    def get_unread_event_push_actions_by_room_for_user(
            self, room_id, user_id, last_read_event_id
    ):
        def _get_unread_event_push_actions_by_room(txn):
            sql = (
                "SELECT stream_ordering, topological_ordering"
                " FROM events"
                " WHERE room_id = ? AND event_id = ?"
            )
            txn.execute(
                sql, (room_id, last_read_event_id)
            )
            results = txn.fetchall()
            if len(results) == 0:
                return {"notify_count": 0, "highlight_count": 0}

            stream_ordering = results[0][0]
            topological_ordering = results[0][1]

            sql = (
                "SELECT sum(notif), sum(highlight)"
                " FROM event_push_actions ea"
                " WHERE"
                " user_id = ?"
                " AND room_id = ?"
                " AND ("
                "       topological_ordering > ?"
                "       OR (topological_ordering = ? AND stream_ordering > ?)"
                ")"
            )
            txn.execute(sql, (
                user_id, room_id,
                topological_ordering, topological_ordering, stream_ordering
            ))
            row = txn.fetchone()
            if row:
                return {
                    "notify_count": row[0] or 0,
                    "highlight_count": row[1] or 0,
                }
            else:
                return {"notify_count": 0, "highlight_count": 0}

        ret = yield self.runInteraction(
            "get_unread_event_push_actions_by_room",
            _get_unread_event_push_actions_by_room
        )
        defer.returnValue(ret)

    def _remove_push_actions_for_event_id_txn(self, txn, room_id, event_id):
        # Sad that we have to blow away the cache for the whole room here
        txn.call_after(
            self.get_unread_event_push_actions_by_room_for_user.invalidate_many,
            (room_id,)
        )
        txn.execute(
            "DELETE FROM event_push_actions WHERE room_id = ? AND event_id = ?",
            (room_id, event_id)
        )

def _action_has_highlight(actions):
    for action in actions:
        try:
            if action.get("set_tweak", None) == "highlight":
                return action.get("value", True)
        except AttributeError:
            pass

    return False
<|endoftext|>


<commit_before>"""
Various models for calculating emission from multiple ions
"""
import numpy as np
from scipy.interpolate import splrep, splev
import astropy.units as u
from astropy.utils.console import ProgressBar
import h5py
import fiasco

class EmissionModel(fiasco.IonCollection):
    """
    Model for how atomic data is used to calculate emission from
    coronal plasma.
    """

    @u.quantity_input
    def __init__(self, density: u.cm**(-3), *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.temperature = self[0].temperature
        self.density = density
        self.resolved_wavelengths = kwargs.get('resolved_wavelengths',{})

    def interpolate_to_mesh_indices(self, loop):
        """
        Return interpolated loop indices to the temperature and density meshes defined for
        the atomic data. For use with `~scipy.ndimage.map_coordinates`.
        """
        nots_itemperature = splrep(self.temperature.value, np.arange(self.temperature.shape[0]))
        nots_idensity = splrep(self.density.value, np.arange(self.density.shape[0]))
        itemperature = splev(np.ravel(loop.electron_temperature.value), nots_itemperature)
        idensity = splev(np.ravel(loop.density.value), nots_idensity)

        return itemperature, idensity

    def calculate_emissivity(self, savefile, **kwargs):
        """
        Calculate and store emissivity for every ion in the model
        """
        self.emissivity_savefile = savefile
        with h5py.File(savefile, 'w') as hf:
            with ProgressBar(len(self._ion_list), ipython_widget=kwargs.get('notebook', True)) as progress:
                for ion in self:
                    wavelength, emissivity = ion.emissivity(self.density, include_energy=False)
                    if wavelength is None or emissivity is None:
                        continue
                    emissivity = emissivity[:,:,np.argsort(wavelength)]
                    wavelength = np.sort(wavelength)
                    grp = hf.create_group(ion.ion_name)
                    ds = grp.create_dataset('wavelength', data=wavelength.value)
                    ds.attrs['units'] = wavelength.unit.to_string()
                    ds = grp.create_dataset('emissivity', data=emissivity.data)
                    ds.attrs['units'] = emissivity.unit.to_string()
                    progress.update()

    def get_emissivity(self, ion):
        with h5py.File(self.emissivity_savefile, 'r') as hf:
            if ion.ion_name not in hf:
                return (None, None)
            ds = hf['/'.join([ion.ion_name, 'wavelength'])]
            wavelength = u.Quantity(ds, ds.attrs['units'])
            ds = hf['/'.join([ion.ion_name, 'emissivity'])]
            emissivity = u.Quantity(ds,  ds.attrs['units'])

        return wavelength, emissivity

    def calculate_emission(self, loop, **kwargs):
        """
        Calculate power per unit volume for a given temperature and density for every transition,
        :math:`\lambda`, in every ion :math:`X^{+m}`, as given by the equation,

        .. math::
           P_{\lambda}(n,T) = \\frac{1}{4\pi}0.83\mathrm{Ab}(X)\\varepsilon_{\lambda}(n,T)\\frac{N(X^{+m})}{N(X)}n

        where :math:`\\mathrm{Ab}(X)` is the abundance of element :math:`X`,
        :math:`\\varepsilon_{\lambda}` is the emissivity for transition :math:`\lambda`,
        and :math:`N(X^{+m})/N(X)` is the ionization fraction of ion :math:`X^{+m}`.
        :math:`P_{\lambda}` is in units of erg cm\ :sup:`-3` s\ :sup:`-1` sr\ :sup:`-1` if
        `energy_unit` is set to `erg` and in units of photons
        cm\ :sup:`-3` s\ :sup:`-1` sr\ :sup:`-1` if `energy_unit` is set to `photon`.
        """
        raise NotImplementedError('Have not yet reimplemented emission calculation.')<commit_msg>separate emission model and field<commit_after>"""
Various models for calculating emission from multiple ions
"""
import numpy as np
from scipy.interpolate import splrep, splev, interp1d
import astropy.units as u
from astropy.utils.console import ProgressBar
import h5py
import fiasco

class EmissionModel(fiasco.IonCollection):
    """
    Model for how atomic data is used to calculate emission from
    coronal plasma.
    """

    @u.quantity_input
    def __init__(self, density: u.cm**(-3), *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.temperature = self[0].temperature
        self.density = density
        self.resolved_wavelengths = kwargs.get('resolved_wavelengths',{})

    def interpolate_to_mesh_indices(self, loop):
        """
        Return interpolated loop indices to the temperature and density meshes defined for
        the atomic data. For use with `~scipy.ndimage.map_coordinates`.
        """
        nots_itemperature = splrep(self.temperature.value, np.arange(self.temperature.shape[0]))
        nots_idensity = splrep(self.density.value, np.arange(self.density.shape[0]))
        itemperature = splev(np.ravel(loop.electron_temperature.value), nots_itemperature)
        idensity = splev(np.ravel(loop.density.value), nots_idensity)

        return itemperature, idensity

    def calculate_emissivity(self, savefile, **kwargs):
        """
        Calculate and store emissivity for every ion in the model
        """
        self.emissivity_savefile = savefile
        with h5py.File(savefile, 'w') as hf:
            with ProgressBar(len(self._ion_list), ipython_widget=kwargs.get('notebook', True)) as progress:
                for ion in self:
                    wavelength, emissivity = ion.emissivity(self.density, include_energy=False)
                    if wavelength is None or emissivity is None:
                        continue
                    emissivity = emissivity[:,:,np.argsort(wavelength)]
                    wavelength = np.sort(wavelength)
                    grp = hf.create_group(ion.ion_name)
                    ds = grp.create_dataset('wavelength', data=wavelength.value)
                    ds.attrs['units'] = wavelength.unit.to_string()
                    ds = grp.create_dataset('emissivity', data=emissivity.data)
                    ds.attrs['units'] = emissivity.unit.to_string()
                    progress.update()

    def get_emissivity(self, ion):
        with h5py.File(self.emissivity_savefile, 'r') as hf:
            if ion.ion_name not in hf:
                return (None, None)
            ds = hf['/'.join([ion.ion_name, 'wavelength'])]
            wavelength = u.Quantity(ds, ds.attrs['units'])
            ds = hf['/'.join([ion.ion_name, 'emissivity'])]
            emissivity = u.Quantity(ds,  ds.attrs['units'])

        return wavelength, emissivity

    def calculate_ionization_fraction(self, field, savefile, interface=None, **kwargs):
        """
        Find the fractional ionization for each loop in the model as defined by the loop
        model interface.
        """
        self.ionization_fraction_savefile = savefile
        with h5py.File(self.ionization_fraction_savefile, 'a') as hf:
            for loop in field.loops:
                if loop.name not in hf:
                    grp = hf.create_group(loop.name)
                else:
                    grp = hf[loop.name]
                if interface:
                    description = 'non-equilibrium ionization fractions'
                    ionization_fraction = interface.get_ionization_fraction(loop, emission_model, **kwargs)
                else:
                    description = 'equilibrium ionization fractions'
                    ionization_fraction = {}
                    for ion in self:
                        f_ioneq = interp1d(ion.temperature, ion.ioneq, fill_value='extrapolate')
                        tmp = f_ioneq(loop.electron_temperature)
                        ionization_fraction[ion.ion_name] = u.Quantity(np.where(tmp < 0., 0., tmp), ion.ioneq.unit)

                for key in ionization_fraction:
                    if key not in grp:
                        dset = grp.create_dataset(key, data=ionization_fraction[key].value)
                    else:
                        dset = grp[key]
                    dset[:,:] = ionization_fraction[key].value
                    dset.attrs['units'] = ionization_fraction[key].unit.to_string()
                    dset.attrs['description'] = description

    def get_ionization_fraction(self, loop, ion):
        """
        Get ionization state from the ionization balance equations.

        Note
        ----
        This can be either the equilibrium or the non-equilibrium ionization
        fraction, depending on which was calculated.
        """
        with h5py.File(self.ionization_fraction_savefile, 'r') as hf:
            dset = hf['/'.join([loop.name, ion.ion_name])]
            ionization_fraction = u.Quantity(dset, dset.attrs['units'])

        return ionization_fraction

    def calculate_emission(self, loop, **kwargs):
        """
        Calculate power per unit volume for a given temperature and density for every transition,
        :math:`\lambda`, in every ion :math:`X^{+m}`, as given by the equation,

        .. math::
           P_{\lambda}(n,T) = \\frac{1}{4\pi}0.83\mathrm{Ab}(X)\\varepsilon_{\lambda}(n,T)\\frac{N(X^{+m})}{N(X)}n

        where :math:`\\mathrm{Ab}(X)` is the abundance of element :math:`X`,
        :math:`\\varepsilon_{\lambda}` is the emissivity for transition :math:`\lambda`,
        and :math:`N(X^{+m})/N(X)` is the ionization fraction of ion :math:`X^{+m}`.
        :math:`P_{\lambda}` is in units of erg cm\ :sup:`-3` s\ :sup:`-1` sr\ :sup:`-1` if
        `energy_unit` is set to `erg` and in units of photons
        cm\ :sup:`-3` s\ :sup:`-1` sr\ :sup:`-1` if `energy_unit` is set to `photon`.
        """
        raise NotImplementedError('Have not yet reimplemented emission calculation.')<|endoftext|>


<commit_before>import hashlib
import struct

# Length of encoded index in bytes.
INDEX_ENCODE_LENGTH = 6
# Maximum supported pad index.
INDEX_MAX = (2**(INDEX_ENCODE_LENGTH * 8)) - 1
# Length of tag in bytes.
TAG_LENGTH = 64

# TODO(miles): make asserts raise cryptoerrors.

class CryptoError(Exception):
    pass

def package(index, message, p_text, p_body):
    """
    Secure a message with encryption and integrity protection.

    package := i || body
    body := ciphertext || tag
    tag := SHA(i || ciphertext)
    ciphertext := p_text XOR message

    - '||' means concatenation.
    - i is a fixed-length encoding of the pad index.
    - SHA refers to a SHA256 hash.

    Args:
        index: Index into the pad used. (int)
        message: Plaintext message to encrypt. (string)
        p_cipher: Bytes to use for XORing with message. (string)
        p_body: Bytes to use for XORing with message || tag. (string)

    Returns:
        The secured message, ready for sending. (string)
    """
    # Assert parameter types.
    assert isinstance(index, int)
    assert isinstance(message, str)
    assert isinstance(p_text, str)
    assert isinstance(p_body, str)

    # Assert various properties we know should be true.
    assert 0 <= index <= INDEX_MAX
    assert len(p_text) == len(message)
    assert len(p_body) == len(message) + TAG_LENGTH

    # Encrypt the plaintext.
    ciphertext = encrypt(message, p_text)

    # Encode the index.
    i_enc = encode_index(index)
    assert len(i_enc) == INDEX_ENCODE_LENGTH

    # Create the integrity tag.
    tag = sha(i_enc + ciphertext)
    assert len(tag) == TAG_LENGTH

    return i_enc + ciphertext + tag

def unpackage(package, p_text, p_body):
    """ Extract the message and verify it's integrity """

    # extract package contents
    index = package[:INDEX_ENCODE_LENGTH]
    encrypted_body = package[INDEX_ENCODE_LENGTH:]
    plain_body = decrypt(encrypted_body,p_body)
    ## the starting index of the body pad: index + length of cipher text
    ## length of cipher text = length of body - length of tag
    tag = plain_body[-1*TAG_LENGTH:]
    ciphertext = plain_body[:-1*TAG_LENGTH]

    # integrity check
    if sha(index + ciphertext) == tag:
        message = decrypt(ciphertext,p_text)
        return message
    else:
        raise CryptoError("Integrity Error")

def pre_unpackage(package):
    """Extract what information is necessary to fetch before unpackaging.

    Returns:
    {
        p_text_index
        p_body_index
    }
    Raises:
        CryptoError on failure.
    """
    message_length = len(package) - INDEX_ENCODE_LENGTH - TAG_LENGTH
    assert message_length > 0
    body_length = message_length + TAG_LENGTH
    p_text_index = decode_index(package[:INDEX_ENCODE_LENGTH])
    p_body_index = p_text_index + message_length
    return {
        "message_length": message_length,
        "body_length": body_length,
        "p_text_index": p_text_index,
        "p_body_index": p_body_index,
    }

def encode_index(index_num):
    """
    Encode an index as a fixed length string.
    """
    assert 0 <= index_num <= INDEX_MAX
    assert INDEX_ENCODE_LENGTH < 8
    try:
        eight_pack = struct.pack("<Q", index_num)
    except struct.error as ex:
        raise CryptoError(ex)
    assert len(eight_pack) == 8
    return eight_pack[:INDEX_ENCODE_LENGTH]

def decode_index(index_str):
    """
    Decode an index encoded as a fixed length string.
    """
    assert isinstance(index_str, str)
    assert len(index_str) == INDEX_ENCODE_LENGTH
    eight_pack = index_str + "\x00" + "\x00"
    try:
        return struct.unpack("<Q", eight_pack)[0]
    except struct.error as ex:
        raise CryptoError(ex)

def xor(charA, charB):
    return chr(ord(charA) ^ ord(charB))

def xorHelper((charA,charB)):
    return xor(charA, charB)

def encrypt(message, pad):
    return "".join(map(xorHelper, zip(message, pad)))

def decrypt(cipher, pad):
    return "".join(map(xorHelper, zip(cipher, pad)))

def sha(stuff):
    return hashlib.sha512(stuff).digest()
<commit_msg>Notating some issues with package function<commit_after>import hashlib
import struct

# Length of encoded index in bytes.
INDEX_ENCODE_LENGTH = 6
# Maximum supported pad index.
INDEX_MAX = (2**(INDEX_ENCODE_LENGTH * 8)) - 1
# Length of tag in bytes.
TAG_LENGTH = 64

# TODO(miles): make asserts raise cryptoerrors.

class CryptoError(Exception):
    pass

def package(index, message, p_text, p_body):
    """
    Secure a message with encryption and integrity protection.

    package := i || body
    body := ciphertext || tag
    tag := SHA(i || ciphertext)
    ciphertext := p_text XOR message

    as defined aboue, total package:
       i || (p_text XOR message) || SHA(i || (p_text || ciphertext))
    but it SHOULD be:
       i || (p_text XOR message) || p_body XOR SHA(i || (p_text || ciphertext))

    - '||' means concatenation.
    - i is a fixed-length encoding of the pad index.
    - SHA refers to a SHA256 hash.

    Args:
        index: Index into the pad used. (int)
        message: Plaintext message to encrypt. (string)
        p_cipher: Bytes to use for XORing with message. (string)
        p_body: Bytes to use for XORing with message || tag. (string)

    Returns:
        The secured message, ready for sending. (string)
    """
    # Assert parameter types.
    assert isinstance(index, int)
    assert isinstance(message, str)
    assert isinstance(p_text, str)
    assert isinstance(p_body, str)

    # Assert various properties we know should be true.
    assert 0 <= index <= INDEX_MAX
    assert len(p_text) == len(message)
    assert len(p_body) == len(message) + TAG_LENGTH

    # Encrypt the plaintext.
    ciphertext = encrypt(message, p_text)

    # Encode the index.
    i_enc = encode_index(index)
    assert len(i_enc) == INDEX_ENCODE_LENGTH

    # Create the integrity tag.
    tag = sha(i_enc + ciphertext)
    assert len(tag) == TAG_LENGTH

    return i_enc + ciphertext + tag

def unpackage(package, p_text, p_body):
    """ Extract the message and verify it's integrity """

    # extract package contents
    index = package[:INDEX_ENCODE_LENGTH]
    encrypted_body = package[INDEX_ENCODE_LENGTH:]
    plain_body = decrypt(encrypted_body,p_body)
    ## the starting index of the body pad: index + length of cipher text
    ## length of cipher text = length of body - length of tag
    tag = plain_body[-1*TAG_LENGTH:]
    ciphertext = plain_body[:-1*TAG_LENGTH]

    # integrity check
    if sha(index + ciphertext) == tag:
        message = decrypt(ciphertext,p_text)
        return message
    else:
        raise CryptoError("Integrity Error")

def pre_unpackage(package):
    """Extract what information is necessary to fetch before unpackaging.

    Returns:
    {
        p_text_index
        p_body_index
    }
    Raises:
        CryptoError on failure.
    """
    message_length = len(package) - INDEX_ENCODE_LENGTH - TAG_LENGTH
    assert message_length > 0
    body_length = message_length + TAG_LENGTH
    p_text_index = decode_index(package[:INDEX_ENCODE_LENGTH])
    p_body_index = p_text_index + message_length
    return {
        "message_length": message_length,
        "body_length": body_length,
        "p_text_index": p_text_index,
        "p_body_index": p_body_index,
    }

def encode_index(index_num):
    """
    Encode an index as a fixed length string.
    """
    assert 0 <= index_num <= INDEX_MAX
    assert INDEX_ENCODE_LENGTH < 8
    try:
        eight_pack = struct.pack("<Q", index_num)
    except struct.error as ex:
        raise CryptoError(ex)
    assert len(eight_pack) == 8
    return eight_pack[:INDEX_ENCODE_LENGTH]

def decode_index(index_str):
    """
    Decode an index encoded as a fixed length string.
    """
    assert isinstance(index_str, str)
    assert len(index_str) == INDEX_ENCODE_LENGTH
    eight_pack = index_str + "\x00" + "\x00"
    try:
        return struct.unpack("<Q", eight_pack)[0]
    except struct.error as ex:
        raise CryptoError(ex)

def xor(charA, charB):
    return chr(ord(charA) ^ ord(charB))

def xorHelper((charA,charB)):
    return xor(charA, charB)

def encrypt(message, pad):
    return "".join(map(xorHelper, zip(message, pad)))

def decrypt(cipher, pad):
    return "".join(map(xorHelper, zip(cipher, pad)))

def sha(stuff):
    return hashlib.sha512(stuff).digest()
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go

https://github.com/ethereum/cpp-ethereum/blob/develop/libp2p/RLPxHandshake.cpp#L48
https://github.com/ethereum/cpp-ethereum/blob/develop/libdevcrypto/CryptoPP.cpp#L149

ECIES
http://www.cryptopp.com/wiki/Elliptic_Curve_Integrated_Encryption_Scheme
https://en.wikipedia.org/wiki/Integrated_Encryption_Scheme
"""
CURVE = 'secp256k1'
CIPHERNAME = 'aes-256-ctr'

# FIX PATH ON OS X ()
# https://github.com/yann2192/pyelliptic/issues/11
import os
import sys
_openssl_lib_paths = ['/usr/local/Cellar/openssl/']
for p in _openssl_lib_paths:
    if os.path.exists(p):
        p = os.path.join(p, os.listdir(p)[-1], 'lib')
        os.environ['DYLD_LIBRARY_PATH'] = p
        import pyelliptic
        if CIPHERNAME in pyelliptic.Cipher.get_all_cipher():
            break
if CIPHERNAME not in pyelliptic.Cipher.get_all_cipher():
    print 'required cipher %s not available in openssl library' % CIPHERNAME
    if sys.platform == 'darwin':
        print 'use homebrew to install newer openssl'
        print '> brew install openssl'
    sys.exit(1)

import bitcoin
from sha3 import sha3_256
from hashlib import sha256
import struct
import random
import devp2p.utils as utils

hmac_sha256 = pyelliptic.hmac_sha256

class ECCx(pyelliptic.ECC):

    """
    Modified to work with raw_pubkey format used in RLPx
    and binding default curve and cipher
    """

    def __init__(self, raw_pubkey=None, raw_privkey=None):
        if raw_privkey:
            assert not raw_pubkey
            raw_pubkey = privtopub(raw_privkey)
        if raw_pubkey:
            _, pubkey_x, pubkey_y, _ = self._decode_pubkey(raw_pubkey)
        else:
            pubkey_x, pubkey_y = None, None
        pyelliptic.ECC.__init__(self, pubkey_x=pubkey_x, pubkey_y=pubkey_y,
                                raw_privkey=raw_privkey, curve=CURVE)

    @property
    def raw_pubkey(self):
        return self.pubkey_x + self.pubkey_y

    @staticmethod
    def _decode_pubkey(raw_pubkey):
        assert len(raw_pubkey) == 64
        pubkey_x = raw_pubkey[:32]
        pubkey_y = raw_pubkey[32:]
        return CURVE, pubkey_x, pubkey_y, 64

    def get_ecdh_key(self, raw_pubkey):
        "Compute public key with the local private key and returns a 256bits shared key"
        _, pubkey_x, pubkey_y, _ = self._decode_pubkey(raw_pubkey)
        key = self.raw_get_ecdh_key(pubkey_x, pubkey_y)
        assert len(key) == 32
        return key

    @property
    def raw_privkey(self):
        return self.privkey

    @staticmethod
    def encrypt(data, raw_pubkey):
        assert False
        assert len(raw_pubkey) == 64
        px, py = raw_pubkey[:32], raw_pubkey[32:]
        return ECCx.raw_encrypt(data, px, py, curve=CURVE, ciphername=CIPHERNAME)

    def _decrypt(self, data):
        return pyelliptic.ECC.decrypt(self, data, ciphername=CIPHERNAME)

    ecies_ciphername = 'aes-128-ctr'

    @classmethod
    def ecies_encrypt(cls, data, raw_pubkey):
        """
        ECIES Encrypt, where P = recipient public key is:
        1) generate r = random value
        2) generate shared-secret = kdf( ecdhAgree(r, P) )
        3) generate R = rG [same op as generating a public key]
        4) send 0x04 || R || AsymmetricEncrypt(shared-secret, plaintext) || tag

        https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/params.go#L41

        currently used by go:
        ECIES_AES128_SHA256 = &ECIESParams{
            Hash: sha256.New,
            hashAlgo: crypto.SHA256,
            Cipher: aes.NewCipher,
            BlockSize: aes.BlockSize,
            KeyLen: 16,
            }

        """
        # 1) generate r = random value
        ephem = ECCx()

        # 2) generate shared-secret = kdf( ecdhAgree(r, P) )
        key_material = ephem.raw_get_ecdh_key(pubkey_x=raw_pubkey[:32], pubkey_y=raw_pubkey[32:])
        assert len(key_material) == 32
        key = eciesKDF(key_material, 32)
        assert len(key) == 32
        key_enc, key_mac = key[:16], key[16:]

        key_mac = sha256(key_mac).digest()  # !!!
        assert len(key_mac) == 32
        # 3) generate R = rG [same op as generating a public key]
        ephem_pubkey = ephem.raw_pubkey

        # encrypt
        iv = utils.ienc(random.getrandbits(16 * 8))
        assert len(iv) == 16
        ctx = pyelliptic.Cipher(key_enc, iv, 1, cls.ecies_ciphername)
        ciphertext = ctx.ciphering(data)
        assert len(ciphertext) == len(data)

        # 4) send 0x04 || R || AsymmetricEncrypt(shared-secret, plaintext) || tag
        msg = chr(0x04) + ephem_pubkey + iv + ciphertext

        # the MAC of a message (called the tag) as per SEC 1, 3.5.
        # https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go#L162
        tag = hmac_sha256(key_mac, msg[1 + 64:])
        assert len(tag) == 32
        msg += tag

        assert len(msg) == 1 + 64 + 16 + 32 + len(data) == 113 + len(data)
        return msg

    def ecies_decrypt(self, data):
        """
        Decrypt data with ECIES method using the local private key

        ECIES Decrypt (performed by recipient):
        1) generate shared-secret = kdf( ecdhAgree(myPrivKey, msg[1:65]) )
        2) verify tag
        3) decrypt

        ecdhAgree(r, recipientPublic) == ecdhAgree(recipientPrivate, R)
        [where R = r*G, and recipientPublic = recipientPrivate*G]

        """
        assert data[0] == chr(0x04)

        #  1) generate shared-secret = kdf( ecdhAgree(myPrivKey, msg[1:65]) )
        _shared = data[1:1 + 64]
        # FIXME, check that _shared_pub is a valid one (on curve)

        key_material = self.raw_get_ecdh_key(pubkey_x=_shared[:32], pubkey_y=_shared[32:])
        assert len(key_material) == 32
        key = eciesKDF(key_material, 32)
        assert len(key) == 32
        key_enc, key_mac = key[:16], key[16:]

        key_mac = sha256(key_mac).digest()
        assert len(key_mac) == 32

        tag = data[-32:]
        assert len(tag) == 32

        # 2) verify tag
        if not pyelliptic.equals(hmac_sha256(key_mac, data[1 + 64:- 32]), tag):
            raise RuntimeError("Fail to verify data")

        # 3) decrypt
        blocksize = pyelliptic.OpenSSL.get_cipher(self.ecies_ciphername).get_blocksize()
        iv = data[1 + 64:1 + 64 + blocksize]
        assert len(iv) == 16
        ciphertext = data[1 + 64 + blocksize:- 32]
        assert 1 + len(_shared) + len(iv) + len(ciphertext) + len(tag) == len(data)
        ctx = pyelliptic.Cipher(key_enc, iv, 0, self.ecies_ciphername)
        return ctx.ciphering(ciphertext)

    def sign(self, data):
        """
        pyelliptic.ECC.sign is DER-encoded
        https://bitcoin.stackexchange.com/questions/12554
        """
        signature = ecdsa_sign(data, self.raw_privkey)
        assert len(signature) == 65
        return signature

    def verify(self, signature, message):
        assert len(signature) == 65
        return ecdsa_verify(self.raw_pubkey, signature, message)

def lzpad32(x):
    return '\x00' * (32 - len(x)) + x

def _encode_sig(v, r, s):
    assert isinstance(v, (int, long))
    assert v in (27, 28)
    vb, rb, sb = chr(v - 27), bitcoin.encode(r, 256), bitcoin.encode(s, 256)
    return lzpad32(rb) + lzpad32(sb) + vb

def _decode_sig(sig):
    return ord(sig[64]) + 27, bitcoin.decode(sig[0:32], 256), bitcoin.decode(sig[32:64], 256)

def ecdsa_verify(pubkey, signature, message):
    assert len(signature) == 65
    assert len(pubkey) == 64
    return bitcoin.ecdsa_raw_verify(message, _decode_sig(signature), pubkey)
verify = ecdsa_verify

def ecdsa_sign(message, privkey):
    s = _encode_sig(*bitcoin.ecdsa_raw_sign(message, privkey))
    return s

sign = ecdsa_sign

def ecdsa_recover(message, signature):
    assert len(signature) == 65
    pub = bitcoin.ecdsa_raw_recover(message, _decode_sig(signature))
    assert pub, 'pubkey could not be recovered'
    pub = bitcoin.encode_pubkey(pub, 'bin_electrum')
    assert len(pub) == 64
    return pub
recover = ecdsa_recover

def sha3(seed):
    return sha3_256(seed).digest()

def mk_privkey(seed):
    return sha3(seed)

def privtopub(raw_privkey):
    raw_pubkey = bitcoin.encode_pubkey(bitcoin.privtopub(raw_privkey), 'bin_electrum')
    assert len(raw_pubkey) == 64
    return raw_pubkey

def encrypt(data, raw_pubkey):
    """
    Encrypt data with ECIES method using the public key of the recipient.
    """
    assert len(raw_pubkey) == 64
    return ECCx.encrypt(data, raw_pubkey)

def eciesKDF(key_material, key_len):
    """
    interop w/go ecies implementation

    for sha3, blocksize is 136 bytes
    for sha256, blocksize is 64 bytes

    NIST SP 800-56a Concatenation Key Derivation Function (see section 5.8.1).

    https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go#L134
    https://github.com/ethereum/cpp-ethereum/blob/develop/libdevcrypto/CryptoPP.cpp#L36
    """
    s1 = ""
    key = ""
    hash_blocksize = 64
    reps = ((key_len + 7) * 8) / (hash_blocksize * 8)
    counter = 0
    while counter <= reps:
        counter += 1
        ctx = sha256()
        ctx.update(struct.pack('>I', counter))
        ctx.update(key_material)
        ctx.update(s1)
        key += ctx.digest()
    return key[:key_len]
<commit_msg>Fixed pyelliptic import on non OS X systems<commit_after>#!/usr/bin/python
"""
https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go

https://github.com/ethereum/cpp-ethereum/blob/develop/libp2p/RLPxHandshake.cpp#L48
https://github.com/ethereum/cpp-ethereum/blob/develop/libdevcrypto/CryptoPP.cpp#L149

ECIES
http://www.cryptopp.com/wiki/Elliptic_Curve_Integrated_Encryption_Scheme
https://en.wikipedia.org/wiki/Integrated_Encryption_Scheme
"""
CURVE = 'secp256k1'
CIPHERNAME = 'aes-256-ctr'

import os
import sys
if sys.platform not in ('darwin',):
    import pyelliptic
else:
    # FIX PATH ON OS X ()
    # https://github.com/yann2192/pyelliptic/issues/11
    _openssl_lib_paths = ['/usr/local/Cellar/openssl/']
    for p in _openssl_lib_paths:
        if os.path.exists(p):
            p = os.path.join(p, os.listdir(p)[-1], 'lib')
            os.environ['DYLD_LIBRARY_PATH'] = p
            import pyelliptic
            if CIPHERNAME in pyelliptic.Cipher.get_all_cipher():
                break
    if CIPHERNAME not in pyelliptic.Cipher.get_all_cipher():
        print 'required cipher %s not available in openssl library' % CIPHERNAME
        if sys.platform == 'darwin':
            print 'use homebrew to install newer openssl'
            print '> brew install openssl'
        sys.exit(1)

import bitcoin
from sha3 import sha3_256
from hashlib import sha256
import struct
import random
import devp2p.utils as utils

hmac_sha256 = pyelliptic.hmac_sha256

class ECCx(pyelliptic.ECC):

    """
    Modified to work with raw_pubkey format used in RLPx
    and binding default curve and cipher
    """

    def __init__(self, raw_pubkey=None, raw_privkey=None):
        if raw_privkey:
            assert not raw_pubkey
            raw_pubkey = privtopub(raw_privkey)
        if raw_pubkey:
            _, pubkey_x, pubkey_y, _ = self._decode_pubkey(raw_pubkey)
        else:
            pubkey_x, pubkey_y = None, None
        pyelliptic.ECC.__init__(self, pubkey_x=pubkey_x, pubkey_y=pubkey_y,
                                raw_privkey=raw_privkey, curve=CURVE)

    @property
    def raw_pubkey(self):
        return self.pubkey_x + self.pubkey_y

    @staticmethod
    def _decode_pubkey(raw_pubkey):
        assert len(raw_pubkey) == 64
        pubkey_x = raw_pubkey[:32]
        pubkey_y = raw_pubkey[32:]
        return CURVE, pubkey_x, pubkey_y, 64

    def get_ecdh_key(self, raw_pubkey):
        "Compute public key with the local private key and returns a 256bits shared key"
        _, pubkey_x, pubkey_y, _ = self._decode_pubkey(raw_pubkey)
        key = self.raw_get_ecdh_key(pubkey_x, pubkey_y)
        assert len(key) == 32
        return key

    @property
    def raw_privkey(self):
        return self.privkey

    @staticmethod
    def encrypt(data, raw_pubkey):
        assert False
        assert len(raw_pubkey) == 64
        px, py = raw_pubkey[:32], raw_pubkey[32:]
        return ECCx.raw_encrypt(data, px, py, curve=CURVE, ciphername=CIPHERNAME)

    def _decrypt(self, data):
        return pyelliptic.ECC.decrypt(self, data, ciphername=CIPHERNAME)

    ecies_ciphername = 'aes-128-ctr'

    @classmethod
    def ecies_encrypt(cls, data, raw_pubkey):
        """
        ECIES Encrypt, where P = recipient public key is:
        1) generate r = random value
        2) generate shared-secret = kdf( ecdhAgree(r, P) )
        3) generate R = rG [same op as generating a public key]
        4) send 0x04 || R || AsymmetricEncrypt(shared-secret, plaintext) || tag

        https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/params.go#L41

        currently used by go:
        ECIES_AES128_SHA256 = &ECIESParams{
            Hash: sha256.New,
            hashAlgo: crypto.SHA256,
            Cipher: aes.NewCipher,
            BlockSize: aes.BlockSize,
            KeyLen: 16,
            }

        """
        # 1) generate r = random value
        ephem = ECCx()

        # 2) generate shared-secret = kdf( ecdhAgree(r, P) )
        key_material = ephem.raw_get_ecdh_key(pubkey_x=raw_pubkey[:32], pubkey_y=raw_pubkey[32:])
        assert len(key_material) == 32
        key = eciesKDF(key_material, 32)
        assert len(key) == 32
        key_enc, key_mac = key[:16], key[16:]

        key_mac = sha256(key_mac).digest()  # !!!
        assert len(key_mac) == 32
        # 3) generate R = rG [same op as generating a public key]
        ephem_pubkey = ephem.raw_pubkey

        # encrypt
        iv = utils.ienc(random.getrandbits(16 * 8))
        assert len(iv) == 16
        ctx = pyelliptic.Cipher(key_enc, iv, 1, cls.ecies_ciphername)
        ciphertext = ctx.ciphering(data)
        assert len(ciphertext) == len(data)

        # 4) send 0x04 || R || AsymmetricEncrypt(shared-secret, plaintext) || tag
        msg = chr(0x04) + ephem_pubkey + iv + ciphertext

        # the MAC of a message (called the tag) as per SEC 1, 3.5.
        # https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go#L162
        tag = hmac_sha256(key_mac, msg[1 + 64:])
        assert len(tag) == 32
        msg += tag

        assert len(msg) == 1 + 64 + 16 + 32 + len(data) == 113 + len(data)
        return msg

    def ecies_decrypt(self, data):
        """
        Decrypt data with ECIES method using the local private key

        ECIES Decrypt (performed by recipient):
        1) generate shared-secret = kdf( ecdhAgree(myPrivKey, msg[1:65]) )
        2) verify tag
        3) decrypt

        ecdhAgree(r, recipientPublic) == ecdhAgree(recipientPrivate, R)
        [where R = r*G, and recipientPublic = recipientPrivate*G]

        """
        assert data[0] == chr(0x04)

        #  1) generate shared-secret = kdf( ecdhAgree(myPrivKey, msg[1:65]) )
        _shared = data[1:1 + 64]
        # FIXME, check that _shared_pub is a valid one (on curve)

        key_material = self.raw_get_ecdh_key(pubkey_x=_shared[:32], pubkey_y=_shared[32:])
        assert len(key_material) == 32
        key = eciesKDF(key_material, 32)
        assert len(key) == 32
        key_enc, key_mac = key[:16], key[16:]

        key_mac = sha256(key_mac).digest()
        assert len(key_mac) == 32

        tag = data[-32:]
        assert len(tag) == 32

        # 2) verify tag
        if not pyelliptic.equals(hmac_sha256(key_mac, data[1 + 64:- 32]), tag):
            raise RuntimeError("Fail to verify data")

        # 3) decrypt
        blocksize = pyelliptic.OpenSSL.get_cipher(self.ecies_ciphername).get_blocksize()
        iv = data[1 + 64:1 + 64 + blocksize]
        assert len(iv) == 16
        ciphertext = data[1 + 64 + blocksize:- 32]
        assert 1 + len(_shared) + len(iv) + len(ciphertext) + len(tag) == len(data)
        ctx = pyelliptic.Cipher(key_enc, iv, 0, self.ecies_ciphername)
        return ctx.ciphering(ciphertext)

    def sign(self, data):
        """
        pyelliptic.ECC.sign is DER-encoded
        https://bitcoin.stackexchange.com/questions/12554
        """
        signature = ecdsa_sign(data, self.raw_privkey)
        assert len(signature) == 65
        return signature

    def verify(self, signature, message):
        assert len(signature) == 65
        return ecdsa_verify(self.raw_pubkey, signature, message)

def lzpad32(x):
    return '\x00' * (32 - len(x)) + x

def _encode_sig(v, r, s):
    assert isinstance(v, (int, long))
    assert v in (27, 28)
    vb, rb, sb = chr(v - 27), bitcoin.encode(r, 256), bitcoin.encode(s, 256)
    return lzpad32(rb) + lzpad32(sb) + vb

def _decode_sig(sig):
    return ord(sig[64]) + 27, bitcoin.decode(sig[0:32], 256), bitcoin.decode(sig[32:64], 256)

def ecdsa_verify(pubkey, signature, message):
    assert len(signature) == 65
    assert len(pubkey) == 64
    return bitcoin.ecdsa_raw_verify(message, _decode_sig(signature), pubkey)
verify = ecdsa_verify

def ecdsa_sign(message, privkey):
    s = _encode_sig(*bitcoin.ecdsa_raw_sign(message, privkey))
    return s

sign = ecdsa_sign

def ecdsa_recover(message, signature):
    assert len(signature) == 65
    pub = bitcoin.ecdsa_raw_recover(message, _decode_sig(signature))
    assert pub, 'pubkey could not be recovered'
    pub = bitcoin.encode_pubkey(pub, 'bin_electrum')
    assert len(pub) == 64
    return pub
recover = ecdsa_recover

def sha3(seed):
    return sha3_256(seed).digest()

def mk_privkey(seed):
    return sha3(seed)

def privtopub(raw_privkey):
    raw_pubkey = bitcoin.encode_pubkey(bitcoin.privtopub(raw_privkey), 'bin_electrum')
    assert len(raw_pubkey) == 64
    return raw_pubkey

def encrypt(data, raw_pubkey):
    """
    Encrypt data with ECIES method using the public key of the recipient.
    """
    assert len(raw_pubkey) == 64
    return ECCx.encrypt(data, raw_pubkey)

def eciesKDF(key_material, key_len):
    """
    interop w/go ecies implementation

    for sha3, blocksize is 136 bytes
    for sha256, blocksize is 64 bytes

    NIST SP 800-56a Concatenation Key Derivation Function (see section 5.8.1).

    https://github.com/ethereum/go-ethereum/blob/develop/crypto/ecies/ecies.go#L134
    https://github.com/ethereum/cpp-ethereum/blob/develop/libdevcrypto/CryptoPP.cpp#L36
    """
    s1 = ""
    key = ""
    hash_blocksize = 64
    reps = ((key_len + 7) * 8) / (hash_blocksize * 8)
    counter = 0
    while counter <= reps:
        counter += 1
        ctx = sha256()
        ctx.update(struct.pack('>I', counter))
        ctx.update(key_material)
        ctx.update(s1)
        key += ctx.digest()
    return key[:key_len]
<|endoftext|>


<commit_before>from flask import Flask

app = Flask(__name__)

@app.route('/')
def index():
    return 'Hi!'

<commit_msg>Add dhondt endpoint<commit_after>from flask import Flask, request, jsonify, render_template

from . import dhondt

app = Flask(__name__)

@app.route('/')
def index():
    return 'Hi!'

@app.route('/dhondt')
def calculate_dhondt():
    votes = request.args.getlist('v', type=int)
    result = dhondt(votes, 3)
    return jsonify(result=result)<|endoftext|>


<commit_before>#!/usr/bin/python
# Created by Christopher Nguyen
# 5/17/2010

#import modules
import time
import sys, os, traceback, optparse
import numpy as np
import scipy as sp

#HBCAN utilities
#sys.path.append('/home/cnguyen/scripts/HBCAN/utils')
#from io import get_coord_4D_to_3D

class tensor(object):
    """
    Tensor object that when initialized calculates single self diffusion tensor[1]_ in
    each voxel using selected fitting algorithm (DEFAULT: weighted least squares[1]_)

    Requires a given gradient table, b value for each diffusion-weighted gradient vector,
    and image data given all as numpy ndarrays.

    Key Methods
    -----------
    evals : ndarray (X,Y,Z,EV1,EV2,EV3)Z
        Returns cached eigenvalues of self diffusion tensor [1]_ for given index.
    evecs : ndarray (X,Y,Z,EV1x,EV1y,EV1z,EV2x,EV2y,EV2z,EV3x,EV3y,EV3z)
        Returns cached associated eigenvector of self diffusion tensor [1]_ for given
        index.
    FA : ndarray
        Calculates fractional anisotropy [2]_ for given index.
    ADC : ndarray
        Calculates apparent diffusion coefficient or mean diffusitivity [2]_ for given
        index.

    References
    ----------
    ..    [1] Basser, P.J., Mattiello, J., LeBihan, D., 1994. Estimation of the effective
        self-diffusion tensor from the NMR spin echo. J Magn Reson B 103, 247-254.
    ..    [2] Basser, P., Pierpaoli, C., 1996. Microstructural and physiological features
        of tissues elucidated by quantitative-diffusion-tensor MRI. Journal of Magnetic
        Resonance 111, 209-219.

    """
    def _getshape(self):
        pass

    def _getndim(self):
        pass

    def __getitem__(self,index):
        pass

    def __init__(self, data, grad_table, b_values, mask=None,thresh=5,verbose=False):
        eig_decomp, design_mat = WLS_fit(data,grad_table,b_values,verbose=verbose,mask=mask,thresh=thresh)
        self.evals = eig_decomp[:,:,:,0:3]
        self.evecs = eig_decomp[:,:,:,3:12]
        self.prime_evec = eig_decomp[:,:,:,3:6]
        self.adc = self.calc_adc()
        self.fa = self.calc_fa()

    def calc_D(self):
        D = np.dot(np.dot(Q,delta),np.linalg.pinv(Q))
        return D
        pass

    def calc_adc(self):
        #adc = (ev1+ev2+ev3)/3
        return (self.evals[:,:,:,0] + self.evals[:,:,:,1] + self.evals[:,:,:,2]) / 3

    def calc_fa(self):
        adc = self.calc_adc()
        ev1 = self.evals[:,:,:,0]
        ev2 = self.evals[:,:,:,1]
        ev3 = self.evals[:,:,:,2]
        ss_ev = ev1**2+ev2**2+ev3**2

        fa = np.zeros(ev1.shape,dtype='float32') #'int16')
        fa = np.sqrt( 1.5 * ( (ev1-adc)**2+(ev2-adc)**2+(ev3-adc)**2 ) / ss_ev )
        fa[ss_ev == 0] = 0
        return fa

def WLS_fit (data,gtab,bval,mask=None,thresh=None,verbose=False):
    """
    Computes weighted least squares (WLS) fit to calculate self-diffusion tensor.
    (Basser et al., 1994a)

    Parameters
    ----------
    data : ndarray (X,Y,Z,g)
        The image data as a numpy ndarray.
    gtab : ndarray (3,g)
        Diffusion gradient table found in DICOM header as a numpy ndarray.
    bval : ndarray (g,1)
        Diffusion weighting factor b for each vector in gtab.
    mask : ndarray (X,Y,Z)
        Mask that excludes fit in voxels where mask == 0
    thresh : integer (0,np.max(data))
        Simple threshold to exclude fit in voxels where b0 < thresh
    verbose : boolean
        Boolean to indicate verbose output such as timing.

    """
    start_time = time.time()
    ####main part of code
    dims = data.shape

    if mask != None and dims[0:3] != mask.shape :
        raise ValueError('Data image and mask MUST have same 3D volume shape!')

    ###Create log of signal and reshape it to be x:y:z by grad
    fit_dim = (dims[0]*dims[1]*dims[2],dims[3])
    data = data.reshape(fit_dim) #direct reshape for some reason does not work
    #log_s = np.zeros(fit_dim,dtype='float32')
    #log_s_ols = np.zeros(fit_dim,dtype='int16')

    # Y matrix from Chris' paper
    # enforcing positive values to allow for natural log
    data[data <= 0] = 1
    # need to set this seperately to enforce int16 precision
    # set precision to 3 significant figures...to save memory
    # instead of later calculating it with log_s_ols
    log_s = np.int16(np.log(data) * 1000)

    ###Construct design matrix
    #For DTI this is the so called B matrix
    # X matrix from Chris' paper
    B = design_matrix(gtab,bval) # [g by 7]

    ###Weighted Least Squares (WLS) to solve "linear" regression
    # Y hat OLS from Chris' paper
    #  ( [x*y*z by g] [g by 7] [7 by g ] ) = [x*y*z by g]
    log_s_ols = np.int16( np.dot(log_s, np.dot(B, np.linalg.pinv(B))) )
    #del log_s #freeing up memory

    #Setting these arrays later to allow the previous step to have all memory
    #fit_data = np.zeros(fit_dim,dtype='int16') #original data is int16
    eig_decomp = np.zeros((fit_dim[0],12),dtype='float32')#'int16')

    time_diff = list((0,0))
    time_iter = time.time()
    # This step is because we cannot vectorize diagonal vector and tensor fit
    for i in range(np.size(log_s_ols,axis=0)):
        #Check every 1 slices
        if verbose and i % (dims[0]*dims[1]*1) == 0:
            slice = i/dims[0]/dims[1]+1.
            time_diff.append(time.time()-time_iter)
            min = np.mean(time_diff[2:len(time_diff)])/60.0/5*(dims[2]-slice)
            sec = np.round((min - np.fix(min)) * 60.0/5)
            min = np.fix(min)
            percent = 100.*slice/dims[2]
            print str(np.round(percent)) + '% ... time left ' + str(min) + ' MIN ' + str(sec) + ' SEC'
            time_iter=time.time()

        if mask != None and mask.flat[i] == 0:
            continue

        if thresh != None and np.exp(log_s[i,0]/1000.) < thresh:
            continue

        #Split up weighting vector into little w to perform pinv
        w = np.exp(log_s_ols[i,:]/1000.)[:,np.newaxis]

        #pointwise broadcasting to avoid diagonal matrix multiply!
        D = np.dot(np.linalg.pinv(B*w), w.ravel()*log_s[i,:]/1000.) #np.log(data[:,i]))
        #fit_data[i,:] = np.int16(np.exp(np.dot(B,D)))

        ###Obtain eigenvalues and eigenvectors
        eig_decomp[i,:] = decompose_tensor(D[0:6],scale=1)

    #clear variables not needed to save memory
    del log_s_ols
    del log_s

    # Reshape the output
    eig_decomp = eig_decomp.reshape((dims[0],dims[1],dims[2],12))

    #Report how long it took to make the fit
    if verbose:
        min = (time.time() - start_time) / 60.0
        sec = (min - np.fix(min)) * 60.0
        print 'TOTAL TIME: ' + str(np.fix(min)) + ' MIN ' + str(np.round(sec)) + ' SEC'

    return(eig_decomp, B)

def decompose_tensor(D,scale=1):
    """
    Computes tensor eigen decomposition to calculate eigenvalues and eigenvectors of self-diffusion tensor. Assumes D has units on order of ~ 10^-4 mm^2/s
    (Basser et al., 1994a)

    Parameters
    ----------
    D : ndarray (X,Y,Z,g)
        The six unique diffusitivities (Dxx, Dyy,Dzz,Dxy,Dxz,Dyz)
    scale : integer range(1,N)
        Simple scaling parameter since diffusitivities are small.

    """
    tensor = np.zeros((3,3))
    tensor[0,0] = D[0]  #Dxx
    tensor[1,1] = D[1]  #Dyy
    tensor[2,2] = D[2]  #Dzz
    tensor[1,0] = D[3]  #Dxy
    tensor[2,0] = D[4]  #Dxz
    tensor[2,1] = D[5]  #Dyz
    tensor[0,1] = tensor[1,0] #Dyx
    tensor[0,2] = tensor[2,0] #Dzx
    tensor[1,2] = tensor[2,1] #Dzy

    #outputs multiplicity as well so need to unique
    eigenvals, eigenvecs = np.linalg.eig(tensor)

    if np.size(eigenvals) != 3:
        raise ValueError('not 3 eigenvalues : ' + str(eigenvals))

    #need to sort the eigenvalues and associated eigenvectors
    eigenvecs = eigenvecs[:,eigenvals.argsort()[::-1]]
    eigenvals.sort() #very fast
    eigenvals = eigenvals[::-1]

    #Forcing negative eigenvalues to 0
    eigenvals[eigenvals <0] = 0
    # b ~ 10^3 s/mm^2 and D ~ 10^-4 mm^2/s
    # eigenvecs: each vector is columnar

    eig_params = np.concatenate((eigenvals,eigenvecs.T.flat[:]))*scale

    return(eig_params)

def design_matrix(gtab,bval,dtype='float32'):
    """
    Constructs design matrix for DTI weighted least squares or least squares fitting.
    (Basser et al., 1994a)

    Parameters
    ----------
    gtab : ndarray (3,g)
        Diffusion gradient table found in DICOM header as a numpy ndarray.
    bval : ndarray (g,1)
        Diffusion weighting factor b for each vector in gtab.
    dtype : string
        Parameter to control the dtype of returned designed matrix

    """

    B = np.zeros((bval.size,7),dtype=dtype)
    G = gtab

    if gtab.shape[1] != bval.shape[0] :
        print 'Gradient table size is not consistent with bval vector... could be b/c of b0 images'
        print 'Will try to set nonzero bval index with gradient table to construct B matrix'

        G = np.zeros((3,np.size(bval)))
        G[:,np.where(bval > 0)]=gtab

    B[:,0] = G[0,:]*G[0,:]*1.*bval   #Bxx
    B[:,1] = G[1,:]*G[1,:]*1.*bval   #Byy
    B[:,2] = G[2,:]*G[2,:]*1.*bval   #Bzz
    B[:,3] = G[0,:]*G[1,:]*2.*bval   #Bxy
    B[:,4] = G[0,:]*G[2,:]*2.*bval   #Bxz
    B[:,5] = G[1,:]*G[2,:]*2.*bval   #Byz
    B[:,6] = np.ones(np.size(bval),dtype=dtype)

    #Need to return [g by 7]
    return -B

def save_scalar_maps(scalar_maps, img=None, coordmap=None):
    #for io of writing and reading nifti images
    from nipy import load_image, save_image
    from nipy.core.api import fromarray #data --> image

    #For writing out with save_image with appropriate affine matrix
    if img != None:
        coordmap = get_coord_4D_to_3D(img.affine)
        header = img.header.copy()

    ###Save scalar maps if requested
    print ''
    print 'Saving t2di map ... '+out_root+'_t2di.nii.gz'

    #fyi the dtype flag for save image does not appear to work right now...
    t2di_img = fromarray(data[:,:,:,0],'ijk','xyz',coordmap=coordmap)
    if img != []:
        t2di_img.header = header
    save_image(t2di_img,out_root+'_t2di.nii.gz',dtype=np.int16)

    scalar_fnames = ('ev1','ev2','ev3','adc','fa','ev1p','ev1f','ev1s')
    for i in range(np.size(scalar_maps,axis=3)):
        #need to send in 4 x 4 affine matrix for 3D image not 5 x 5 from original 4D image
        print 'Saving '+ scalar_fnames[i] + ' map ... '+out_root+'_'+scalar_fnames[i]+'.nii.gz'
        scalar_img = fromarray(np.int16(scalar_maps[:,:,:,i]),'ijk' ,'xyz',coordmap=coordmap)
        if img != []:
            scalar_img.header = header
        save_image(scalar_img,out_root+'_'+scalar_fnames[i]+'.nii.gz',dtype=np.int16)

    print ''
    print 'Saving D = [Dxx,Dyy,Dzz,Dxy,Dxz,Dyz] map ... '+out_root+'_self_diffusion.nii.gz'
    #Saving 4D matrix holding diffusion coefficients
    if img != [] :
        coordmap = img.coordmap
        header = img.header.copy()
    tensor_img = fromarray(tensor_data,'ijkl','xyzt',coordmap=coordmap)
    tensor_img.header = header
    save_image(tensor_img,out_root+'_self_diffusion.nii.gz',dtype=np.int16)

    print

    return
<commit_msg>RF - introducing maskedview into dti.py to allow for efficient memory handling<commit_after>#!/usr/bin/python
# Created by Christopher Nguyen
# 5/17/2010

#import modules
import time
import sys, os, traceback, optparse
import numpy as np
import scipy as sp

#dipy modules
from dipy.core.maskedview import MaskedView
from dipy.core.memory import memory

class tensor(object):
    """
    Tensor object that when initialized calculates single self diffusion tensor[1]_ in
    each voxel using selected fitting algorithm (DEFAULT: weighted least squares[1]_)

    Requires a given gradient table, b value for each diffusion-weighted gradient vector,
    and image data given all as numpy ndarrays.

    Parameters
    ----------
    data : ndarray (X,Y,Z,g)
        The image data as a numpy ndarray.
    gtab : ndarray (3,g)
        Diffusion gradient table found in DICOM header as a numpy ndarray.
    bval : ndarray (g,1)
        Diffusion weighting factor b for each vector in gtab.
    mask : ndarray (X,Y,Z)
        Mask that excludes fit in voxels where mask == 0
    thresh : integer (0,np.max(data))
        Simple threshold to exclude fit in voxels where b0 < thresh

    Key Methods
    -----------
    evals : ndarray (X,Y,Z,EV1,EV2,EV3)Z
        Returns cached eigenvalues of self diffusion tensor [1]_ for given index.
    evecs : ndarray (X,Y,Z,EV1x,EV1y,EV1z,EV2x,EV2y,EV2z,EV3x,EV3y,EV3z)
        Returns cached associated eigenvector of self diffusion tensor [1]_ for given
        index.
    FA : ndarray
        Calculates fractional anisotropy [2]_ for given index.
    ADC : ndarray
        Calculates apparent diffusion coefficient or mean diffusitivity [2]_ for given
        index.

    References
    ----------
    ..    [1] Basser, P.J., Mattiello, J., LeBihan, D., 1994. Estimation of the effective
        self-diffusion tensor from the NMR spin echo. J Magn Reson B 103, 247-254.
    ..    [2] Basser, P., Pierpaoli, C., 1996. Microstructural and physiological features
        of tissues elucidated by quantitative-diffusion-tensor MRI. Journal of Magnetic
        Resonance 111, 209-219.

    """
    def _getshape(self):
        pass

    def _getndim(self):
        pass

    def __getitem__(self,index):
        pass

    def __init__(self, data, grad_table, b_values, mask=None,thresh=5,verbose=False):
        if mask == None:
            mask = data[:,:,:,0]
            mask[mask<thresh] = 0

        dims = data.shape
        if mask != None and dims[0:3] != mask.shape:
            raise ValueError('Data image and mask MUST have same 3D volume shape!')

        mask4d = mask[:,:,:,np.newaxis]
        for i in range(dims[3]-1):
            mask4d = np.concatenate((mask4d,mask[:,:,:,np.newaxis]),axis=3)

        #no support from maskedview to handle 3D mask for 4D data
        #   therefore input mask must be 4D
        #no maskedview.reshape
        #   need to reshape the 4d mask
        #also data needs to be masked first
        #   mask needs to be broadcasted to 4D to match and then mask data
        data = MaskedView(mask4d.reshape((np.prod(dims[0:3]),dims[3])),data[mask4d>0]) #mask[:,:,:,np.newaxis]>0])

        eig_decomp, design_mat = WLS_fit(data,grad_table,b_values,verbose=verbose)

        mask = mask.reshape(dims)
        eig_decomp = MaskedView(mask,eig_decomp)
        self.evals = eig_decomp[:,:,:,0:3]
        self.evecs = eig_decomp[:,:,:,3:12]
        self.prime_evec = eig_decomp[:,:,:,3:6]

        #this is for convenience (does not add much memory)
        self.adc = self.calc_adc()
        self.fa = self.calc_fa()
        #self.D = self.calc_D()

    def calc_D(self):
        D = np.dot(np.dot(Q,delta),np.linalg.pinv(Q))
        return D
        pass

    def calc_adc(self):
        #adc = (ev1+ev2+ev3)/3
        return (self.evals[:,:,:,0] + self.evals[:,:,:,1] + self.evals[:,:,:,2]) / 3

    def calc_fa(self):
        adc = self.calc_adc()
        ev1 = self.evals[:,:,:,0]
        ev2 = self.evals[:,:,:,1]
        ev3 = self.evals[:,:,:,2]
        ss_ev = ev1**2+ev2**2+ev3**2

        fa = np.zeros(ev1.shape,dtype='float32') #'int16')
        fa = np.sqrt( 1.5 * ( (ev1-adc)**2+(ev2-adc)**2+(ev3-adc)**2 ) / ss_ev )
        fa[ss_ev == 0] = 0
        return fa

def WLS_fit (data,gtab,bval,verbose=False):
    """
    Computes weighted least squares (WLS) fit to calculate self-diffusion tensor.
    (Basser et al., 1994a)

    Parameters
    ----------
    data : ndarray (X,Y,Z,g) OR Maskedview (X*Y*Z,g) [preferred]
        The image data as a numpy ndarray.
    gtab : ndarray (3,g)
        Diffusion gradient table found in DICOM header as a numpy ndarray.
    bval : ndarray (g,1)
        Diffusion weighting factor b for each vector in gtab.
    verbose : boolean
        Boolean to indicate verbose output such as timing.

    Returns
    -------
    eig_decomp : ndarray (X,Y,Z,12) OR Maskedview (X*Y*Z,g)
        Eigenvalues and eigenvectors from eigen decomposition of the tensor
    design_mat : ndarray (g,7)
        DTI design matrix to reconstruct fitted data if desired

    """
    start_time = time.time()
    ####main part of code
    dims = data.shape

    if len(dims) == 4:
        fit_dim = (dims[0]*dims[1]*dims[2],dims[3])
        # Y matrix from Chris' paper
        data = data.reshape(fit_dim) #direct reshape for some reason does not work
    elif len(dims) == 2:
        fit_dim = dims

    ###Create log of signal and reshape it to be x:y:z by grad
    if isinstance(data,np.ndarray):
        data[data <= 0] = 1 # enforcing positive values to allow for natural log
        ###Create log of signal and reshape it to be x:y:z by grad
        # need to set this seperately to enforce int16 precision
        # set precision to 3 significant figures...to save memory
        # instead of later calculating it with log_s_ols
        log_s = np.int16(np.log(data) * 1000)
        scale=1/1000.
    else:
        log_s = np.log(data)
        scale=1

    ###Construct design matrix
    #For DTI this is the so called B matrix
    # X matrix from Chris' paper
    B = design_matrix(gtab,bval) # [g by 7]

    ###Weighted Least Squares (WLS) to solve "linear" regression
    # Y hat OLS from Chris' paper
    #  ( [x*y*z by g] [g by 7] [7 by g ] ) = [x*y*z by g]
    log_s_ols = np.dot(log_s, np.dot(B, np.linalg.pinv(B)))
    del log_s #freeing up memory

    #handling if ndarray is sent
    if isinstance(data,np.ndarray):
        log_s_ols = np.int16(log_s_ols)

    #Setting these arrays later to allow the previous step to have all memory
    #fit_data = np.zeros(fit_dim,dtype='int16') #original data is int16
    eig_decomp = np.zeros((fit_dim[0],12),dtype='float32')#'int16')

    time_diff = list((0,0))
    time_iter = time.time()
    # This step is because we cannot vectorize diagonal vector and tensor fit
    for i in range(log_s_ols.shape[0]): #range(np.size(log_s_ols,axis=0)):
        #Check every 1 slices
        if verbose and i % (dims[0]*dims[1]*1) == 0:
            slice = i/dims[0]/dims[1]+1.
            time_diff.append(time.time()-time_iter)
            min = np.mean(time_diff[2:len(time_diff)])/60.0/5*(dims[2]-slice)
            sec = np.round((min - np.fix(min)) * 60.0/5)
            min = np.fix(min)
            percent = 100.*slice/dims[2]
            print str(np.round(percent)) + '% ... time left: ' + str(min) + ' MIN ' \
                                + str(sec) + ' SEC ... memory: ' + memory()/1024. + 'MB'
            time_iter=time.time()

        #Split up weighting vector into little w to perform pinv
        w = np.exp(log_s_ols[i,:]*scale)[:,np.newaxis]

        #pointwise broadcasting to avoid diagonal matrix multiply!
        D = np.dot(np.linalg.pinv(B*w), w.ravel()*np.log(data[:,i])*scale) #log_s[i,:]

        ###Obtain eigenvalues and eigenvectors
        eig_decomp[i,:] = decompose_tensor(D[0:6],scale=1)

    #clear variables not needed to save memory
    del log_s_ols

    # Reshape the output
    if len(dims) == 4:
        eig_decomp = eig_decomp.reshape((dims[0],dims[1],dims[2],12))

    #Report how long it took to make the fit
    if verbose:
        min = (time.time() - start_time) / 60.0
        sec = (min - np.fix(min)) * 60.0
        print 'TOTAL TIME: ' + str(np.fix(min)) + ' MIN ' + str(np.round(sec)) + ' SEC'

    return(eig_decomp, B)

def decompose_tensor(D,scale=1):
    """
    Computes tensor eigen decomposition to calculate eigenvalues and eigenvectors of
    self-diffusion tensor. Assumes D has units on order of ~ 10^-4 mm^2/s
    (Basser et al., 1994a)

    Parameters
    ----------
    D : ndarray (X,Y,Z,g)
        The six unique diffusitivities (Dxx, Dyy,Dzz,Dxy,Dxz,Dyz)
    scale : integer range(1,N)
        Simple scaling parameter since diffusitivities are small.

    """
    tensor = np.zeros((3,3))
    tensor[0,0] = D[0]  #Dxx
    tensor[1,1] = D[1]  #Dyy
    tensor[2,2] = D[2]  #Dzz
    tensor[1,0] = D[3]  #Dxy
    tensor[2,0] = D[4]  #Dxz
    tensor[2,1] = D[5]  #Dyz
    tensor[0,1] = tensor[1,0] #Dyx
    tensor[0,2] = tensor[2,0] #Dzx
    tensor[1,2] = tensor[2,1] #Dzy

    #outputs multiplicity as well so need to unique
    eigenvals, eigenvecs = np.linalg.eig(tensor)

    if np.size(eigenvals) != 3:
        raise ValueError('not 3 eigenvalues : ' + str(eigenvals))

    #need to sort the eigenvalues and associated eigenvectors
    eigenvecs = eigenvecs[:,eigenvals.argsort()[::-1]]
    eigenvals.sort() #very fast
    eigenvals = eigenvals[::-1]

    #Forcing negative eigenvalues to 0
    eigenvals[eigenvals <0] = 0
    # b ~ 10^3 s/mm^2 and D ~ 10^-4 mm^2/s
    # eigenvecs: each vector is columnar

    eig_params = np.concatenate((eigenvals,eigenvecs.T.flat[:]))*scale

    return(eig_params)

def design_matrix(gtab,bval,dtype='float32'):
    """
    Constructs design matrix for DTI weighted least squares or least squares fitting.
    (Basser et al., 1994a)

    Parameters
    ----------
    gtab : ndarray (3,g)
        Diffusion gradient table found in DICOM header as a numpy ndarray.
    bval : ndarray (g,1)
        Diffusion weighting factor b for each vector in gtab.
    dtype : string
        Parameter to control the dtype of returned designed matrix

    """

    B = np.zeros((bval.size,7),dtype=dtype)
    G = gtab

    if gtab.shape[1] != bval.shape[0] :
        print 'Gradient table size is not consistent with bval vector... could be b/c \
               of b0 images'
        print 'Will try to set nonzero bval index with gradient table to construct \
               B matrix'

        G = np.zeros((3,np.size(bval)))
        G[:,np.where(bval > 0)]=gtab

    B[:,0] = G[0,:]*G[0,:]*1.*bval   #Bxx
    B[:,1] = G[1,:]*G[1,:]*1.*bval   #Byy
    B[:,2] = G[2,:]*G[2,:]*1.*bval   #Bzz
    B[:,3] = G[0,:]*G[1,:]*2.*bval   #Bxy
    B[:,4] = G[0,:]*G[2,:]*2.*bval   #Bxz
    B[:,5] = G[1,:]*G[2,:]*2.*bval   #Byz
    B[:,6] = np.ones(np.size(bval),dtype=dtype)

    #Need to return [g by 7]
    return -B

def save_scalar_maps(scalar_maps, img=None, coordmap=None):
    #for io of writing and reading nifti images
    from nipy import load_image, save_image
    from nipy.core.api import fromarray #data --> image

    #For writing out with save_image with appropriate affine matrix
    if img != None:
        coordmap = get_coord_4D_to_3D(img.affine)
        header = img.header.copy()

    ###Save scalar maps if requested
    print ''
    print 'Saving t2di map ... '+out_root+'_t2di.nii.gz'

    #fyi the dtype flag for save image does not appear to work right now...
    t2di_img = fromarray(data[:,:,:,0],'ijk','xyz',coordmap=coordmap)
    if img != []:
        t2di_img.header = header
    save_image(t2di_img,out_root+'_t2di.nii.gz',dtype=np.int16)

    scalar_fnames = ('ev1','ev2','ev3','adc','fa','ev1p','ev1f','ev1s')
    for i in range(np.size(scalar_maps,axis=3)):
        #need to send in 4 x 4 affine matrix for 3D image not 5 x 5 from original 4D image
        print 'Saving '+ scalar_fnames[i] + ' map ... '+out_root+'_'+scalar_fnames[i]+'.nii.gz'
        scalar_img = fromarray(np.int16(scalar_maps[:,:,:,i]),'ijk' ,'xyz',coordmap=coordmap)
        if img != []:
            scalar_img.header = header
        save_image(scalar_img,out_root+'_'+scalar_fnames[i]+'.nii.gz',dtype=np.int16)

    print ''
    print 'Saving D = [Dxx,Dyy,Dzz,Dxy,Dxz,Dyz] map ... '+out_root+'_self_diffusion.nii.gz'
    #Saving 4D matrix holding diffusion coefficients
    if img != [] :
        coordmap = img.coordmap
        header = img.header.copy()
    tensor_img = fromarray(tensor_data,'ijkl','xyzt',coordmap=coordmap)
    tensor_img.header = header
    save_image(tensor_img,out_root+'_self_diffusion.nii.gz',dtype=np.int16)

    print

    return
<|endoftext|>


<commit_before>import pygame
import numpy as np

class Graph():

	def __init__(self, display, sensor):
		self.display = display
		self.surface = self.display.screen
		self.width = self.display.width
		self.height = self.display.height / 3
		self.sensor = sensor
		self.surface = pygame.Surface((self.width, self.height))
		self.drawBG()
		self.curX = 0
		self.x = 0
		self.colors = (
		(128, 128, 128),
		(128, 0, 0),
		(128, 128, 0),
		(0, 128, 0),
		(0, 128, 128),
		(255, 255, 255),
		(255, 255, 255),
		(128, 255, 255),
		(255, 128, 255),
		(255, 255, 128),
		(255, 255, 255),
		(255, 255, 255),
		(0, 128, 128)
		)

		#self.sensor.channels[1].onUpdate = self.updateBuffer

	def render(self, channels=None):
		self.checkBounds()
		for i, channel in enumerate(self.sensor.channels):
			self.renderChannel(channel, self.colors[i])
		#self.renderChannel(self.sensor.channels[1], self.colors[1])
		#self.renderChannel(self.sensor.channels[2], self.colors[2])
		#self.renderChannel(self.sensor.channels[3], self.colors[3])

		if self.x % 10 == 0:
			self.display.screen.blit(self.surface, (0, self.display.height * 2 / 3))
			pygame.display.flip()

	def renderChannel(self, channel, color):
		min = channel.min
		max = channel.max

		value = channel.getValue()
		y = int(self.map_value(value, min, max, self.height, 0))
		self.surface.set_at((self.x, y), color)

		avg = channel.getBufferAvg()
		yAvg = int(self.map_value(avg, min, max, self.height, 0))
		#self.surface.set_at((self.x, yAvg), color)

		rng = channel.getRng()
		yRng = int(self.map_value(rng, 0, 100, self.height, 0))
		#self.surface.set_at((self.x, yRng), [x / 2 for x in color])

	def checkBounds(self, scroll=False):
		if self.x >= (self.width - 1):
			if scroll:
				self.surface.scroll(-1,0)
				pygame.draw.rect(self.surface, (0,0,0), (self.width-1, 0, 2, self.height))
				self.x -= 1
			else:
				self.x = 0
				self.drawBG()
		else:
			self.x += 1

	def drawBG(self):
		self.surface.fill((0, 0, 0))
		for i in range(10):
			y = int(self.map_value(i, 0, 10, self.height, 0))
			pygame.draw.line(self.surface, (32, 32, 32), (0, y), (self.width, y))

	def map_value(self, value, in_min, in_max, out_min, out_max):
		return (value - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

	"""
	def updateBuffer(self, channel):

		min = 1200
		max = 2000
		buffer = channel.npBuffer
		buffer = channel.smoothed
		std = np.std(channel.npBuffer) + 1200
		var = np.var(channel.npBuffer) * 0.1 + 1200
		avg = np.average(channel.npBuffer)
		mean = np.mean(channel.npBuffer)
		quant = 50
		val = int(avg / quant) * quant

		#for i in np.nditer(buffer):
		for i in range(buffer.shape[0]):
			self.checkBounds()
			value = buffer[i]
			y = int(self.map_value(value, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[1])

			#smooth = smoothed[i]
			#y = int(self.map_value(smooth, min, max, self.height, 0))
			#self.surface.set_at((self.x, y), self.colors[1])

			y = int(self.map_value(var, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[2])
			y = int(self.map_value(std, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[3])
			y = int(self.map_value(val, min, max, self.height, 0))
			#self.surface.set_at((self.x, y), self.colors[4])

			self.x += 1
		pygame.display.flip()

	def draw(self, channel):
		pass

	def setValues(self, values):
		self.values = values

	def render(self, channels=None):
		self.checkBounds()

		self.renderChannel(self.sensor.channels[0], self.colors[1])
		self.renderChannel(self.sensor.channels[1], self.colors[2])
		self.renderChannel(self.sensor.channels[2], self.colors[3])

		pygame.display.flip()
		#self.surface.blit(self.surface, (0,0))

	"""<commit_msg>Start Cleanup, NumPy Smooth<commit_after>import pygame
import numpy as np

class Graph():

	def __init__(self, display, sensor):
		self.display = display
		self.surface = self.display.screen
		self.width = self.display.width
		self.height = self.display.height / 3
		self.sensor = sensor
		self.surface = pygame.Surface((self.width, self.height))
		self.drawBG()
		self.curX = 0
		self.x = 0
		self.colors = (
		(128, 128, 128),
		(128, 0, 0),
		(128, 128, 0),
		(0, 128, 0),
		(0, 128, 128),
		(255, 255, 255),
		(255, 255, 255),
		(128, 255, 255),
		(255, 128, 255),
		(255, 255, 128),
		(255, 255, 255),
		(255, 255, 255),
		(0, 128, 128)
		)

		#self.sensor.channels[1].onUpdate = self.updateBuffer

	def render(self, channels=None):
		self.checkBounds()
		for i, channel in enumerate(self.sensor.channels):
			if channels is not None or i in channels:
				self.renderChannel(channel, self.colors[i])
		#self.renderChannel(self.sensor.channels[1], self.colors[1])
		#self.renderChannel(self.sensor.channels[2], self.colors[2])
		#self.renderChannel(self.sensor.channels[3], self.colors[3])

		if self.x % 10 == 0:
			self.display.screen.blit(self.surface, (0, self.display.height * 2 / 3))
			pygame.display.flip()

	def renderChannel(self, channel, color):
		min = channel.min
		max = channel.max

		value = channel.getValue()
		y = int(self.map_value(value, min, max, self.height, 0))
		self.surface.set_at((self.x, y), color)

		avg = channel.getBufferAvg()
		yAvg = int(self.map_value(avg, min, max, self.height, 0))
		#self.surface.set_at((self.x, yAvg), color)

		rng = channel.getRng()
		yRng = int(self.map_value(rng, 0, 100, self.height, 0))
		#self.surface.set_at((self.x, yRng), [x / 2 for x in color])

	def checkBounds(self, scroll=False):
		if self.x >= (self.width - 1):
			if scroll:
				self.surface.scroll(-1,0)
				pygame.draw.rect(self.surface, (0,0,0), (self.width-1, 0, 2, self.height))
				self.x -= 1
			else:
				self.x = 0
				self.drawBG()
		else:
			self.x += 1

	def drawBG(self):
		self.surface.fill((0, 0, 0))
		for i in range(10):
			y = int(self.map_value(i, 0, 10, self.height, 0))
			pygame.draw.line(self.surface, (32, 32, 32), (0, y), (self.width, y))

	def map_value(self, value, in_min, in_max, out_min, out_max):
		return (value - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

	"""
	def updateBuffer(self, channel):

		min = 1200
		max = 2000
		buffer = channel.npBuffer
		buffer = channel.smoothed
		std = np.std(channel.npBuffer) + 1200
		var = np.var(channel.npBuffer) * 0.1 + 1200
		avg = np.average(channel.npBuffer)
		mean = np.mean(channel.npBuffer)
		quant = 50
		val = int(avg / quant) * quant

		#for i in np.nditer(buffer):
		for i in range(buffer.shape[0]):
			self.checkBounds()
			value = buffer[i]
			y = int(self.map_value(value, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[1])

			#smooth = smoothed[i]
			#y = int(self.map_value(smooth, min, max, self.height, 0))
			#self.surface.set_at((self.x, y), self.colors[1])

			y = int(self.map_value(var, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[2])
			y = int(self.map_value(std, min, max, self.height, 0))
			self.surface.set_at((self.x, y), self.colors[3])
			y = int(self.map_value(val, min, max, self.height, 0))
			#self.surface.set_at((self.x, y), self.colors[4])

			self.x += 1
		pygame.display.flip()

	def draw(self, channel):
		pass

	def setValues(self, values):
		self.values = values

	def render(self, channels=None):
		self.checkBounds()

		self.renderChannel(self.sensor.channels[0], self.colors[1])
		self.renderChannel(self.sensor.channels[1], self.colors[2])
		self.renderChannel(self.sensor.channels[2], self.colors[3])

		pygame.display.flip()
		#self.surface.blit(self.surface, (0,0))

	"""<|endoftext|>


<commit_before>
from SoftLayer import Client, API_PUBLIC_ENDPOINT
from oslo.config import cfg

from jumpgate.common.sl.auth import get_auth, get_token_details
from jumpgate.common.sl.errors import handle_softlayer_errors
from SoftLayer import SoftLayerAPIError

opts = [
    cfg.StrOpt('endpoint', default=API_PUBLIC_ENDPOINT),
    cfg.StrOpt('catalog_template_file', default='identity.templates'),
]

cfg.CONF.register_opts(opts, group='softlayer')
cfg.CONF.register_opts([cfg.StrOpt('admin_token', secret=True,
                                   default='ADMIN')], group='DEFAULT')

def hook_get_client(req, resp, kwargs):
    client = Client(endpoint_url=cfg.CONF['softlayer']['endpoint'])
    client.auth = None
    req.env['tenant_id'] = None
    req.env['sl_client'] = client

    if req.headers.get('X-AUTH-TOKEN'):
        if 'X-AUTH-TOKEN' in req.headers:
            admin_token = cfg.CONF['DEFAULT']['admin_token']
            if admin_token and admin_token == req.headers.get('X-AUTH-TOKEN'):
                # authn for jumpgate api, but no mapping into slapi
                return
            tenant_id = kwargs.get('tenant_id',
                                   req.headers.get('X-AUTH-PROJECT-ID'))
            token_details = get_token_details(req.headers['X-AUTH-TOKEN'],
                                              tenant_id=tenant_id)

            client.auth = get_auth(token_details)

            req.env['tenant_id'] = token_details['tenant_id']

def add_hooks(app):
    if hook_get_client not in app.before_hooks:
        app.before_hooks.append(hook_get_client)

    app.add_error_handler(SoftLayerAPIError, handle_softlayer_errors)
<commit_msg>admin_token support<commit_after>
from SoftLayer import Client, API_PUBLIC_ENDPOINT
from oslo.config import cfg

from jumpgate.common.sl.auth import get_auth, get_token_details
from jumpgate.common.sl.errors import handle_softlayer_errors
from SoftLayer import SoftLayerAPIError

opts = [
    cfg.StrOpt('endpoint', default=API_PUBLIC_ENDPOINT),
    cfg.StrOpt('catalog_template_file', default='identity.templates'),
]

cfg.CONF.register_opts(opts, group='softlayer')
cfg.CONF.register_opts([cfg.StrOpt('admin_token', secret=True,
                                   default='ADMIN')], group='DEFAULT')

def hook_get_client(req, resp, kwargs):
    client = Client(endpoint_url=cfg.CONF['softlayer']['endpoint'])
    client.auth = None
    req.env['tenant_id'] = None
    req.env['sl_client'] = client

    if req.headers.get('X-AUTH-TOKEN'):
        if 'X-AUTH-TOKEN' in req.headers:
            admin_token = cfg.CONF['DEFAULT']['admin_token']
            if admin_token and admin_token == req.headers.get('X-AUTH-TOKEN'):
                # admin_token authenticates to Jumpgate API, but does not
                # provide SLAPI access
                return
            tenant_id = kwargs.get('tenant_id',
                                   req.headers.get('X-AUTH-PROJECT-ID'))
            token_details = get_token_details(req.headers['X-AUTH-TOKEN'],
                                              tenant_id=tenant_id)

            client.auth = get_auth(token_details)

            req.env['tenant_id'] = token_details['tenant_id']

def add_hooks(app):
    if hook_get_client not in app.before_hooks:
        app.before_hooks.append(hook_get_client)

    app.add_error_handler(SoftLayerAPIError, handle_softlayer_errors)
<|endoftext|>


<commit_before>"""
A simple log mechanism styled after PEP 282.

Retained for compatibility and should not be used.
"""

import logging

from ._log import log as _global_log

DEBUG = logging.DEBUG
INFO = logging.INFO
WARN = logging.WARN
ERROR = logging.ERROR
FATAL = logging.FATAL

log = _global_log.log
debug = _global_log.debug
info = _global_log.info
warn = _global_log.warning
error = _global_log.error
fatal = _global_log.fatal

def set_threshold(level):
    orig = _global_log.level
    _global_log.setLevel(level)
    return orig

def set_verbosity(v):
    if v <= 0:
        set_threshold(logging.WARN)
    elif v == 1:
        set_threshold(logging.INFO)
    elif v >= 2:
        set_threshold(logging.DEBUG)
<commit_msg>Add distutils.log.Log back for compatibility<commit_after>"""
A simple log mechanism styled after PEP 282.

Retained for compatibility and should not be used.
"""

import logging
import warnings

from ._log import log as _global_log

DEBUG = logging.DEBUG
INFO = logging.INFO
WARN = logging.WARN
ERROR = logging.ERROR
FATAL = logging.FATAL

log = _global_log.log
debug = _global_log.debug
info = _global_log.info
warn = _global_log.warning
error = _global_log.error
fatal = _global_log.fatal

def set_threshold(level):
    orig = _global_log.level
    _global_log.setLevel(level)
    return orig

def set_verbosity(v):
    if v <= 0:
        set_threshold(logging.WARN)
    elif v == 1:
        set_threshold(logging.INFO)
    elif v >= 2:
        set_threshold(logging.DEBUG)

class Log(logging.Logger):
    """distutils.log.Logger is deprecated, please use an alternative from `logging`."""

    def __init__(self, threshold=WARN):
        warnings.warn(Log.__doc__)  # avoid DeprecationWarning to ensure warn is shown
        super().__init__(__name__, level=threshold)

    @property
    def threshold(self):
        return self.level

    @threshold.setter
    def threshold(self, level):
        self.setLevel(level)

    warn = logging.Logger.warning
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.conf import settings
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.db.models import get_model

def aloha_init(request):
    """
    Build the javascript file which is initializing the aloha-editor
    Run the javascript code for the AlohaInput widget
    """

    links = []
    link_models = getattr(settings, 'DJALOHA_LINK_MODELS', ())
    for full_model_name in link_models:
        app_name, model_name = full_model_name.split('.')
        model = get_model(app_name, model_name)
        if model:
            links.extend(model.objects.all())

    return render_to_response(
        'djaloha/aloha_init.js',
        {'links': links},
        context_instance=RequestContext(request)
    )
<commit_msg>Fixed mimetype on aloha_init view.<commit_after># -*- coding: utf-8 -*-
from django.conf import settings
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.db.models import get_model

def aloha_init(request):
    """
    Build the javascript file which is initializing the aloha-editor
    Run the javascript code for the AlohaInput widget
    """

    links = []
    link_models = getattr(settings, 'DJALOHA_LINK_MODELS', ())
    for full_model_name in link_models:
        app_name, model_name = full_model_name.split('.')
        model = get_model(app_name, model_name)
        if model:
            links.extend(model.objects.all())

    return render_to_response(
        'djaloha/aloha_init.js',
        {'links': links},
        mimetype='text/javascript',
        context_instance=RequestContext(request)
    )
<|endoftext|>


<commit_before>#encoding: utf-8
"""
Djolar searcher mixins
"""

class DjangoSearchMixin(object):
    '''
    Make the view class searchable
    '''
    searcher_class = None

    def get_searcher_class(self):
        """
        Return the class to use for the search.
        Defaults to using `self.searcher_class`.
        """
        assert self.searcher_class is not None, (
            "'%s' should either include a `searcher_class` attribute, "
            "or override the `get_searcher_class()` method."
            % self.__class__.__name__
        )

        return self.searcher_class

    def get_searcher(self, *args, **kwargs):
        '''
        Return the search instance for building model search query object
        '''
        searcher_class = self.get_searcher_class()
        return searcher_class(*args, **kwargs)

    def get_search_queryset(self, *args, **kwargs):
        '''
        Return the queryset for result filter
        '''
        assert self.queryset is not None, (
            "'%s' should either include a `queryset` attribute, "
            "or override the `get_search_queryset()` method."
            % self.__class__.__name__
        )

        return self.queryset

    def get_queryset(self, *args, **kwargs):
        searcher = self.get_searcher()

        # Build search field
        queryQ = searcher.get_query_fields(self.request.GET)

        # Get order by field
        orderBy = searcher.get_order_fields(self.request.GET)

        # Make queryset
        queryset = self.get_search_queryset().filter(queryQ).order_by(*orderBy)

        return queryset

<commit_msg>fix PEP8<commit_after># encoding: utf-8
"""
Djolar searcher mixins
"""

class DjangoSearchMixin(object):
    '''
    Make the view class searchable
    '''
    searcher_class = None

    def get_searcher_class(self):
        """
        Return the class to use for the search.
        Defaults to using `self.searcher_class`.
        """
        assert self.searcher_class is not None, (
            "'%s' should either include a `searcher_class` attribute, "
            "or override the `get_searcher_class()` method."
            % self.__class__.__name__
        )

        return self.searcher_class

    def get_searcher(self, *args, **kwargs):
        '''
        Return the search instance for building model search query object
        '''
        searcher_class = self.get_searcher_class()
        return searcher_class(*args, **kwargs)

    def get_search_queryset(self, *args, **kwargs):
        '''
        Return the queryset for result filter
        '''
        assert self.queryset is not None, (
            "'%s' should either include a `queryset` attribute, "
            "or override the `get_search_queryset()` method."
            % self.__class__.__name__
        )

        return self.queryset

    def get_queryset(self, *args, **kwargs):
        searcher = self.get_searcher()

        # Build search field
        queryQ = searcher.get_query_fields(self.request.GET)

        # Get order by field
        orderBy = searcher.get_order_fields(self.request.GET)

        # Make queryset
        queryset = self.get_search_queryset().filter(queryQ).order_by(*orderBy)

        return queryset
<|endoftext|>


<commit_before>'''

This is a thin veneer around a pyfilesystem. It adds a few bits of
functionality:

1) Django configuration. This can go to Amazon S3 or a static
filesystem.
2) The ability to get URLs for objects stored on the filesystem.
3) The ability to create objects with a limited lifetime. A
task can garbage-collect those objects.

'''

import json
import os
import os.path
import types

from django.conf import settings

from models import FSExpirations

if hasattr(settings, 'DJFS'):
    djfs_settings = settings.DJFS
else:
    djfs_settings = {'type' : 'osfs',
                     'directory_root' : 'django-pyfs/static/django-pyfs',
                     'url_root' : '/static/django-pyfs'}

if djfs_settings['type'] == 'osfs':
    from fs.osfs import OSFS
elif djfs_settings['type'] == 's3fs':
    from fs.s3fs import S3FS
    from boto.s3.connection import S3Connection
    from boto.s3.key import Key
    s3conn = S3Connection()
else:
    raise AttributeError("Bad filesystem: "+str(djfs_settings['type']))

def get_filesystem(namespace):
    ''' Returns a pyfilesystem for static module storage.

    The file system will have two additional properties:
    1) get_url: A way to get a URL for a static file download
    2) expire: A way to expire files (so they are automatically destroyed)
    '''
    if djfs_settings['type'] == 'osfs':
        return get_osfs( namespace )
    elif djfs_settings['type'] == 's3fs':
        return get_s3fs( namespace )
    else:
        raise AttributeError("Bad filesystem: "+str(djfs_settings['type']))

def expire_objects():
    ''' Remove all obsolete objects from the file systems. Untested. '''
    objects = sorted(FSExpirations.expired(), key=lambda x:x.module)
    fs = None
    module = None
    for o in objects:
        if module != o.module:
            module = o.module
            fs = get_filesystem(module)
        if fs.exists(o.filename):
            fs.remove(o.filename)
        o.delete()

def patch_fs(fs, namespace, url_method):
    ''' Patch a filesystem object to add two methods:
          get_url returns a URL for a resource stored on that filesystem. It takes two parameters:
              filename: Which resource
              timeout: How long that resource is available for
          expire sets a timeout on how long the system should keep the resource. It takes four parameters:
              filename: Which resource
              seconds: How long we will keep it
              days: (optional) More user-friendly if a while
              expires: (optional) boolean; if set to False, we keep the resource forever.
          Without calling this method, we provide no guarantees on how long resources will stick around.
    '''
    def expire(self, filename, seconds, days=0, expires = True):
        ''' Set the lifespan of a file on the filesystem.

        filename: Name of file
        expire: False means the file will never be removed
        seconds and days give time to expiration.
        '''
        FSExpirations.create_expiration(namespace, filename, seconds, days=days, expires = expires)

    fs.expire = types.MethodType(expire, fs)
    fs.get_url = types.MethodType(url_method, fs)
    return fs

def get_osfs(namespace):
    ''' Helper method to get_filesystem for a file system on disk '''
    full_path = os.path.join(djfs_settings['directory_root'], namespace)
    if not os.path.exists(full_path):
        os.makedirs(full_path)
    osfs = OSFS(full_path)
    osfs = patch_fs(osfs, namespace, lambda self, filename, timeout=0:os.path.join(djfs_settings['url_root'], namespace, filename))
    return osfs

def get_s3fs(namespace):
    ''' Helper method to get_filesystem for a file system on S3 '''
    fullpath = namespace
    if 'prefix' in djfs_settings:
        fullpath = os.path.join(djfs_settings['prefix'], fullpath)
    s3fs = S3FS(djfs_settings['bucket'], fullpath)

    def get_s3_url(self, filename, timeout=60):
        global s3conn
        try:
            return s3conn.generate_s3_url(timeout, 'GET', bucket = djfs_settings['bucket'], key = filename)
        except: # If connection has timed out
            s3conn = S3Connection()
            return s3conn.generate_s3_url(timeout, 'GET', bucket = djfs_settings['bucket'], key = filename)

    s3fs = patch_fs(s3fs, namespace, get_s3_url)
    return s3fs

<commit_msg>Adds support for auth. Note: UNTESTED (but will be in a few minutes, so if you see this later, tested)<commit_after>'''

This is a thin veneer around a pyfilesystem. It adds a few bits of
functionality:

1) Django configuration. This can go to Amazon S3 or a static
filesystem.
2) The ability to get URLs for objects stored on the filesystem.
3) The ability to create objects with a limited lifetime. A
task can garbage-collect those objects.

'''

import json
import os
import os.path
import types

from django.conf import settings

from models import FSExpirations

if hasattr(settings, 'DJFS'):
    djfs_settings = settings.DJFS
else:
    djfs_settings = {'type' : 'osfs',
                     'directory_root' : 'django-pyfs/static/django-pyfs',
                     'url_root' : '/static/django-pyfs'}

if djfs_settings['type'] == 'osfs':
    from fs.osfs import OSFS
elif djfs_settings['type'] == 's3fs':
    from fs.s3fs import S3FS
    from boto.s3.connection import S3Connection
    from boto.s3.key import Key
    key_id = djfs_settings.get('AWS_ACCESS_KEY_ID', None)
    key_secret = djfs_settings.get('AWS_SECRET_ACCESS_KEY', None)
    s3conn = S3Connection(aws_access_key_id=key_id, aws_secret_access_key=key_secret)
else:
    raise AttributeError("Bad filesystem: "+str(djfs_settings['type']))

def get_filesystem(namespace):
    ''' Returns a pyfilesystem for static module storage.

    The file system will have two additional properties:
    1) get_url: A way to get a URL for a static file download
    2) expire: A way to expire files (so they are automatically destroyed)
    '''
    if djfs_settings['type'] == 'osfs':
        return get_osfs( namespace )
    elif djfs_settings['type'] == 's3fs':
        return get_s3fs( namespace )
    else:
        raise AttributeError("Bad filesystem: "+str(djfs_settings['type']))

def expire_objects():
    ''' Remove all obsolete objects from the file systems. Untested. '''
    objects = sorted(FSExpirations.expired(), key=lambda x:x.module)
    fs = None
    module = None
    for o in objects:
        if module != o.module:
            module = o.module
            fs = get_filesystem(module)
        if fs.exists(o.filename):
            fs.remove(o.filename)
        o.delete()

def patch_fs(fs, namespace, url_method):
    ''' Patch a filesystem object to add two methods:
          get_url returns a URL for a resource stored on that filesystem. It takes two parameters:
              filename: Which resource
              timeout: How long that resource is available for
          expire sets a timeout on how long the system should keep the resource. It takes four parameters:
              filename: Which resource
              seconds: How long we will keep it
              days: (optional) More user-friendly if a while
              expires: (optional) boolean; if set to False, we keep the resource forever.
          Without calling this method, we provide no guarantees on how long resources will stick around.
    '''
    def expire(self, filename, seconds, days=0, expires = True):
        ''' Set the lifespan of a file on the filesystem.

        filename: Name of file
        expire: False means the file will never be removed
        seconds and days give time to expiration.
        '''
        FSExpirations.create_expiration(namespace, filename, seconds, days=days, expires = expires)

    fs.expire = types.MethodType(expire, fs)
    fs.get_url = types.MethodType(url_method, fs)
    return fs

def get_osfs(namespace):
    ''' Helper method to get_filesystem for a file system on disk '''
    full_path = os.path.join(djfs_settings['directory_root'], namespace)
    if not os.path.exists(full_path):
        os.makedirs(full_path)
    osfs = OSFS(full_path)
    osfs = patch_fs(osfs, namespace, lambda self, filename, timeout=0:os.path.join(djfs_settings['url_root'], namespace, filename))
    return osfs

def get_s3fs(namespace):
    ''' Helper method to get_filesystem for a file system on S3 '''
    fullpath = namespace
    if 'prefix' in djfs_settings:
        fullpath = os.path.join(djfs_settings['prefix'], fullpath)
    s3fs = S3FS(djfs_settings['bucket'], fullpath)

    def get_s3_url(self, filename, timeout=60):
        global s3conn
        try:
            return s3conn.generate_s3_url(timeout, 'GET', bucket = djfs_settings['bucket'], key = filename)
        except: # If connection has timed out
            s3conn = S3Connection()
            return s3conn.generate_s3_url(timeout, 'GET', bucket = djfs_settings['bucket'], key = filename)

    s3fs = patch_fs(s3fs, namespace, get_s3_url)
    return s3fs

<|endoftext|>