<commit_before># stdlib
from typing import Dict as TypeDict
from typing import KeysView as TypeKeysView
from typing import List as TypeList
from typing import Set as TypeSet

# third party
from autodp.autodp_core import Mechanism
from autodp.transformer_zoo import Composition

# relative
from .entity import Entity
from .scalar import PhiScalar

class AdversarialAccountant:
    def __init__(self, max_budget: float = 10, delta: float = 1e-6) -> None:
        self.entity2ledger: TypeDict[Entity, Mechanism] = {}
        self.max_budget = max_budget
        self.delta = delta

    def append(self, entity2mechanisms: TypeDict[str, TypeList[Mechanism]]) -> None:
        for key, ms in entity2mechanisms.items():
            if key not in self.entity2ledger.keys():
                self.entity2ledger[key] = list()
            for m in ms:
                self.entity2ledger[key].append(m)

    def get_eps_for_entity(self, entity: Entity) -> PhiScalar:
        # compose them with the transformation: compose
        compose = Composition()
        mechanisms = self.entity2ledger[entity]
        composed_mech = compose(mechanisms, [1] * len(mechanisms))

        return composed_mech.get_approxDP(self.delta)

        # # Query for eps given delta
        # return PhiScalar(
        #     value=composed_mech.get_approxDP(self.delta),
        #     min_val=0,
        #     max_val=self.max_budget,
        #     entity=entity,
        # )

    def has_budget(self, entity_name: str) -> bool:
        eps = self.get_eps_for_entity(entity_name)
        if eps is not None:
            return eps < self.max_budget
        # if eps.value is not None:
        #     return eps.value < self.max_budget

    @property
    def entities(self) -> TypeKeysView[str]:
        return self.entity2ledger.keys()

    @property
    def overbudgeted_entities(self) -> TypeSet[str]:
        entities = set()

        for entity_name in self.entities:
            if not self.has_budget(entity_name):
                entities.add(entity_name)

        return entities

    def print_ledger(self, delta: float = 1e-6) -> None:
        for entity, mechanisms in self.entity2ledger.items():
            print(str(entity) + "\t" + str(self.get_eps_for_entity(entity)))
<commit_msg>Update Adversarial_accountant<commit_after># stdlib
from typing import Dict as TypeDict
from typing import KeysView as TypeKeysView
from typing import List as TypeList
from typing import Set as TypeSet

# third party
from autodp.autodp_core import Mechanism
from autodp.transformer_zoo import Composition
from sqlalchemy.engine import Engine

# relative
from .entity import Entity
from .scalar import PhiScalar
from ..node.common.node_manager.ledger_manager import LedgerManager

class AdversarialAccountant:
    def __init__(self, db_engine: Engine, max_budget: float = 10, delta: float = 1e-6) -> None:
        self.entity2ledger = LedgerManager(db_engine)
        self.max_budget = max_budget
        self.delta = delta

    def append(self, entity2mechanisms: TypeDict[str, TypeList[Mechanism]]) -> None:
        for key, ms in entity2mechanisms.items():
            if key not in self.entity2ledger.keys():
                self.entity2ledger[key] = list()
            for m in ms:
                self.entity2ledger[key].append(m)

    def get_eps_for_entity(self, entity_name: str) -> PhiScalar:
        # compose them with the transformation: compose
        compose = Composition()
        mechanisms = self.entity2ledger[entity_name]
        composed_mech = compose(mechanisms, [1] * len(mechanisms))

        return composed_mech.get_approxDP(self.delta)

        # # Query for eps given delta
        # return PhiScalar(
        #     value=composed_mech.get_approxDP(self.delta),
        #     min_val=0,
        #     max_val=self.max_budget,
        #     entity=entity,
        # )

    def has_budget(self, entity_name: str) -> bool:
        eps = self.get_eps_for_entity(entity_name)
        if eps is not None:
            return eps < self.max_budget
        # if eps.value is not None:
        #     return eps.value < self.max_budget

    @property
    def entities(self) -> TypeKeysView[str]:
        return self.entity2ledger.keys()

    @property
    def overbudgeted_entities(self) -> TypeSet[str]:
        entities = set()

        for entity_name in self.entities:
            if not self.has_budget(entity_name):
                entities.add(entity_name)

        return entities

    def print_ledger(self, delta: float = 1e-6) -> None:
        for entity, mechanisms in self.entity2ledger.items():
            print(str(entity) + "\t" + str(self.get_eps_for_entity(entity)))
<|endoftext|>


<commit_before>import pytest

def test_tor_mirror_present(File):
    """
    Ensure the FPF mirror of the Tor apt repo, tor-apt.ops.freedom.press,
    is configured. This repository required manual updating with current
    tor releases, to avoid breakage of untested updates.
    """
    f = File('/etc/apt/sources.list.d/tor_apt_ops_apt_freedom_press.list')

    regex = ('^deb http:\/\/tor-apt\.ops\.freedom\.press trusty main$')
    assert f.contains(regex)

def test_tor_mirror_fingerprint(Command):
    """
    Ensure the FPF tor mirror repo has the correct fingerprint on the
    associated signing pubkey. We don't use the SecureDrop Release Signing
    Key for the Tor mirror: packages are identical to the official Tor repo,
    so signatures match, as well.
    """
    c = Command('apt-key finger')
    tor_gpg_pub_key_info = """/etc/apt/trusted.gpg.d/deb.torproject.org-keyring.gpg
-----------------------------------------------------
pub   2048R/886DDD89 2009-09-04 [expires: 2020-08-29]
      Key fingerprint = A3C4 F0F9 79CA A22C DBA8  F512 EE8C BC9E 886D DD89
uid                  deb.torproject.org archive signing key
sub   2048R/219EC810 2009-09-04 [expires: 2018-08-30]"""

    assert c.rc == 0
    assert tor_gpg_pub_key_info in c.stdout

@pytest.mark.parametrize('filename', [
    '/etc/apt/security.list',
    '/etc/apt/sources.list.d',
])
def test_tor_project_repo_absent(Command):
    """
    Ensure that no apt source list files contain the entry for
    the official Tor apt repo, since we don't control issuing updates
    in that repo. We're mirroring it to avoid breakage caused by
    untested updates (which has broken prod twice to date).
    """
    c = Command("grep -riP 'deb\.torproject\.org' /etc/apt*")
    # Grep returns non-zero when no matches, and we want no matches.
    assert c.rc != 0
    assert c.stdout == ""

def test_tor_project_repo_files_absent(File):
    """
    Ensure that specific apt source list files are absent,
    having been 'hidden' via the securedrop-config package.
    """
    f = File("/etc/apt/sources.list.d/deb_torproject_org_torproject_org.list")
    assert not f.exists
<commit_msg>Repairs parametrization on testinfra tor tests<commit_after>import pytest

def test_tor_mirror_present(File):
    """
    Ensure the FPF mirror of the Tor apt repo, tor-apt.ops.freedom.press,
    is configured. This repository required manual updating with current
    tor releases, to avoid breakage of untested updates.
    """
    f = File('/etc/apt/sources.list.d/tor_apt_ops_apt_freedom_press.list')

    regex = ('^deb http:\/\/tor-apt\.ops\.freedom\.press trusty main$')
    assert f.contains(regex)

def test_tor_mirror_fingerprint(Command):
    """
    Ensure the FPF tor mirror repo has the correct fingerprint on the
    associated signing pubkey. We don't use the SecureDrop Release Signing
    Key for the Tor mirror: packages are identical to the official Tor repo,
    so signatures match, as well.
    """
    c = Command('apt-key finger')
    tor_gpg_pub_key_info = """/etc/apt/trusted.gpg.d/deb.torproject.org-keyring.gpg
-----------------------------------------------------
pub   2048R/886DDD89 2009-09-04 [expires: 2020-08-29]
      Key fingerprint = A3C4 F0F9 79CA A22C DBA8  F512 EE8C BC9E 886D DD89
uid                  deb.torproject.org archive signing key
sub   2048R/219EC810 2009-09-04 [expires: 2018-08-30]"""

    assert c.rc == 0
    assert tor_gpg_pub_key_info in c.stdout

@pytest.mark.parametrize('filename', [
    '/etc/apt/security.list',
    '/etc/apt/sources.list.d',
])
def test_tor_project_repo_absent(Command, filename):
    """
    Ensure that no apt source list files contain the entry for
    the official Tor apt repo, since we don't control issuing updates
    in that repo. We're mirroring it to avoid breakage caused by
    untested updates (which has broken prod twice to date).
    """
    c = Command("grep -riP 'deb\.torproject\.org' {}".format(filename))
    # Grep returns non-zero when no matches, and we want no matches.
    assert c.rc != 0
    assert c.stdout == ""

def test_tor_project_repo_files_absent(File):
    """
    Ensure that specific apt source list files are absent,
    having been 'hidden' via the securedrop-config package.
    """
    f = File("/etc/apt/sources.list.d/deb_torproject_org_torproject_org.list")
    assert not f.exists
<|endoftext|>


<commit_before># Copyright (c) Ecometrica. All rights reserved.
# Distributed under the BSD license. See LICENSE for details.
from collections import namedtuple
from decorator import decorator
from functools import wraps

from django.db.models import get_apps
from piston.utils import rc

from wapiti.conf import ID_RE

_RegisteredType = namedtuple('RegisteredType', ('model', 'api'))

_registered_types = {}

def register(name, modelapi):
    """Register a model with the API"""
    global _registered_types
    if modelapi.__name__ in _registered_types:
        return
    if not modelapi.objects:
        modelapi.objects = modelapi.model.objects

    _registered_types[name] = _RegisteredType(model=modelapi.model,
                                              api=modelapi)

def api_method(f):
    """Decorator to declare a method api-accessible"""
    f.api = True
    def wrapper(*args, **kwargs):
        return f(*args, **kwargs)
    return decorator(wrapper, f)

def _is_id(id):
    return ID_RE.match(id)

def _register_models():
    """Find app api submodules and register models"""
    for a in get_apps():
        try:
            _temp = __import__('.'.join(a.__name__.split('.')[:-1] + ['api']),
                               globals(), locals())
        except ImportError:
            pass

<commit_msg>reorganize wrapper<commit_after># Copyright (c) Ecometrica. All rights reserved.
# Distributed under the BSD license. See LICENSE for details.
from collections import namedtuple
from decorator import decorator
from functools import wraps

from django.db.models import get_apps
from piston.utils import rc

from wapiti.conf import ID_RE

_RegisteredType = namedtuple('RegisteredType', ('model', 'api'))

_registered_types = {}

def register(name, modelapi):
    """Register a model with the API"""
    global _registered_types
    if modelapi.__name__ in _registered_types:
        return
    if not modelapi.objects:
        modelapi.objects = modelapi.model.objects

    _registered_types[name] = _RegisteredType(model=modelapi.model,
                                              api=modelapi)

def _api_method(f, *args, **kwargs):
    return f(*args, **kwargs)

def api_method(f):
    """Decorator to declare a method api-accessible"""
    f.api = True
    return decorator(_api_method, f)

def _is_id(id):
    return ID_RE.match(id)

def _register_models():
    """Find app api submodules and register models"""
    for a in get_apps():
        try:
            _temp = __import__('.'.join(a.__name__.split('.')[:-1] + ['api']),
                               globals(), locals())
        except ImportError:
            pass

<|endoftext|>


<commit_before>import threading
import time
import queue
from functools import wraps, partial
import sys

from . import eventloop

def behavior(f):
    @wraps(f)
    def wrapper(*args):
        def _f(context, msg):
            f(*(args + (context, msg)))
        return _f
    return wrapper

@behavior
def bar(context, msg):
    print("i'm bar", msg)

@behavior
def foo(n, context, msg):
    print(n, context, msg)
    context.become(bar)

@behavior
def gii(context, n):
    print('gii got', n)
    if n > 0:
        context.self << n-1
    else:
        print('finally')
        a = context.sponsor(foo, 3)
        a << 5

q = queue.Queue()

def beh_loop(queue, block=True):
    try:
        while True:
            actor, msg = queue.get(block=block)
            try:
                actor._behavior(actor, msg)
            except Exception as exc:
                print(eventloop._format_exception(sys.exc_info()))
    except StopIteration:
        return

def run_loop():
    try:
        beh_loop(q, block=False)
    except queue.Empty:
        return

# loop = threading.Thread(
#     target=beh_loop,
#     args=(queue,),
#     name='beh_loop')
#loop.start()

def sponsor(behavior, *args):
    a = Actor(sponsor, behavior, *args)
    return Capability(a.send)

class Actor(object):

    def __init__(self, sponsor, behavior, *args):
        self.sponsor = sponsor
        self.self = Capability(self.send)
        self.become(behavior, *args)

    def become(self, behavior, *args):
        self._behavior = behavior(*args)

    def send(self, msg):
        q.put((self, msg))

class Capability(object):

    def __init__(self, f):
        self.f = f

    def __lshift__(self, msg):
        self.f(msg)

def test_receive_message():

    result = None

    @behavior
    def beh(context, msg):
        nonlocal result
        result = msg

    a = sponsor(beh)
    a << 5
    run_loop()
    assert result == 5

def test_create_with_args():

    result = None

    @behavior
    def beh(arg, context, msg):
        nonlocal result
        result = arg

    a = sponsor(beh, True)
    a << 0
    run_loop()
    assert result is True

def test_one_shot():

    sink_beh_done = False
    destination_beh_done = False
    message = None

    @behavior
    def one_shot_beh(destination, context, msg):
        destination << msg
        context.become(sink_beh)

    @behavior
    def sink_beh(context, msg):
        assert msg == 'second'
        nonlocal sink_beh_done
        sink_beh_done = True

    @behavior
    def destination_beh(context, msg):
        nonlocal message, destination_beh_done
        message = msg
        destination_beh_done = True

    destination = sponsor(destination_beh)
    one_shot = sponsor(one_shot_beh, destination)

    one_shot << 'first'
    one_shot << 'second'

    run_loop()
    assert message == 'first'
    assert sink_beh_done and destination_beh_done

def test_serial():

    first_msg = second_msg = third_msg = None
    first_behavior = second_behavior = third_behavior = None
    first = second = third = None

    @behavior
    def first_beh(context, msg):
        context.become(second_beh)
        nonlocal first_msg, first_behavior, first
        first_msg = msg
        first_behavior = record()
        first = True

    @behavior
    def second_beh(context, msg):
        context.become(third_beh)
        nonlocal second_msg, second_behavior, second
        second_msg = msg
        second_behavior = record()
        second = True

    @behavior
    def third_beh(context, msg):
        nonlocal third_msg, third_behavior
        third_msg = msg
        third_behavior = record()

    def record():
        return bool(first), bool(second), bool(third)

    serial = sponsor(first_beh)
    serial << 'foo'
    serial << 'foo'
    serial << 'foo'

    run_loop()

    assert first_msg == 'foo' and second_msg == 'foo' and third_msg == 'foo'
    assert first_behavior == (False, False, False)
    assert second_behavior == (True, False, False)
    assert third_behavior == (True, True, False)

@behavior
def ringlink_beh(next, context, n):
    next << n

@behavior
def ringlast_beh(first, context, n):
    loop_completion_times.append(time.time())
    if n > 1:
        first << n-1
    else:
        context.become(sink_beh)
        report()

@behavior
def sink_beh(context, msg):
    pass

@behavior
def ringbuilder_beh(m, context, msg):
    if m > 0:
        next = context.sponsor(ringbuilder_beh, m-1)
        next << msg
        context.become(ringlink_beh, next)
    else:
        global construction_end_time
        construction_end_time = time.time()
        msg['first'] << msg['n']
        context.become(ringlast_beh, msg['first'])

def erlang_challenge(m, n):
    print('Starting {} actor ring'. format(m))
    global construction_start_time
    construction_start_time = time.time()
    ring = sponsor(ringbuilder_beh, m)
    ring << {'first': ring, 'n': n}

construction_start_time = 0
construction_end_time = 0
loop_completion_times = []

def report():
    print('Construction time: {} seconds'.format(
        construction_end_time - construction_start_time))
    print('Loop times:')
    loop_completion_times.insert(0, construction_end_time)
    intervals = [t1-t0
                 for (t0, t1) in zip(loop_completion_times,
                                     loop_completion_times[1:])]
    for t in intervals:
        print('  {} seconds'.format(t))
    print('Average: {} seconds'.format(sum(intervals)/len(intervals)))
<commit_msg>Rename objects for more sensible explanation<commit_after>import threading
import time
import queue
from functools import wraps, partial
import sys

from . import eventloop

def behavior(f):
    @wraps(f)
    def wrapper(*args):
        def _f(context, msg):
            f(*(args + (context, msg)))
        return _f
    return wrapper

@behavior
def bar(context, msg):
    print("i'm bar", msg)

@behavior
def foo(n, context, msg):
    print(n, context, msg)
    context.become(bar)

@behavior
def gii(context, n):
    print('gii got', n)
    if n > 0:
        context.self << n-1
    else:
        print('finally')
        a = context.sponsor(foo, 3)
        a << 5

q = queue.Queue()

def beh_loop(queue, block=True):
    try:
        while True:
            context, msg = queue.get(block=block)
            try:
                context._behavior(context, msg)
            except Exception as exc:
                print(eventloop._format_exception(sys.exc_info()))
    except StopIteration:
        return

def run_loop():
    try:
        beh_loop(q, block=False)
    except queue.Empty:
        return

# loop = threading.Thread(
#     target=beh_loop,
#     args=(queue,),
#     name='beh_loop')
#loop.start()

def sponsor(behavior, *args):
    a = Context(sponsor, behavior, *args)
    return Actor(a.send)

class Context(object):

    def __init__(self, sponsor, behavior, *args):
        self.sponsor = sponsor
        self.self = Actor(self.send)
        self.become(behavior, *args)

    def become(self, behavior, *args):
        self._behavior = behavior(*args)

    def send(self, msg):
        q.put((self, msg))

class Actor(object):

    def __init__(self, f):
        self.f = f

    def __lshift__(self, msg):
        self.f(msg)

def test_receive_message():

    result = None

    @behavior
    def beh(context, msg):
        nonlocal result
        result = msg

    a = sponsor(beh)
    a << 5
    run_loop()
    assert result == 5

def test_create_with_args():

    result = None

    @behavior
    def beh(arg, context, msg):
        nonlocal result
        result = arg

    a = sponsor(beh, True)
    a << 0
    run_loop()
    assert result is True

def test_one_shot():

    sink_beh_done = False
    destination_beh_done = False
    message = None

    @behavior
    def one_shot_beh(destination, context, msg):
        destination << msg
        context.become(sink_beh)

    @behavior
    def sink_beh(context, msg):
        assert msg == 'second'
        nonlocal sink_beh_done
        sink_beh_done = True

    @behavior
    def destination_beh(context, msg):
        nonlocal message, destination_beh_done
        message = msg
        destination_beh_done = True

    destination = sponsor(destination_beh)
    one_shot = sponsor(one_shot_beh, destination)

    one_shot << 'first'
    one_shot << 'second'

    run_loop()
    assert message == 'first'
    assert sink_beh_done and destination_beh_done

def test_serial():

    first_msg = second_msg = third_msg = None
    first_behavior = second_behavior = third_behavior = None
    first = second = third = None

    @behavior
    def first_beh(context, msg):
        context.become(second_beh)
        nonlocal first_msg, first_behavior, first
        first_msg = msg
        first_behavior = record()
        first = True

    @behavior
    def second_beh(context, msg):
        context.become(third_beh)
        nonlocal second_msg, second_behavior, second
        second_msg = msg
        second_behavior = record()
        second = True

    @behavior
    def third_beh(context, msg):
        nonlocal third_msg, third_behavior
        third_msg = msg
        third_behavior = record()

    def record():
        return bool(first), bool(second), bool(third)

    serial = sponsor(first_beh)
    serial << 'foo'
    serial << 'foo'
    serial << 'foo'

    run_loop()

    assert first_msg == 'foo' and second_msg == 'foo' and third_msg == 'foo'
    assert first_behavior == (False, False, False)
    assert second_behavior == (True, False, False)
    assert third_behavior == (True, True, False)

@behavior
def ringlink_beh(next, context, n):
    next << n

@behavior
def ringlast_beh(first, context, n):
    loop_completion_times.append(time.time())
    if n > 1:
        first << n-1
    else:
        context.become(sink_beh)
        report()

@behavior
def sink_beh(context, msg):
    pass

@behavior
def ringbuilder_beh(m, context, msg):
    if m > 0:
        next = context.sponsor(ringbuilder_beh, m-1)
        next << msg
        context.become(ringlink_beh, next)
    else:
        global construction_end_time
        construction_end_time = time.time()
        msg['first'] << msg['n']
        context.become(ringlast_beh, msg['first'])

def erlang_challenge(m, n):
    print('Starting {} actor ring'. format(m))
    global construction_start_time
    construction_start_time = time.time()
    ring = sponsor(ringbuilder_beh, m)
    ring << {'first': ring, 'n': n}

construction_start_time = 0
construction_end_time = 0
loop_completion_times = []

def report():
    print('Construction time: {} seconds'.format(
        construction_end_time - construction_start_time))
    print('Loop times:')
    loop_completion_times.insert(0, construction_end_time)
    intervals = [t1-t0
                 for (t0, t1) in zip(loop_completion_times,
                                     loop_completion_times[1:])]
    for t in intervals:
        print('  {} seconds'.format(t))
    print('Average: {} seconds'.format(sum(intervals)/len(intervals)))
<|endoftext|>


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

import sqlite3 as sqlite
import pandas as pd
from pandas.io.sql import to_sql, read_sql
import re
import os
import inspect

class SQLDF(object):
    def __init__(self, env, inmemory=True, udfs={}, udafs={}):
        super(SQLDF, self).__init__()
        self.env = env
        self.inmemory = inmemory
        self.udfs = udfs
        self.udafs = udafs
        if self.inmemory:
            self._dbname = ":memory:"
        else:
            self._dbname = ".pysqldf.db"
        self.conn = sqlite.connect(self._dbname, detect_types=sqlite.PARSE_DECLTYPES|sqlite.PARSE_COLNAMES)
        self._set_udf(udfs)
        self._set_udaf(udafs)

    def __del__(self):
        self.conn.close()
        if not self.inmemory:
            os.remove(self._dbname)

    def execute(self, query):
        tables = self._extract_table_names(query)
        for table in tables:
            if table not in self.env:
                raise Exception("%s not found" % table)
            df = self.env[table]
            df = self._ensure_data_frame(df, table)
            self._write_table(table, df)

        try:
            result = read_sql(query, self.conn, index_col=None)
            if 'index' in result:
                del result['index']
        except Exception:
            result = None
        finally:
            self._del_table(tables)
        return result

    def _extract_table_names(self, query):
        "extracts table names from a sql query"
        # a good old fashioned regex. turns out this worked better than actually parsing the code
        rgx = r'(?:from|join)\s+([a-z_][a-z0-9_]*)(?:\s|;|$|\))'
        tables = re.findall(rgx, query, re.IGNORECASE)
        return list(set(tables))

    def _ensure_data_frame(self, obj, name):
        """
        obj a python object to be converted to a DataFrame

        take an object and make sure that it's a pandas data frame
        """
        try:
            df = pd.DataFrame(obj)
            columns = [col if isinstance(col, str) else "c%d" % i
                       for i, col in enumerate(df.columns)]
            df.columns = columns
        except Exception:
            raise Exception("%s is not a convertable data to Dataframe" % name)

        return df

    def _write_table(self, tablename, df):
        "writes a dataframe to the sqlite database"

        for col in df.columns:
            if re.search("[()]", col):
                msg = "please follow SQLite column naming conventions: "
                msg += "http://www.sqlite.org/lang_keywords.html"
                raise Exception(msg)

        dtype = dict((k, str(v)) for k, v in dict(df.dtypes).items())
        to_sql(df, name=tablename, con=self.conn, flavor='sqlite', dtype=dtype)

    def _del_table(self, tablenames):
        for tablename in tablenames:
            self.conn.execute("drop table " + tablename)
        self.conn.commit()

    def _set_udf(self, udfs):
        for name, func in udfs.items():
            num_params = len(inspect.getargspec(func).args)
            self.conn.create_function(name, num_params, func)

    def _set_udaf(self, udafs):
        for name, agg_class in udafs.items():
            if inspect.isfunction(agg_class):
                agg_class = self._create_agg_class(agg_class)
            num_params = len(inspect.getargspec(agg_class.step).args) - 1 # subtract self
            self.conn.create_aggregate(name, num_params, agg_class)

    def _create_agg_class(self, agg_function):
        return type(agg_function.__name__, (), {
            "__init__": lambda self: setattr(self, "l", []),
            "step": lambda self, x: self.l.append(x),
            "finalize": lambda self: agg_function(self.l)
        })

if __name__ == '__main__':
    from pandas import DataFrame

    data = [
        ["hoge", 0,  1,  2,  3,  4],
        ["fuga", 5,  6,  7,  8,  9],
        ["hoge", 10, 11, 12, 13, 14],
        ["fuga", 15, 16, 17, 18, 19],
        ["hoge", 20, 21, 22, 23, 24]
    ]
    df = DataFrame(data, columns=["label", "a", "b", "c", "d", "e"])
    class mysum(object):
        def __init__(self):
            self.count = 0
        def step(self, val):
            self.count += val
        def finalize(self):
            return self.count

    sqldf = SQLDF(locals(), udfs={"half": lambda x: x / 2}, udafs={"mysum": mysum})
    print(sqldf.execute("select * from data;"))
    print(sqldf.execute("select sum(c1) from data;"))
    print(sqldf.execute("select half(c1) from data;"))
    print(sqldf.execute("select mysum(c2) from data;"))
    print(sqldf.execute("select MySum(c2) from data;"))
    print(sqldf.execute("select * from df;"))
    print(sqldf.execute("select sum(a) from df;"))
    print(sqldf.execute("select half(a) from df;"))
    print(sqldf.execute("select mysum(b) from df;"))
    print(sqldf.execute("select MySum(b) from df;"))
<commit_msg>use anonymous class<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sqlite3 as sqlite
import pandas as pd
from pandas.io.sql import to_sql, read_sql
import re
import os
import inspect

class SQLDF(object):
    def __init__(self, env, inmemory=True, udfs={}, udafs={}):
        super(SQLDF, self).__init__()
        self.env = env
        self.inmemory = inmemory
        self.udfs = udfs
        self.udafs = udafs
        if self.inmemory:
            self._dbname = ":memory:"
        else:
            self._dbname = ".pysqldf.db"
        self.conn = sqlite.connect(self._dbname, detect_types=sqlite.PARSE_DECLTYPES|sqlite.PARSE_COLNAMES)
        self._set_udf(udfs)
        self._set_udaf(udafs)

    def __del__(self):
        self.conn.close()
        if not self.inmemory:
            os.remove(self._dbname)

    def execute(self, query):
        tables = self._extract_table_names(query)
        for table in tables:
            if table not in self.env:
                raise Exception("%s not found" % table)
            df = self.env[table]
            df = self._ensure_data_frame(df, table)
            self._write_table(table, df)

        try:
            result = read_sql(query, self.conn, index_col=None)
            if 'index' in result:
                del result['index']
        except Exception:
            result = None
        finally:
            self._del_table(tables)
        return result

    def _extract_table_names(self, query):
        "extracts table names from a sql query"
        # a good old fashioned regex. turns out this worked better than actually parsing the code
        rgx = r'(?:from|join)\s+([a-z_][a-z0-9_]*)(?:\s|;|$|\))'
        tables = re.findall(rgx, query, re.IGNORECASE)
        return list(set(tables))

    def _ensure_data_frame(self, obj, name):
        """
        obj a python object to be converted to a DataFrame

        take an object and make sure that it's a pandas data frame
        """
        try:
            df = pd.DataFrame(obj)
            columns = [col if isinstance(col, str) else "c%d" % i
                       for i, col in enumerate(df.columns)]
            df.columns = columns
        except Exception:
            raise Exception("%s is not a convertable data to Dataframe" % name)

        return df

    def _write_table(self, tablename, df):
        "writes a dataframe to the sqlite database"

        for col in df.columns:
            if re.search("[()]", col):
                msg = "please follow SQLite column naming conventions: "
                msg += "http://www.sqlite.org/lang_keywords.html"
                raise Exception(msg)

        dtype = dict((k, str(v)) for k, v in dict(df.dtypes).items())
        to_sql(df, name=tablename, con=self.conn, flavor='sqlite', dtype=dtype)

    def _del_table(self, tablenames):
        for tablename in tablenames:
            self.conn.execute("drop table " + tablename)
        self.conn.commit()

    def _set_udf(self, udfs):
        for name, func in udfs.items():
            num_params = len(inspect.getargspec(func).args)
            self.conn.create_function(name, num_params, func)

    def _set_udaf(self, udafs):
        for name, agg_class in udafs.items():
            if inspect.isfunction(agg_class):
                agg_class = self._create_agg_class(agg_class)
            num_params = len(inspect.getargspec(agg_class.step).args) - 1 # subtract self
            self.conn.create_aggregate(name, num_params, agg_class)

    def _create_agg_class(self, agg_function):
        # return anonymous class
        return type("", (), {
            "__init__": lambda self: setattr(self, "l", []),
            "step": lambda self, x: self.l.append(x),
            "finalize": lambda self: agg_function(self.l)
        })

if __name__ == '__main__':
    from pandas import DataFrame

    data = [
        ["hoge", 0,  1,  2,  3,  4],
        ["fuga", 5,  6,  7,  8,  9],
        ["hoge", 10, 11, 12, 13, 14],
        ["fuga", 15, 16, 17, 18, 19],
        ["hoge", 20, 21, 22, 23, 24]
    ]
    df = DataFrame(data, columns=["label", "a", "b", "c", "d", "e"])
    class mysum(object):
        def __init__(self):
            self.count = 0
        def step(self, val):
            self.count += val
        def finalize(self):
            return self.count

    sqldf = SQLDF(locals(), udfs={"half": lambda x: x / 2}, udafs={"mysum": mysum})
    print(sqldf.execute("select * from data;"))
    print(sqldf.execute("select sum(c1) from data;"))
    print(sqldf.execute("select half(c1) from data;"))
    print(sqldf.execute("select mysum(c2) from data;"))
    print(sqldf.execute("select MySum(c2) from data;"))
    print(sqldf.execute("select * from df;"))
    print(sqldf.execute("select sum(a) from df;"))
    print(sqldf.execute("select half(a) from df;"))
    print(sqldf.execute("select mysum(b) from df;"))
    print(sqldf.execute("select MySum(b) from df;"))
<|endoftext|>


<commit_before># 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.

import functools
import hmac
import urllib.parse

from pyramid.config.views import DefaultViewMapper
from pyramid.interfaces import IViewMapperFactory
from pyramid.httpexceptions import HTTPForbidden, HTTPMethodNotAllowed

from warehouse.utils.http import add_vary

REASON_NO_ORIGIN = "Origin checking failed - no Origin or Referer."
REASON_BAD_ORIGIN = "Origin checking failed - {} does not match {}."
REASON_BAD_TOKEN = "CSRF token missing or incorrect."

class InvalidCSRF(HTTPForbidden):
    pass

def csrf_exempt(view):
    @functools.wraps(view)
    def wrapped(context, request):
        request._process_csrf = False
        return view(context, request)
    return wrapped

def csrf_protect(view_or_scope):
    scope = None
    if isinstance(view_or_scope, str):
        scope = view_or_scope

    def inner(view):
        @functools.wraps(view)
        def wrapped(context, request):
            request._process_csrf = True
            request._csrf_scope = scope
            return view(context, request)
        return wrapped

    if scope is None:
        return inner(view_or_scope)
    else:
        return inner

def csrf_mapper_factory(mapper):
    class CSRFMapper(mapper):

        def __call__(self, view):
            view = super().__call__(view)

            # Check if the view has CSRF exempted, and if it is then we just
            # want to return the view without wrapping it.
            if not getattr(view, "_process_csrf", True):
                return view

            @functools.wraps(view)
            def wrapped(context, request):
                # Assign our view to an innerview function so that we can
                # modify it inside of the wrapped function.
                innerview = view

                # Check if we're processing CSRF for this request at all or
                # if it has been exempted from CSRF.
                if not getattr(request, "_process_csrf", True):
                    return innerview(context, request)

                # If we're processing CSRF for this request, then we want to
                # set a Vary: Cookie header on every response to ensure that
                # we don't cache the result of a CSRF check or a form with a
                # CSRF token in it.
                if getattr(request, "_process_csrf", None):
                    innerview = add_vary("Cookie")(innerview)

                # Assume that anything not defined as 'safe' by RFC2616 needs
                # protection
                if request.method not in {"GET", "HEAD", "OPTIONS", "TRACE"}:
                    # Determine if this request has set itself so that it
                    # should be protected against CSRF. If it has not and it's
                    # gotten one of these methods, then we want to raise an
                    # error stating that this resource does not support this
                    # method.
                    if not getattr(request, "_process_csrf", None):
                        raise HTTPMethodNotAllowed

                    if request.scheme == "https":
                        # Determine the origin of this request
                        origin = request.headers.get("Origin")
                        if origin is None:
                            origin = request.headers.get("Referer")

                        # Fail if we were not able to locate an origin at all
                        if not origin:
                            raise InvalidCSRF(REASON_NO_ORIGIN)

                        # Parse the origin and host for comparison
                        originp = urllib.parse.urlparse(origin)
                        hostp = urllib.parse.urlparse(request.host_url)

                        # Actually check our Origin against our Current
                        # Host URL.
                        if ((originp.scheme, originp.hostname, originp.port)
                                != (hostp.scheme, hostp.hostname, hostp.port)):
                            reason_origin = origin
                            if origin != "null":
                                reason_origin = urllib.parse.urlunparse(
                                    originp[:2] + ("", "", "", ""),
                                )

                            reason = REASON_BAD_ORIGIN.format(
                                reason_origin, request.host_url,
                            )

                            raise InvalidCSRF(reason)

                    session = getattr(request, "_session", request.session)

                    # Get the provided CSRF token from the request.
                    request_token = request.POST.get("csrf_token", "")
                    if not request_token:
                        request_token = request.headers.get("CSRFToken", "")

                    # Get our CSRF token from the session, scoped or not
                    # depending on if our @csrf_protect header was registered
                    # with a scope or not.
                    scope = request._csrf_scope
                    if scope is None:
                        csrf_token = session.get_csrf_token()
                    else:
                        csrf_token = session.get_scoped_csrf_token(scope)

                    if not hmac.compare_digest(csrf_token, request_token):
                        raise InvalidCSRF(REASON_BAD_TOKEN)

                return innerview(context, request)

            return wrapped
    return CSRFMapper

def includeme(config):
    # We need to commit what's happened so far so that we can get the current
    # default ViewMapper
    config.commit()

    # Get the current default ViewMapper, and create a subclass of it that
    # will wrap our view with CSRF checking.
    mapper = config.registry.queryUtility(IViewMapperFactory)
    if mapper is None:
        mapper = DefaultViewMapper
    config.set_view_mapper(csrf_mapper_factory(mapper))
<commit_msg>Remove a leftover conditional<commit_after># 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.

import functools
import hmac
import urllib.parse

from pyramid.config.views import DefaultViewMapper
from pyramid.interfaces import IViewMapperFactory
from pyramid.httpexceptions import HTTPForbidden, HTTPMethodNotAllowed

from warehouse.utils.http import add_vary

REASON_NO_ORIGIN = "Origin checking failed - no Origin or Referer."
REASON_BAD_ORIGIN = "Origin checking failed - {} does not match {}."
REASON_BAD_TOKEN = "CSRF token missing or incorrect."

class InvalidCSRF(HTTPForbidden):
    pass

def csrf_exempt(view):
    @functools.wraps(view)
    def wrapped(context, request):
        request._process_csrf = False
        return view(context, request)
    return wrapped

def csrf_protect(view_or_scope):
    scope = None
    if isinstance(view_or_scope, str):
        scope = view_or_scope

    def inner(view):
        @functools.wraps(view)
        def wrapped(context, request):
            request._process_csrf = True
            request._csrf_scope = scope
            return view(context, request)
        return wrapped

    if scope is None:
        return inner(view_or_scope)
    else:
        return inner

def csrf_mapper_factory(mapper):
    class CSRFMapper(mapper):

        def __call__(self, view):
            view = super().__call__(view)

            @functools.wraps(view)
            def wrapped(context, request):
                # Assign our view to an innerview function so that we can
                # modify it inside of the wrapped function.
                innerview = view

                # Check if we're processing CSRF for this request at all or
                # if it has been exempted from CSRF.
                if not getattr(request, "_process_csrf", True):
                    return innerview(context, request)

                # If we're processing CSRF for this request, then we want to
                # set a Vary: Cookie header on every response to ensure that
                # we don't cache the result of a CSRF check or a form with a
                # CSRF token in it.
                if getattr(request, "_process_csrf", None):
                    innerview = add_vary("Cookie")(innerview)

                # Assume that anything not defined as 'safe' by RFC2616 needs
                # protection
                if request.method not in {"GET", "HEAD", "OPTIONS", "TRACE"}:
                    # Determine if this request has set itself so that it
                    # should be protected against CSRF. If it has not and it's
                    # gotten one of these methods, then we want to raise an
                    # error stating that this resource does not support this
                    # method.
                    if not getattr(request, "_process_csrf", None):
                        raise HTTPMethodNotAllowed

                    if request.scheme == "https":
                        # Determine the origin of this request
                        origin = request.headers.get("Origin")
                        if origin is None:
                            origin = request.headers.get("Referer")

                        # Fail if we were not able to locate an origin at all
                        if not origin:
                            raise InvalidCSRF(REASON_NO_ORIGIN)

                        # Parse the origin and host for comparison
                        originp = urllib.parse.urlparse(origin)
                        hostp = urllib.parse.urlparse(request.host_url)

                        # Actually check our Origin against our Current
                        # Host URL.
                        if ((originp.scheme, originp.hostname, originp.port)
                                != (hostp.scheme, hostp.hostname, hostp.port)):
                            reason_origin = origin
                            if origin != "null":
                                reason_origin = urllib.parse.urlunparse(
                                    originp[:2] + ("", "", "", ""),
                                )

                            reason = REASON_BAD_ORIGIN.format(
                                reason_origin, request.host_url,
                            )

                            raise InvalidCSRF(reason)

                    session = getattr(request, "_session", request.session)

                    # Get the provided CSRF token from the request.
                    request_token = request.POST.get("csrf_token", "")
                    if not request_token:
                        request_token = request.headers.get("CSRFToken", "")

                    # Get our CSRF token from the session, scoped or not
                    # depending on if our @csrf_protect header was registered
                    # with a scope or not.
                    scope = request._csrf_scope
                    if scope is None:
                        csrf_token = session.get_csrf_token()
                    else:
                        csrf_token = session.get_scoped_csrf_token(scope)

                    if not hmac.compare_digest(csrf_token, request_token):
                        raise InvalidCSRF(REASON_BAD_TOKEN)

                return innerview(context, request)

            return wrapped
    return CSRFMapper

def includeme(config):
    # We need to commit what's happened so far so that we can get the current
    # default ViewMapper
    config.commit()

    # Get the current default ViewMapper, and create a subclass of it that
    # will wrap our view with CSRF checking.
    mapper = config.registry.queryUtility(IViewMapperFactory)
    if mapper is None:
        mapper = DefaultViewMapper
    config.set_view_mapper(csrf_mapper_factory(mapper))
<|endoftext|>


<commit_before># Copyright (c) 2012 Israel Herraiz <isra@herraiz.org>

# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import wx
from wx import xrc
from grid import Grid
from image_importer import ImageImporter
from math import sqrt

class MyApp(wx.App):

    def __init__ (self, xrcfn, colorsfn):

        self._xrcfn = xrcfn
        self._colorsfn = colorsfn
        self._scroll_rate = 10
        self._erase_tool = False
        self._grid = Grid()

        self._operations = []
        self._current_operation = None
        self._max_undo = 100

        self._timer = None
        self._current_mouse_pos = (-1, -1)

        wx.App.__init__ (self)

    def OnInit (self):

        # Colors must be imported before creating the frame
        self._import_colors ()
        self.current_color = None

        # Create main frame
        self._res = xrc.XmlResource (self._xrcfn)
        self._init_frame()

        return True

    def _import_colors (self):

        f = open(self._colorsfn, 'r')
        ls = f.readlines()
        f.close()

        self._colors = {}
        for l in ls:
            dmc, name, code = l.split(',')
            self._colors[dmc] = (code, name)

    def _find_dmc_color (self, color):
        for dmc in self._colors.keys():
            code, name = self._colors[dmc]

            red = int(code[1:3], 16)
            green = int(code[3:5], 16)
            blue = int(code[5:7], 16)

            if red == color.Red() and green == color.Green() and blue == color.Blue():
                return name

        return 'None'

    def _find_closest_dmc_color (self, color):

        distance = 1000

        red = color.Red()
        green = color.Green()
        blue = color.Blue()

        bestred, bestgreen, bestblue = (None, None, None)

        for dmc in self._colors.keys():
            code, name = self._colors[dmc]

            dred = int(code[1:3], 16)
            dgreen = int(code[3:5], 16)
            dblue = int(code[5:7], 16)

            r = (dred + red) / 2
            dr = dred - red
            dg = dgreen - green
            db = dblue - blue

            ndistance = (2+r/256)*dr**2 + 4*dg**2 + (2+(255-r)/256)*db**2
            ndistance = sqrt(ndistance)

            if ndistance < distance:
                distance = ndistance
                bestred, bestgreen, bestblue = (dred, dgreen, dblue)

        return wx.Colour (bestred, bestgreen, bestblue)

    def OnPaint (self, event):

        dc = wx.PaintDC (event.GetEventObject())
        dc.Clear()
        self._panel.DoPrepareDC(dc)
        self._grid.draw_grid (dc)

        event.Skip()

    def _init_frame (self):
        self._frame = self._res.LoadFrame (None,'MyMainFrame')
        self._panel = xrc.XRCCTRL (self._frame, 'MainPanel')
        self._panel.SetScrollRate (self._scroll_rate, self._scroll_rate)
        self._panel.SetVirtualSize (self._grid.get_size ())

        self._toolbar = self._frame.GetToolBar ()
        self._toolbar.ToggleTool (xrc.XRCID('editgrid'), not self._erase_tool)
        self._toolbar.ToggleTool (xrc.XRCID('erase'), self._erase_tool)

        color_choice_id = 54 # Random int
        color_list = []
        for k in self._colors.keys():
            dmc = k
            code, name = self._colors[k]

            color_list.append('%s (%s)' % (name, dmc))

        self._color_choice = wx.Choice (self._toolbar, color_choice_id, (-1,-1), (-1,-1), color_list )
        self._toolbar.AddControl (self._color_choice)
        self._change_color(None)

        self._menubar = self._frame.GetMenuBar()
        self._statusbar = self._frame.GetStatusBar()

        self._panel.Bind(wx.EVT_PAINT, self.OnPaint)
        self._panel.Bind(wx.EVT_MOUSE_EVENTS, self._print_cell)
        self._toolbar.Bind(wx.EVT_TOOL, self._undo,     id = xrc.XRCID('undo'))
        self._toolbar.Bind(wx.EVT_TOOL, self._redo,     id = xrc.XRCID('redo'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_zoom, id = xrc.XRCID('zoomout'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_zoom, id = xrc.XRCID('zoomin'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_edit, id = xrc.XRCID('editgrid'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_edit, id = xrc.XRCID('erase'))
        self._toolbar.Bind,wx.EVT_CHOICE(self, color_choice_id, self._change_color)

        self._frame.Bind(wx.EVT_MENU, self._import_image, id = xrc.XRCID('importimage'))

        self._timer = wx.Timer()
        self._timer.Bind(wx.EVT_TIMER, self._show_tooltip)

        self._frame.SetSize ((800,600))
        self._panel.FitInside()
        self._frame.SetTitle ("Stitchy Studio")

        self._timer.Start(3000,True)

        self.SetTopWindow (self._frame)
        self._frame.Show()

    def _import_image (self, event):

        path = wx.FileSelector ('Choose an image',
                               wildcard = "BMP|*.bmp|GIF|*.gif|JPEG|*.jp*g|PNG|*.png|PCX|*.pcx|TIFF|*.tiff|Other|*",
                               flags = wx.FD_FILE_MUST_EXIST,
                               parent = self._frame)

        importer = ImageImporter ()
        importer.load_image (path)
        importer.scale_image()

        height, width = importer.get_size ()

        dc = wx.ClientDC (self._panel)

        self._panel.DoPrepareDC (dc)

        for x in range (0, width):
            for y in range (0, height):
                color = importer.get_color (x, y)
                bestcolor = self._find_closest_dmc_color (color)
                self._grid.add_cell (x, y, dc, bestcolor, False)

        event.Skip()

    def _change_color (self, event):

        selection = self._color_choice.GetStringSelection()

        dmc = selection.split("(")[1].split(")")[0]

        color, _name = self._colors[dmc]

        red = int(color[1:3], 16)
        green = int(color[3:5], 16)
        blue = int(color[5:7], 16)

        self.current_color = wx.Colour (red=red, green=green, blue=blue)

        if event:
            event.Skip()

    def _print_cell (self, event):

        mousex, mousey = self._panel.CalcUnscrolledPosition(event.GetX(), event.GetY())

        self._current_mouse_pos = (mousex, mousey)
        color = self._grid.get_color_by_mouse (mousex, mousey)
        if not color:
            color = 'None'
        else:
            color = self._find_dmc_color (color)

        self._statusbar.SetStatusText('Color: %s' % str(color))

        if event.ButtonDown(wx.MOUSE_BTN_LEFT) or event.Dragging():

            dc = wx.ClientDC (event.GetEventObject())
            self._panel.DoPrepareDC (dc)

            if self._erase_tool:
                current_color = self._grid.get_color_by_mouse(mousex, mousey)
            else:
                current_color = self.current_color

            xcell, ycell = self._grid.mouse2cell (mousex, mousey)
            self._grid.add_cell (xcell, ycell, dc, current_color, self._erase_tool)
            # Add operation for undo and redo
            op = (xcell, ycell, current_color, self._erase_tool)
            if (len(self._operations) == 0) or (not op in self._operations):
                self._operations.append (op)
                self._current_operation = len(self._operations) - 1
        elif event.Moving():
            self._timer.Start(3000,True)

        event.Skip()

    def _undo (self, event):

        if self._current_operation:
            op = self._operations[self._current_operation]

            xcell, ycell, color, erase = op
            dc = wx.ClientDC (self._panel)
            self._panel.DoPrepareDC (dc)

            if color:
                self._grid.add_cell (xcell, ycell, dc, color, not erase)
            else:
                self._grid.add_cell (xcell, ycell, dc, color, True)

            self._current_operation = self._current_operation - 1
            if self._current_operation < 0:
                self._current_operation = None

    def _redo (self, event):
        if not self._current_operation:
            self._current_operation = 0

        try:
            op = self._operations[self._current_operation]
            xcell, ycell, color, erase = op

            dc = wx.ClientDC (self._panel)
            self._panel.DoPrepareDC (dc)
            self._grid.add_cell (xcell, ycell, dc, color, erase)

            self._current_operation += 1
            if self._current_operation >= len(self._operations):
                self._current_operation = len(self._operations) - 1
        except IndexError:
            # No actions to redo
            pass

    def _set_zoom (self, event):

        if event.GetId() == xrc.XRCID('zoomout'):
            self._grid.decrease_zoom()
        elif event.GetId() == xrc.XRCID('zoomin'):
            self._grid.increase_zoom()

        size = self._grid.get_size()
        self._panel.SetVirtualSize(size)
        self._panel.FitInside()

        self._panel.SetScrollRate(size[0]/10, size[1]/10)
        self._panel.Refresh()
        event.Skip()

    def _set_edit (self, event):

        if event.GetId() == xrc.XRCID('editgrid'):
            self._erase_tool = False
        elif event.GetId() == xrc.XRCID('erase'):
            self._erase_tool = True

        self._toolbar.ToggleTool (xrc.XRCID('editgrid'), not self._erase_tool)
        self._toolbar.ToggleTool (xrc.XRCID('erase'), self._erase_tool)

        event.Skip()

    def _show_tooltip (self, event):

        x, y = self._current_mouse_pos

        color = self._grid.get_color_by_mouse (x, y)

        if color:

            red = color.Red()
            green = color.Green()
            blue = color.Blue()

            color = self._find_dmc_color (color)
            tip = wx.TipWindow (self._frame, "Color: %s\n\nRGB: (%s, %s, %s)\n\nClick tooltip to close" % (str(color),red,green,blue))

        event.Skip()
<commit_msg>Report color number too<commit_after># Copyright (c) 2012 Israel Herraiz <isra@herraiz.org>

# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import wx
from wx import xrc
from grid import Grid
from image_importer import ImageImporter
from math import sqrt

class MyApp(wx.App):

    def __init__ (self, xrcfn, colorsfn):

        self._xrcfn = xrcfn
        self._colorsfn = colorsfn
        self._scroll_rate = 10
        self._erase_tool = False
        self._grid = Grid()

        self._operations = []
        self._current_operation = None
        self._max_undo = 100

        self._timer = None
        self._current_mouse_pos = (-1, -1)

        wx.App.__init__ (self)

    def OnInit (self):

        # Colors must be imported before creating the frame
        self._import_colors ()
        self.current_color = None

        # Create main frame
        self._res = xrc.XmlResource (self._xrcfn)
        self._init_frame()

        return True

    def _import_colors (self):

        f = open(self._colorsfn, 'r')
        ls = f.readlines()
        f.close()

        self._colors = {}
        for l in ls:
            dmc, name, code = l.split(',')
            self._colors[dmc] = (code, name)

    def _find_dmc_color (self, color):
        for dmc in self._colors.keys():
            code, name = self._colors[dmc]

            red = int(code[1:3], 16)
            green = int(code[3:5], 16)
            blue = int(code[5:7], 16)

            if red == color.Red() and green == color.Green() and blue == color.Blue():
                return '%s #%s' % (name, dmc)

        return 'None'

    def _find_closest_dmc_color (self, color):

        distance = 1000

        red = color.Red()
        green = color.Green()
        blue = color.Blue()

        bestred, bestgreen, bestblue = (None, None, None)

        for dmc in self._colors.keys():
            code, name = self._colors[dmc]

            dred = int(code[1:3], 16)
            dgreen = int(code[3:5], 16)
            dblue = int(code[5:7], 16)

            r = (dred + red) / 2
            dr = dred - red
            dg = dgreen - green
            db = dblue - blue

            ndistance = (2+r/256)*dr**2 + 4*dg**2 + (2+(255-r)/256)*db**2
            ndistance = sqrt(ndistance)

            if ndistance < distance:
                distance = ndistance
                bestred, bestgreen, bestblue = (dred, dgreen, dblue)

        return wx.Colour (bestred, bestgreen, bestblue)

    def OnPaint (self, event):

        dc = wx.PaintDC (event.GetEventObject())
        dc.Clear()
        self._panel.DoPrepareDC(dc)
        self._grid.draw_grid (dc)

        event.Skip()

    def _init_frame (self):
        self._frame = self._res.LoadFrame (None,'MyMainFrame')
        self._panel = xrc.XRCCTRL (self._frame, 'MainPanel')
        self._panel.SetScrollRate (self._scroll_rate, self._scroll_rate)
        self._panel.SetVirtualSize (self._grid.get_size ())

        self._toolbar = self._frame.GetToolBar ()
        self._toolbar.ToggleTool (xrc.XRCID('editgrid'), not self._erase_tool)
        self._toolbar.ToggleTool (xrc.XRCID('erase'), self._erase_tool)

        color_choice_id = 54 # Random int
        color_list = []
        for k in self._colors.keys():
            dmc = k
            code, name = self._colors[k]

            color_list.append('%s (%s)' % (name, dmc))

        self._color_choice = wx.Choice (self._toolbar, color_choice_id, (-1,-1), (-1,-1), color_list )
        self._toolbar.AddControl (self._color_choice)
        self._change_color(None)

        self._menubar = self._frame.GetMenuBar()
        self._statusbar = self._frame.GetStatusBar()

        self._panel.Bind(wx.EVT_PAINT, self.OnPaint)
        self._panel.Bind(wx.EVT_MOUSE_EVENTS, self._print_cell)
        self._toolbar.Bind(wx.EVT_TOOL, self._undo,     id = xrc.XRCID('undo'))
        self._toolbar.Bind(wx.EVT_TOOL, self._redo,     id = xrc.XRCID('redo'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_zoom, id = xrc.XRCID('zoomout'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_zoom, id = xrc.XRCID('zoomin'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_edit, id = xrc.XRCID('editgrid'))
        self._toolbar.Bind(wx.EVT_TOOL, self._set_edit, id = xrc.XRCID('erase'))
        self._toolbar.Bind,wx.EVT_CHOICE(self, color_choice_id, self._change_color)

        self._frame.Bind(wx.EVT_MENU, self._import_image, id = xrc.XRCID('importimage'))

        self._timer = wx.Timer()
        self._timer.Bind(wx.EVT_TIMER, self._show_tooltip)

        self._frame.SetSize ((800,600))
        self._panel.FitInside()
        self._frame.SetTitle ("Stitchy Studio")

        self._timer.Start(3000,True)

        self.SetTopWindow (self._frame)
        self._frame.Show()

    def _import_image (self, event):

        path = wx.FileSelector ('Choose an image',
                               wildcard = "BMP|*.bmp|GIF|*.gif|JPEG|*.jp*g|PNG|*.png|PCX|*.pcx|TIFF|*.tiff|Other|*",
                               flags = wx.FD_FILE_MUST_EXIST,
                               parent = self._frame)

        importer = ImageImporter ()
        importer.load_image (path)
        importer.scale_image()

        height, width = importer.get_size ()

        dc = wx.ClientDC (self._panel)

        self._panel.DoPrepareDC (dc)

        for x in range (0, width):
            for y in range (0, height):
                color = importer.get_color (x, y)
                bestcolor = self._find_closest_dmc_color (color)
                self._grid.add_cell (x, y, dc, bestcolor, False)

        event.Skip()

    def _change_color (self, event):

        selection = self._color_choice.GetStringSelection()

        dmc = selection.split("(")[1].split(")")[0]

        color, _name = self._colors[dmc]

        red = int(color[1:3], 16)
        green = int(color[3:5], 16)
        blue = int(color[5:7], 16)

        self.current_color = wx.Colour (red=red, green=green, blue=blue)

        if event:
            event.Skip()

    def _print_cell (self, event):

        mousex, mousey = self._panel.CalcUnscrolledPosition(event.GetX(), event.GetY())

        self._current_mouse_pos = (mousex, mousey)
        color = self._grid.get_color_by_mouse (mousex, mousey)
        if not color:
            color = 'None'
        else:
            color = self._find_dmc_color (color)

        self._statusbar.SetStatusText('Color: %s' % str(color))

        if event.ButtonDown(wx.MOUSE_BTN_LEFT) or event.Dragging():

            dc = wx.ClientDC (event.GetEventObject())
            self._panel.DoPrepareDC (dc)

            if self._erase_tool:
                current_color = self._grid.get_color_by_mouse(mousex, mousey)
            else:
                current_color = self.current_color

            xcell, ycell = self._grid.mouse2cell (mousex, mousey)
            self._grid.add_cell (xcell, ycell, dc, current_color, self._erase_tool)
            # Add operation for undo and redo
            op = (xcell, ycell, current_color, self._erase_tool)
            if (len(self._operations) == 0) or (not op in self._operations):
                self._operations.append (op)
                self._current_operation = len(self._operations) - 1
        elif event.Moving():
            self._timer.Start(3000,True)

        event.Skip()

    def _undo (self, event):

        if self._current_operation:
            op = self._operations[self._current_operation]

            xcell, ycell, color, erase = op
            dc = wx.ClientDC (self._panel)
            self._panel.DoPrepareDC (dc)

            if color:
                self._grid.add_cell (xcell, ycell, dc, color, not erase)
            else:
                self._grid.add_cell (xcell, ycell, dc, color, True)

            self._current_operation = self._current_operation - 1
            if self._current_operation < 0:
                self._current_operation = None

    def _redo (self, event):
        if not self._current_operation:
            self._current_operation = 0

        try:
            op = self._operations[self._current_operation]
            xcell, ycell, color, erase = op

            dc = wx.ClientDC (self._panel)
            self._panel.DoPrepareDC (dc)
            self._grid.add_cell (xcell, ycell, dc, color, erase)

            self._current_operation += 1
            if self._current_operation >= len(self._operations):
                self._current_operation = len(self._operations) - 1
        except IndexError:
            # No actions to redo
            pass

    def _set_zoom (self, event):

        if event.GetId() == xrc.XRCID('zoomout'):
            self._grid.decrease_zoom()
        elif event.GetId() == xrc.XRCID('zoomin'):
            self._grid.increase_zoom()

        size = self._grid.get_size()
        self._panel.SetVirtualSize(size)
        self._panel.FitInside()

        self._panel.SetScrollRate(size[0]/10, size[1]/10)
        self._panel.Refresh()
        event.Skip()

    def _set_edit (self, event):

        if event.GetId() == xrc.XRCID('editgrid'):
            self._erase_tool = False
        elif event.GetId() == xrc.XRCID('erase'):
            self._erase_tool = True

        self._toolbar.ToggleTool (xrc.XRCID('editgrid'), not self._erase_tool)
        self._toolbar.ToggleTool (xrc.XRCID('erase'), self._erase_tool)

        event.Skip()

    def _show_tooltip (self, event):

        x, y = self._current_mouse_pos

        color = self._grid.get_color_by_mouse (x, y)

        if color:

            red = color.Red()
            green = color.Green()
            blue = color.Blue()

            color = self._find_dmc_color (color)
            tip = wx.TipWindow (self._frame, "Color: %s\n\nRGB: (%s, %s, %s)\n\nClick tooltip to close" % (str(color),red,green,blue))

        event.Skip()
<|endoftext|>


<commit_before>SECRET_KEY = "lorem ipsum"

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sites',

    'tango_shared',
    'easy_thumbnails',
    'tango_comments',
    'tests.custom_comments',
    'tests.testapp'
)

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}

SITE_ID = 1

ROOT_URLCONF = 'tango_comments.urls'

AUTH_USER_MODEL = 'auth.User'

#stripped down middleware
MIDDLEWARE = [
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
]
<commit_msg>Template settings<commit_after>SECRET_KEY = "lorem ipsum"

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sites',

    'tango_shared',
    'easy_thumbnails',
    'tango_comments',
    'tests.custom_comments',
    'tests.testapp'
)

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': ':memory:',
    }
}

SITE_ID = 1

ROOT_URLCONF = 'tango_comments.urls'

AUTH_USER_MODEL = 'auth.User'

#stripped down middleware
MIDDLEWARE = [
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
]

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'APP_DIRS': True,
    },
]
<|endoftext|>


<commit_before>"""Representation of a lattice object which contains all the elements of the
    machine.
"""
import numpy
from pytac.exceptions import LatticeException

class Lattice(object):
    """Representation of a lattice.

    Represents a lattice object that contains all elements of the ring. It has
    a name and a control system to be used for unit conversion.

    **Attributes:**

    Attributes:
        name (str): The name of the lattice.

    .. Private Attributes:
           _lattice (list): The list of all the element objects in the lattice.
           _cs (ControlSystem): The control system used to store the values on
                                 a PV.
           _energy (int): The total energy of the lattice.
           _model (Model): A pytac model object associated with the lattice.
    """
    def __init__(self, name, energy):
        """.. The constructor method for the class, called whenever a 'Lattice'
               object is constructed.

        Args:
            name (str): The name of the lattice.
            energy (int): The total energy of the lattice.

        **Methods:**
        """
        self.name = name
        self._lattice = []
        self._energy = energy
        self._models = {}
        self._uc = {}

    def set_model(self, model, model_type):
        """Add a model to the lattice.

        Args:
            model (Model): instance of Model.
            model_type (str): EpicsModel or ATModel.
        """
        self.models[model_type] = model

    def get_fields(self):
        """Get the fields defined on the lattice.

        Includes all fields defined by all models.

        Returns:
            set: A sequence of all the fields defined on the lattice.
        """
        fields = set()
        for model in self._models:
            fields.update(self._models[model].get_fields())
        return fields

    def add_device(self, field, device, uc):
        """Add device and unit conversion objects to a given field.

        A DeviceModel must be set before calling this method.

        Args:
            field (str): The key to store the unit conversion and device
                          objects.
            device (Device): The device object used for this field.
            uc (UnitConv): The unit conversion object used for this field.

        Raises:
            KeyError: if no DeviceModel is set.
        """
        self._models[pytac.LIVE].add_device(field, device)
        self._uc[field] = uc

    def get_device(self, field):
        """Get the device for the given field.

        A DeviceModel must be set before calling this method.

        Args:
            field (str): The lookup key to find the device on the lattice.

        Returns:
            Device: The device on the given field.

        Raises:
            KeyError: if no DeviceModel is set.
        """
        return self._models[pytac.LIVE].get_device(field)

    def get_unitconv(self, field):
        """Get the unit conversion option for the specified field.

        Args:
            field (str): The field associated with this conversion.

        Returns:
            UnitConv: The object associated with the specified field.

        Raises:
            KeyError: if no unit conversion object is present.
        """
        return self._uc[field]

    def get_value(self, field, handle=pytac.RB, units=pytac.ENG,
                  model=pytac.LIVE):
        """Get the value for a field on the lattice.

        Args:
            field (str): The requested field.
            handle (str): pytac.SP or pytac.RB.
            units (str): pytac.ENG or pytac.PHYS returned.
            model (str): pytac.LIVE or pytac.SIM.

        Returns:
            float: The value of the requested field

        Raises:
            DeviceException: if there is no device on the given field.
            FieldException: if the lattice does not have the specified field.
        """
        try:
            model = self._models[model]
            value = model.get_value(field, handle)
            return self._uc[field].convert(value, origin=model.units,
                                           target=units)
        except KeyError:
            raise DeviceException('No model type {} on lattice {}'.format(model,
                                                                          self))
        except FieldException:
            raise FieldException('No field {} on lattice {}'.format(field, self))

    def set_value(self, field, value, handle=pytac.SP, units=pytac.ENG,
                  model=pytac.LIVE):
        """Set the value for a field.

        This value can be set on the machine or the simulation.

        Args:
            field (str): The requested field.
            value (float): The value to set.
            handle (str): pytac.SP or pytac.RB.
            units (str): pytac.ENG or pytac.PHYS.
            model (str): pytac.LIVE or pytac.SIM.

        Raises:
            DeviceException: if arguments are incorrect.
            FieldException: if the lattice does not have the specified field.
        """
        if handle != pytac.SP:
            raise HandleException('Must write using {}'.format(pytac.SP))
        try:
            model = self._models[model]
        except KeyError:
            raise DeviceException('No model type {} on lattice {}'.format(model,
                                                                          self))
        try:
            value = self._uc[field].convert(value, origin=units, target=model.units)
            model.set_value(field, value)
        except KeyError:
            raise FieldException('No field {} on lattice {}'.format(model, self))
        except FieldException:
            raise FieldException('No field {} on lattice {}'.format(field, self))

    def get_energy(self):
        """Function to get the total energy of the lattice.

        Returns:
            int: energy of the lattice.
        """
        return self._energy

    def __getitem__(self, n):
        """Get the (n + 1)th element of the lattice - i.e. index 0 represents
        the first element in the lattice.

        Args:
            n (int): index.

        Returns:
            Element: indexed element.
        """
        return self._lattice[n]

    def __len__(self):
        """The number of elements in the lattice.

        When using the len function returns the number of elements in
        the lattice.

        Returns:
            int: The number of elements in the lattice.
        """
        return len(self._lattice)

    def get_length(self):
        """Returns the length of the lattice.

        Returns:
            float: The length of the lattice.
        """
        total_length = 0
        for e in self._lattice:
            total_length += e.length
        return total_length

    def add_element(self, element):
        """Append an element to the lattice.

        Args:
            element (Element): element to append.
        """
        self._lattice.append(element)

    def get_elements(self, family=None, cell=None):
        """Get the elements of a family from the lattice.

        If no family is specified it returns all elements. Elements are
        returned in the order they exist in the ring.

        Args:
            family (str): requested family.
            cell (int): restrict elements to those in the specified cell.

        Returns:
            list: list containing all elements of the specified family.
        """
        elements = []
        if family is None:
            elements = self._lattice

        for element in self._lattice:
            if family in element.families:
                elements.append(element)

        if cell is not None:
            elements = [e for e in elements if e.cell == cell]

        return elements

    def get_all_families(self):
        """Get all families of elements in the lattice.

        Returns:
            set: all defined families.
        """
        families = set()
        for element in self._lattice:
            families.update(element.families)

        return families

    def get_s(self, elem):
        """Find the s position of an element in the lattice.

        Note that the given element must exist in the lattice.

        Args:
            elem (Element): The element that the position is being asked for.

        Returns:
            float: The position of the given element.

        Raises:
            LatticeException: if element doesn't exist in the lattice.
        """
        s_pos = 0
        for el in self._lattice:
            if el is not elem:
                s_pos += el.length
            else:
                return s_pos
        raise LatticeException('Element {} not in lattice {}'.format(elem, self))

    def get_family_s(self, family):
        """Get s positions for all elements from the same family.

        Args:
            family (str): requested family.

        Returns:
            list: list of s positions for each element.
        """
        elements = self.get_elements(family)
        s_positions = []
        for element in elements:
            s_positions.append(self.get_s(element))
        return s_positions

    def get_element_devices(self, family, field):
        """Get devices for a specific field for elements in the specfied
        family.

        Typically all elements of a family will have devices associated with
        the same fields - for example, BPMs each have a device for fields 'x'
        and 'y'.

        Args:
            family (str): family of elements.
            field (str): field specifying the devices.

        Returns:
            list: devices for specified family and field.
        """
        elements = self.get_elements(family)
        devices = []
        for element in elements:
            try:
                devices.append(element.get_device(field))
            except KeyError:
                pass

        return devices

    def get_element_device_names(self, family, field):
        """Get the names for devices attached to a specific field for elements
        in the specfied family.

        Typically all elements of a family will have devices associated with
        the same fields - for example, BPMs each have a device for fields 'x'
        and 'y'.

        Args:
            family (str): family of elements.
            field (str): field specifying the devices.

        Returns:
            list: device names for specified family and field.
        """
        devices = self.get_devices(family, field)
        return [device.name for device in devices]

    def get_element_values(self, family, field, handle, dtype=None):
        """Get all values for a family and field.

        Args:
            family (str): family to request the values of.
            field (str): field to request values for.
            handle (str): pytac.RB or pytac.SP.
            dtype (numpy.dtype): if None, return a list. If not None, return a
                                  numpy array of the specified type.

        Returns:
            list or numpy array: sequence of values.
        """
        elements = self.get_elements(family)
        values = [element.get_value(field, handle) for element in elements]
        if dtype is not None:
            values = numpy.array(values, dtype=dtype)
        return values

    def set_element_values(self, family, field, values):
        """Sets the values for a family and field.

        The PVs are determined by family and device. Note that only setpoint
        PVs can be modified.

        Args:
            family (str): family on which to set values.
            field (str):  field to set values for.
            values (sequence): A list of values to assign.

        Raises:
            LatticeException: if the given list of values doesn't match the
                               number of elements in the family.
        """
        elements = self.get_elements(family)
        if len(elements) != len(values):
            raise LatticeException("Number of elements in given array must be"
                                   " equal to the number of elements in the "
                                   "family")
        for element, value in zip(elements, values):
            element.set_value(field, value)
<commit_msg>Import correct exceptions for new lattice methodds.<commit_after>"""Representation of a lattice object which contains all the elements of the
    machine.
"""
import numpy
from pytac.exceptions import LatticeException, DeviceException, FieldException, HandleException

class Lattice(object):
    """Representation of a lattice.

    Represents a lattice object that contains all elements of the ring. It has
    a name and a control system to be used for unit conversion.

    **Attributes:**

    Attributes:
        name (str): The name of the lattice.

    .. Private Attributes:
           _lattice (list): The list of all the element objects in the lattice.
           _cs (ControlSystem): The control system used to store the values on
                                 a PV.
           _energy (int): The total energy of the lattice.
           _model (Model): A pytac model object associated with the lattice.
    """
    def __init__(self, name, energy):
        """.. The constructor method for the class, called whenever a 'Lattice'
               object is constructed.

        Args:
            name (str): The name of the lattice.
            energy (int): The total energy of the lattice.

        **Methods:**
        """
        self.name = name
        self._lattice = []
        self._energy = energy
        self._models = {}
        self._uc = {}

    def set_model(self, model, model_type):
        """Add a model to the lattice.

        Args:
            model (Model): instance of Model.
            model_type (str): EpicsModel or ATModel.
        """
        self.models[model_type] = model

    def get_fields(self):
        """Get the fields defined on the lattice.

        Includes all fields defined by all models.

        Returns:
            set: A sequence of all the fields defined on the lattice.
        """
        fields = set()
        for model in self._models:
            fields.update(self._models[model].get_fields())
        return fields

    def add_device(self, field, device, uc):
        """Add device and unit conversion objects to a given field.

        A DeviceModel must be set before calling this method.

        Args:
            field (str): The key to store the unit conversion and device
                          objects.
            device (Device): The device object used for this field.
            uc (UnitConv): The unit conversion object used for this field.

        Raises:
            KeyError: if no DeviceModel is set.
        """
        self._models[pytac.LIVE].add_device(field, device)
        self._uc[field] = uc

    def get_device(self, field):
        """Get the device for the given field.

        A DeviceModel must be set before calling this method.

        Args:
            field (str): The lookup key to find the device on the lattice.

        Returns:
            Device: The device on the given field.

        Raises:
            KeyError: if no DeviceModel is set.
        """
        return self._models[pytac.LIVE].get_device(field)

    def get_unitconv(self, field):
        """Get the unit conversion option for the specified field.

        Args:
            field (str): The field associated with this conversion.

        Returns:
            UnitConv: The object associated with the specified field.

        Raises:
            KeyError: if no unit conversion object is present.
        """
        return self._uc[field]

    def get_value(self, field, handle=pytac.RB, units=pytac.ENG,
                  model=pytac.LIVE):
        """Get the value for a field on the lattice.

        Args:
            field (str): The requested field.
            handle (str): pytac.SP or pytac.RB.
            units (str): pytac.ENG or pytac.PHYS returned.
            model (str): pytac.LIVE or pytac.SIM.

        Returns:
            float: The value of the requested field

        Raises:
            DeviceException: if there is no device on the given field.
            FieldException: if the lattice does not have the specified field.
        """
        try:
            model = self._models[model]
            value = model.get_value(field, handle)
            return self._uc[field].convert(value, origin=model.units,
                                           target=units)
        except KeyError:
            raise DeviceException('No model type {} on lattice {}'.format(model,
                                                                          self))
        except FieldException:
            raise FieldException('No field {} on lattice {}'.format(field, self))

    def set_value(self, field, value, handle=pytac.SP, units=pytac.ENG,
                  model=pytac.LIVE):
        """Set the value for a field.

        This value can be set on the machine or the simulation.

        Args:
            field (str): The requested field.
            value (float): The value to set.
            handle (str): pytac.SP or pytac.RB.
            units (str): pytac.ENG or pytac.PHYS.
            model (str): pytac.LIVE or pytac.SIM.

        Raises:
            DeviceException: if arguments are incorrect.
            FieldException: if the lattice does not have the specified field.
        """
        if handle != pytac.SP:
            raise HandleException('Must write using {}'.format(pytac.SP))
        try:
            model = self._models[model]
        except KeyError:
            raise DeviceException('No model type {} on lattice {}'.format(model,
                                                                          self))
        try:
            value = self._uc[field].convert(value, origin=units, target=model.units)
            model.set_value(field, value)
        except KeyError:
            raise FieldException('No field {} on lattice {}'.format(model, self))
        except FieldException:
            raise FieldException('No field {} on lattice {}'.format(field, self))

    def get_energy(self):
        """Function to get the total energy of the lattice.

        Returns:
            int: energy of the lattice.
        """
        return self._energy

    def __getitem__(self, n):
        """Get the (n + 1)th element of the lattice - i.e. index 0 represents
        the first element in the lattice.

        Args:
            n (int): index.

        Returns:
            Element: indexed element.
        """
        return self._lattice[n]

    def __len__(self):
        """The number of elements in the lattice.

        When using the len function returns the number of elements in
        the lattice.

        Returns:
            int: The number of elements in the lattice.
        """
        return len(self._lattice)

    def get_length(self):
        """Returns the length of the lattice.

        Returns:
            float: The length of the lattice.
        """
        total_length = 0
        for e in self._lattice:
            total_length += e.length
        return total_length

    def add_element(self, element):
        """Append an element to the lattice.

        Args:
            element (Element): element to append.
        """
        self._lattice.append(element)

    def get_elements(self, family=None, cell=None):
        """Get the elements of a family from the lattice.

        If no family is specified it returns all elements. Elements are
        returned in the order they exist in the ring.

        Args:
            family (str): requested family.
            cell (int): restrict elements to those in the specified cell.

        Returns:
            list: list containing all elements of the specified family.
        """
        elements = []
        if family is None:
            elements = self._lattice

        for element in self._lattice:
            if family in element.families:
                elements.append(element)

        if cell is not None:
            elements = [e for e in elements if e.cell == cell]

        return elements

    def get_all_families(self):
        """Get all families of elements in the lattice.

        Returns:
            set: all defined families.
        """
        families = set()
        for element in self._lattice:
            families.update(element.families)

        return families

    def get_s(self, elem):
        """Find the s position of an element in the lattice.

        Note that the given element must exist in the lattice.

        Args:
            elem (Element): The element that the position is being asked for.

        Returns:
            float: The position of the given element.

        Raises:
            LatticeException: if element doesn't exist in the lattice.
        """
        s_pos = 0
        for el in self._lattice:
            if el is not elem:
                s_pos += el.length
            else:
                return s_pos
        raise LatticeException('Element {} not in lattice {}'.format(elem, self))

    def get_family_s(self, family):
        """Get s positions for all elements from the same family.

        Args:
            family (str): requested family.

        Returns:
            list: list of s positions for each element.
        """
        elements = self.get_elements(family)
        s_positions = []
        for element in elements:
            s_positions.append(self.get_s(element))
        return s_positions

    def get_element_devices(self, family, field):
        """Get devices for a specific field for elements in the specfied
        family.

        Typically all elements of a family will have devices associated with
        the same fields - for example, BPMs each have a device for fields 'x'
        and 'y'.

        Args:
            family (str): family of elements.
            field (str): field specifying the devices.

        Returns:
            list: devices for specified family and field.
        """
        elements = self.get_elements(family)
        devices = []
        for element in elements:
            try:
                devices.append(element.get_device(field))
            except KeyError:
                pass

        return devices

    def get_element_device_names(self, family, field):
        """Get the names for devices attached to a specific field for elements
        in the specfied family.

        Typically all elements of a family will have devices associated with
        the same fields - for example, BPMs each have a device for fields 'x'
        and 'y'.

        Args:
            family (str): family of elements.
            field (str): field specifying the devices.

        Returns:
            list: device names for specified family and field.
        """
        devices = self.get_devices(family, field)
        return [device.name for device in devices]

    def get_element_values(self, family, field, handle, dtype=None):
        """Get all values for a family and field.

        Args:
            family (str): family to request the values of.
            field (str): field to request values for.
            handle (str): pytac.RB or pytac.SP.
            dtype (numpy.dtype): if None, return a list. If not None, return a
                                  numpy array of the specified type.

        Returns:
            list or numpy array: sequence of values.
        """
        elements = self.get_elements(family)
        values = [element.get_value(field, handle) for element in elements]
        if dtype is not None:
            values = numpy.array(values, dtype=dtype)
        return values

    def set_element_values(self, family, field, values):
        """Sets the values for a family and field.

        The PVs are determined by family and device. Note that only setpoint
        PVs can be modified.

        Args:
            family (str): family on which to set values.
            field (str):  field to set values for.
            values (sequence): A list of values to assign.

        Raises:
            LatticeException: if the given list of values doesn't match the
                               number of elements in the family.
        """
        elements = self.get_elements(family)
        if len(elements) != len(values):
            raise LatticeException("Number of elements in given array must be"
                                   " equal to the number of elements in the "
                                   "family")
        for element, value in zip(elements, values):
            element.set_value(field, value)
<|endoftext|>