<commit_before>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011, Paul D. Nation & Robert J. Johansson
#
###########################################################################
from scipy import *
from Qobj import Qobj
from basis import basis
from operators import destroy

def coherent(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    x=basis(N,0)
    a=destroy(N)
    D=(alpha*a.dag()-conj(alpha)*a).expm()
    return D*x

def coherent_dm(N, alpha):
    """
    Generate the density matrix representation of a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    psi = coherent(N,alpha)
    return psi * psi.dag()

def coherent_fast(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    data = zeros([N,1])
    n = arange(N)
    data[:,0] = exp(-(abs(alpha)**2)/2.0)*(alpha**(n))/sqrt(factorial(n))
    return Qobj(data)

def coherent_dm_fast(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    psi = coherent_fast(N, alpha)
    return psi * psi.dag()

def fock_dm(N, m):
    """
    Generate the density matrix for a fock state

    @param N the number of states
    @param m the fock state number
    """
    psi = basis(N, m)
    return psi * psi.dag()

def fock(N, m):
    """
    Generate a state vector for a fock state

    @param N the number of states
    @param m the fock state number
    """
    return basis(N, m)

def thermal_dm(N, n):
    """
    Generates the density matrix for a thermal state of n particles

    @param N: the number of states
    @param n: expectational value for number of particles in thermal state
    """

    i=arange(N)
    rm = diag((1.0+n)**(-1.0)*(n/(1.0+n))**(i)) #populates diagonal terms (the only nonzero terms in matrix)
    return Qobj(rm)

if __name__ == "__main__":
    print (coherent(5,.1)*coherent(5,.1).dag()).tr()

<commit_msg>added ket to density matrix function<commit_after>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011, Paul D. Nation & Robert J. Johansson
#
###########################################################################
from scipy import *
from Qobj import Qobj
from basis import basis
from operators import destroy

def coherent(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    x=basis(N,0)
    a=destroy(N)
    D=(alpha*a.dag()-conj(alpha)*a).expm()
    return D*x

def coherent_dm(N, alpha):
    """
    Generate the density matrix representation of a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    psi = coherent(N,alpha)
    return psi * psi.dag()

def coherent_fast(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    data = zeros([N,1])
    n = arange(N)
    data[:,0] = exp(-(abs(alpha)**2)/2.0)*(alpha**(n))/sqrt(factorial(n))
    return Qobj(data)

def coherent_dm_fast(N,alpha):
    """
    Generate a coherent state
    @param N the number of states
    @param alpha the coherent state amplitude (complex)
    """
    psi = coherent_fast(N, alpha)
    return psi * psi.dag()

def fock_dm(N, m):
    """
    Generate the density matrix for a fock state

    @param N the number of states
    @param m the fock state number
    """
    psi = basis(N, m)
    return psi * psi.dag()

def fock(N, m):
    """
    Generate a state vector for a fock state

    @param N the number of states
    @param m the fock state number
    """
    return basis(N, m)

def thermal_dm(N, n):
    """
    Generates the density matrix for a thermal state of n particles

    @param N: the number of states
    @param n: expectational value for number of particles in thermal state
    """

    i=arange(N)
    rm = diag((1.0+n)**(-1.0)*(n/(1.0+n))**(i)) #populates diagonal terms (the only nonzero terms in matrix)
    return Qobj(rm)

def ket2dm(Q):
    """
    Takes input ket or bra vector and returns density matrix formed by outer product.
    @param Q: Ket or bra vector
    @return dm: Density matrix formed by outer product
    """
    if Q.type=='ket':
        out=Q*Q.dag()
    elif Q.type=='bra':
        out=Q.dag()*Q
    else:
        raise TypeError("Input is not a ket or bra vector.")
    return Qobj(out)

if __name__ == "__main__":
    print (coherent(5,.1)*coherent(5,.1).dag()).tr()
    print ket2dm(coherent(3,.1).dag())

<|endoftext|>


<commit_before># This file is updated automatically, do not edit.
version = '0.4.0'
<commit_msg>Updated version file<commit_after># This file is updated automatically, do not edit.
version = '0.5.0'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

#   Copyright 2014 David Marin
#
#   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.
"""Merge and correct company information."""
import logging
import re
from collections import defaultdict

from .brand import get_brands_for_company
from .category import get_brand_categories
from .category import get_company_categories
from .db import COMPANIES_PREFIX
from .db import output_row
from .db import select_all_companies
from .db import select_brand_ratings
from .db import select_company_ratings
from .db import select_company
from .norm import group_by_keys
from .norm import merge_dicts
from .norm import norm_with_variants
from .norm import simplify_whitespace

log = logging.getLogger(__name__)

# various misspellings of company names
COMPANY_CORRECTIONS = {
    'Delta Airlines': 'Delta Air Lines',
    'GEPA- The Fairtrade Company': 'GEPA - The Fairtrade Company',
    'V.F. Corporation': 'VF Corporation',
    'Wolverine Worldwide': 'Wolverine World Wide',
    'Woolworths Australia': 'Woolworths Limited',
}

# Name changes. May eventually want separate logic for this.
COMPANY_CORRECTIONS.update({
    'Sweet Earth Chocolates': 'Mama Ganache',  # renamed in 2012
    'RIM': 'BlackBerry Limited', # renamed in 2013
    'Research In Motion': 'BlackBerry Limited',
})

COMPANY_ALIASES = [
    ['AB Electrolux', 'Electrolux'],
    ['Anheuser-Busch', 'Anheuser-Busch InBev'],
    ['ASUS', 'ASUSTeK Computer'],
    ['Disney', 'The Walt Disney Company', 'The Walt Disney Co.'],
    ['Gap', 'The Gap'],  # might solve this with "Gap" brand?
    ['HP', 'Hewlett-Packard'],
    ['HTC Electronics', 'HTC'],
    ['Illy', u'illycaffè'],
    ['JetBlue', 'JetBlue Airways'],
    ['Kellogg', "Kellogg's"],  # might solve this with a brand?
    ['Lidl', 'Lidl Stiftung'],
    ['LG', 'LGE', 'LG Electronics'],
    ['New Look', 'New Look Retailers'],
    ['Philips', 'Royal Philips', 'Royal Philips Electronics'],
    ['Rivers Australia', 'Rivers (Australia) Pty Ltd'],
    # technically, Wells Fargo Bank, N.A. is a subsidiary of Wells Fargo
    # the multinational. Not worrying about this now and I don't think this is
    # what they meant anyway.
    ['Wells Fargo', 'Wells Fargo Bank'],
    ['Whole Foods', 'Whole Foods Market'],
    ['Wibra', 'Wibra Supermarkt'],
]

# "The X Co." -- okay to strip
THE_X_CO_RE = re.compile(
    r'^The (?P<company>.*) Co\.$')

# X [&] Co. -- okay to strip
X_CO_RE = re.compile(
    r'^(?P<company>.*?)( \&)? Co\.$'
)

# "The X Company" -- okay for matching, but don't use as short name
THE_X_COMPANY_RE = re.compile(
    r'^The (?P<company>.*) (Co\.|Corporation|Cooperative|Company|Group)$')

# "X Company" -- okay for matching, but don't use as short name
X_COMPANY_RE = re.compile(
    r'^(?P<company>.*) ('
    r'Brands'
    r'|Co.'
    r'|Company'
    r'|Corporation'
    r'|Enterprises'
    r'|Group'
    r'|Gruppe'
    r'|Holdings?'
    r'|Products'
    r'|Ventures?'
    r')$'
)

COMPANY_TYPE_RE = re.compile(
    r'^(?P<company>.*?)(?P<intl1> International)?,? (?P<type>'
    r'A\.?& S\. Klein GmbH \& Co\. KG'
    r'|A/S'
    r'|AB'
    r'|AG'
    r'|AS'
    r'|Ab'
    r'|BV'
    r'|B\.V\.'
    r'|B.V. Nederland'
    r'|C\.V\.'
    r'|Corp.'
    r'|GmbH \& C[oO]\. [oO]HG'
    r'|GmbH \& Co\. ?KG\.?'  # handle typo: Lukas Meindl GmbH & Co.KG
    r'|GmbH \& Co\. KGaA'
    r'|GmbH'
    r'|Inc\.?'
    r'|Incorporated'
    r'|International'
    r'|KG\.?'
    r'|Llc'
    r'|LLC'
    r'|LLP'
    r'|LP'
    r'|Limited'
    r'|Llp'
    r'|Ltd\.?'
    r'|Ltda\.?'
    r'|nv'
    r'|NV'
    r'|N\.V\.'
    r'|PBC'  # "Public Benefit Corporation"? Only on B Corp site
    r'|PLC'
    r'|P\.C\.'
    r'|P\.C\.'
    r'|Pty Ltd'
    r'|Pty'
    r'|S.\L\.'
    r'|SA'
    r'|SAPI DE CV SOFOM ENR'
    r'|SARL'
    r'|SE'
    r'|S\.A\.?'
    r'|S\.A\.U\.'
    r'|S\.R\.L\.'
    r'|S\.p\.A\.'
    r'|Sarl'
    r'|SpA'
    r'|b\.v\.'
    r'|gmbh'
    r'|inc'
    r'|plc'
    r')(?P<intl2> International)?$'
)

COMPANY_TYPE_CORRECTIONS = {
    'BV': 'B.V.',
    'Incorporated': 'Inc',
    'Llc': 'LLC',
    'Llp': 'LLP',
    'NV': 'N.V.',
    'S.A': 'S.A.',
    'Sarl': 'SARL',
    'SpA': 'S.p.A.',
    'b.v.': 'B.V.',
    'gmbh': 'GmbH',
    'inc': 'Inc',
    'nv': 'N.V.',
}

UNSTRIPPABLE_COMPANY_TYPES = {
    'LLP',
}

UNSTRIPPABLE_COMPANIES = {
    'Globe International',
    'Woolworths Limited',
}

def handle_matched_company(cd, category_map):
    """Take in a company dictionary from match_companies(), and
    output rows about that company to the database.

    This handles writing all rows except for the "campaign" table.

    (or will eventually)
    """
    # get brands
    brand_to_row, brand_map = get_brands_for_company(cd['keys'])

    # pick a canonical name for this company
    company_canonical, company_full = name_company(cd, set(brand_to_row))

    if company_canonical == company_full:
        log.info(company_canonical)
    else:
        log.info(u'{} ({})'.format(company_canonical, company_full))

    # merge company rows
    company_row = {}

    for campaign_id, company in sorted(cd['keys']):
         company_row.update(select_company(campaign_id, company))

    del company_row['campaign_id']  # should be at least one match

    company_row['company'] = company_canonical
    company_row['company_full'] = company_full

    # store company row
    output_row(company_row, 'company')

    # store company map and ratings
    for campaign_id, campaign_company in cd['keys']:
        # map
        map_row = dict(
            campaign_id=campaign_id, campaign_company=campaign_company,
            company=company_canonical)
        output_row(map_row, 'campaign_company_map')

        # ratings (may be more than one because of `scope`)
        for rating_row in select_company_ratings(
                campaign_id, campaign_company):
            rating_row['company'] = company_canonical
            output_row(rating_row, 'campaign_company_rating')

    # store company categories
    for cat_row in get_company_categories(
            company_canonical, cd['keys'], category_map):
        output_row(cat_row, 'company_category')

    # store brands
    brand_rows = sorted(brand_to_row.itervalues(), key=lambda r: r['brand'])
    for brand_row in brand_rows:
        brand_row['company'] = company_canonical
        output_row(brand_row, 'brand')

    # store brand map, rating
    campaign_to_company = dict(cd['keys'])
    brand_to_keys = defaultdict(set)

    for (campaign_id, campaign_brand), brand_canonical in brand_map.items():
        campaign_company = campaign_to_company[campaign_id]

        brand_to_keys[brand_canonical].add(
            (campaign_id, campaign_company, campaign_brand))

        # map
        map_row = dict(
            campaign_id=campaign_id, campaign_company=campaign_company,
            campaign_brand=campaign_brand, company=company_canonical,
            brand=brand_canonical)
        output_row(map_row, 'campaign_brand_map')

        # ratings
        for rating_row in select_brand_ratings(
                campaign_id, campaign_company, campaign_brand):
            rating_row['company'] = company_canonical
            rating_row['brand'] = brand_canonical
            output_row(rating_row, 'campaign_brand_rating')

    # store brand categories
    for brand_canonical, keys in sorted(brand_to_keys.items()):
        for cat_row in get_brand_categories(
                company_canonical, brand_canonical, keys, category_map):
            output_row(cat_row, 'brand_category')

def name_company(cd, brands=()):
    """Return short and full names for a company."""
    short_name = pick_short_name_from_variants(cd['display_names'])

    # if matching short name is shorter and matches a brand for one
    # of the companies, use that
    matching_short_name = pick_short_name_from_variants(
        cd['display_names'] | cd['matching_names'])

    if (matching_short_name != short_name and
        matching_short_name in brands):
        short_name = matching_short_name

    full_name = pick_full_name_from_variants(cd['display_names'])

    return short_name, full_name

def match_companies(companies_with_campaign_ids=None, aliases=None):
    """Match up similar company names.

    companies -- sequence of (campaign_id, company). Defaults to
                 get_companies_with_campaign_ids()
    aliases -- sequence of list of matching company names. Defaults to
               COMPANY_ALIASES

    Yields company dicts, which contain:
    keys -- set of tuples of (campaign_id, original_company_name)
    display_names -- set of names appropriate for display
    matching_names -- set of names for matching
    """
    if companies_with_campaign_ids is None:
        companies_with_campaign_ids = select_all_companies()
    if aliases is None:
        aliases = COMPANY_ALIASES

    to_merge = []

    # use variants for matching
    for variants in aliases:
        variants = set(variants)
        to_merge.append({'keys': set(), 'display_names': set(),
                         'matching_names': variants})

    # add in companies
    for campaign_id, company in companies_with_campaign_ids:
        if not company:  # skip blank company names
            continue
        keys = {(campaign_id, company)}
        display, matching = get_company_name_variants(company)

        to_merge.append({'keys': keys, 'display_names': set(display),
                         'matching_names': set(matching)})

    # match up company dicts and merge them together
    def keyfunc(cd):
        normed_variants = set()
        for name in cd['matching_names']:
            normed_variants.update(norm_with_variants(name))
        return normed_variants

    # and merge them together
    for cd_group in group_by_keys(to_merge, keyfunc):
        merged = merge_dicts(cd_group)

        if not merged['keys'] and merged['display_names']:
            # this can happen if hard-coded variants don't match anything
            log.warn('orphaned company dict: {}'.format(repr(merged)))
        else:
            yield merged

def pick_short_name_from_variants(variants):
    # shortest name, ties broken by, not all lower/upper, has accents
    return sorted(variants,
                  key=lambda v: (len(v), v == v.lower(), v == v.upper(),
                                 -len(v.encode('utf8'))))[0]

def pick_full_name_from_variants(variants):
    # longest name, ties broken by, not all lower/upper, has accents
    return sorted(variants,
                  key=lambda v: (-len(v), v == v.lower(), v == v.upper(),
                                 -len(v.encode('utf8'))))[0]

def get_company_name_variants(company):
    """Convert a company name to a set of variants for display,
    and a set of variants for matching only.
    """

    display_variants = set()  # usable as display name
    matching_variants = set()  # usable for matching

    def handle(company):
        company = simplify_whitespace(company)

        company = COMPANY_CORRECTIONS.get(company) or company

        m = COMPANY_TYPE_RE.match(company)
        if m:
            company = m.group('company')
            intl1 = m.group('intl1') or ''
            c_type = m.group('type')
            intl2 = m.group('intl2') or ''
            c_type = COMPANY_TYPE_CORRECTIONS.get(c_type) or c_type
            c_full = company + intl1 + ' ' + c_type + intl2
            display_variants.add(c_full)

            if (c_type in UNSTRIPPABLE_COMPANY_TYPES or
                c_full in UNSTRIPPABLE_COMPANIES):
                return

        display_variants.add(company)

        m = THE_X_CO_RE.match(company)
        if m:
            display_variants.add(m.group('company'))
            return

        m = X_CO_RE.match(company)
        if m:
            display_variants.add(m.group('company'))
            return

        m = THE_X_COMPANY_RE.match(company)
        if m:
            matching_variants.add(m.group('company'))
            return

        m = X_COMPANY_RE.match(company)
        if m:
            matching_variants.add(m.group('company'))
            return

    handle(company)

    matching_variants.update(display_variants)

    # handle slashes in company names
    for mv in list(matching_variants):
        if '/' in mv:
            matching_variants.update(
                filter(None, (part.strip() for part in mv.split('/'))))

    return display_variants, matching_variants

def correct_company_type(c_type):
    if c_type in COMPANY_TYPE_CORRECTIONS:
        return COMPANY_TYPE_CORRECTIONS[c_type]
    elif c_type.endswith('.') and c_type.count('.') == 1:
        return c_type[:-1]
    else:
        return c_type
<commit_msg>use new name for Puma SE<commit_after># -*- coding: utf-8 -*-

#   Copyright 2014 David Marin
#
#   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.
"""Merge and correct company information."""
import logging
import re
from collections import defaultdict

from .brand import get_brands_for_company
from .category import get_brand_categories
from .category import get_company_categories
from .db import COMPANIES_PREFIX
from .db import output_row
from .db import select_all_companies
from .db import select_brand_ratings
from .db import select_company_ratings
from .db import select_company
from .norm import group_by_keys
from .norm import merge_dicts
from .norm import norm_with_variants
from .norm import simplify_whitespace

log = logging.getLogger(__name__)

# various misspellings of company names
COMPANY_CORRECTIONS = {
    'Delta Airlines': 'Delta Air Lines',
    'GEPA- The Fairtrade Company': 'GEPA - The Fairtrade Company',
    'PUMA AG Rudolf Dassler Sport': 'Puma SE',
    'V.F. Corporation': 'VF Corporation',
    'Wolverine Worldwide': 'Wolverine World Wide',
    'Woolworths Australia': 'Woolworths Limited',
}

# Name changes. May eventually want separate logic for this.
COMPANY_CORRECTIONS.update({
    'Sweet Earth Chocolates': 'Mama Ganache',  # renamed in 2012
    'RIM': 'BlackBerry Limited', # renamed in 2013
    'Research In Motion': 'BlackBerry Limited',
})

COMPANY_ALIASES = [
    ['AB Electrolux', 'Electrolux'],
    ['Anheuser-Busch', 'Anheuser-Busch InBev'],
    ['ASUS', 'ASUSTeK Computer'],
    ['Disney', 'The Walt Disney Company', 'The Walt Disney Co.'],
    ['Gap', 'The Gap'],  # might solve this with "Gap" brand?
    ['HP', 'Hewlett-Packard'],
    ['HTC Electronics', 'HTC'],
    ['Illy', u'illycaffè'],
    ['JetBlue', 'JetBlue Airways'],
    ['Kellogg', "Kellogg's"],  # might solve this with a brand?
    ['Lidl', 'Lidl Stiftung'],
    ['LG', 'LGE', 'LG Electronics'],
    ['New Look', 'New Look Retailers'],
    ['Philips', 'Royal Philips', 'Royal Philips Electronics'],
    ['Rivers Australia', 'Rivers (Australia) Pty Ltd'],
    # technically, Wells Fargo Bank, N.A. is a subsidiary of Wells Fargo
    # the multinational. Not worrying about this now and I don't think this is
    # what they meant anyway.
    ['Wells Fargo', 'Wells Fargo Bank'],
    ['Whole Foods', 'Whole Foods Market'],
    ['Wibra', 'Wibra Supermarkt'],
]

# "The X Co." -- okay to strip
THE_X_CO_RE = re.compile(
    r'^The (?P<company>.*) Co\.$')

# X [&] Co. -- okay to strip
X_CO_RE = re.compile(
    r'^(?P<company>.*?)( \&)? Co\.$'
)

# "The X Company" -- okay for matching, but don't use as short name
THE_X_COMPANY_RE = re.compile(
    r'^The (?P<company>.*) (Co\.|Corporation|Cooperative|Company|Group)$')

# "X Company" -- okay for matching, but don't use as short name
X_COMPANY_RE = re.compile(
    r'^(?P<company>.*) ('
    r'Brands'
    r'|Co.'
    r'|Company'
    r'|Corporation'
    r'|Enterprises'
    r'|Group'
    r'|Gruppe'
    r'|Holdings?'
    r'|Products'
    r'|Ventures?'
    r')$'
)

COMPANY_TYPE_RE = re.compile(
    r'^(?P<company>.*?)(?P<intl1> International)?,? (?P<type>'
    r'A\.?& S\. Klein GmbH \& Co\. KG'
    r'|A/S'
    r'|AB'
    r'|AG'
    r'|AS'
    r'|Ab'
    r'|BV'
    r'|B\.V\.'
    r'|B.V. Nederland'
    r'|C\.V\.'
    r'|Corp.'
    r'|GmbH \& C[oO]\. [oO]HG'
    r'|GmbH \& Co\. ?KG\.?'  # handle typo: Lukas Meindl GmbH & Co.KG
    r'|GmbH \& Co\. KGaA'
    r'|GmbH'
    r'|Inc\.?'
    r'|Incorporated'
    r'|International'
    r'|KG\.?'
    r'|Llc'
    r'|LLC'
    r'|LLP'
    r'|LP'
    r'|Limited'
    r'|Llp'
    r'|Ltd\.?'
    r'|Ltda\.?'
    r'|nv'
    r'|NV'
    r'|N\.V\.'
    r'|PBC'  # "Public Benefit Corporation"? Only on B Corp site
    r'|PLC'
    r'|P\.C\.'
    r'|P\.C\.'
    r'|Pty Ltd'
    r'|Pty'
    r'|S.\L\.'
    r'|SA'
    r'|SAPI DE CV SOFOM ENR'
    r'|SARL'
    r'|SE'
    r'|S\.A\.?'
    r'|S\.A\.U\.'
    r'|S\.R\.L\.'
    r'|S\.p\.A\.'
    r'|Sarl'
    r'|SpA'
    r'|b\.v\.'
    r'|gmbh'
    r'|inc'
    r'|plc'
    r')(?P<intl2> International)?$'
)

COMPANY_TYPE_CORRECTIONS = {
    'BV': 'B.V.',
    'Incorporated': 'Inc',
    'Llc': 'LLC',
    'Llp': 'LLP',
    'NV': 'N.V.',
    'S.A': 'S.A.',
    'Sarl': 'SARL',
    'SpA': 'S.p.A.',
    'b.v.': 'B.V.',
    'gmbh': 'GmbH',
    'inc': 'Inc',
    'nv': 'N.V.',
}

UNSTRIPPABLE_COMPANY_TYPES = {
    'LLP',
}

UNSTRIPPABLE_COMPANIES = {
    'Globe International',
    'Woolworths Limited',
}

def handle_matched_company(cd, category_map):
    """Take in a company dictionary from match_companies(), and
    output rows about that company to the database.

    This handles writing all rows except for the "campaign" table.

    (or will eventually)
    """
    # get brands
    brand_to_row, brand_map = get_brands_for_company(cd['keys'])

    # pick a canonical name for this company
    company_canonical, company_full = name_company(cd, set(brand_to_row))

    if company_canonical == company_full:
        log.info(company_canonical)
    else:
        log.info(u'{} ({})'.format(company_canonical, company_full))

    # merge company rows
    company_row = {}

    for campaign_id, company in sorted(cd['keys']):
         company_row.update(select_company(campaign_id, company))

    del company_row['campaign_id']  # should be at least one match

    company_row['company'] = company_canonical
    company_row['company_full'] = company_full

    # store company row
    output_row(company_row, 'company')

    # store company map and ratings
    for campaign_id, campaign_company in cd['keys']:
        # map
        map_row = dict(
            campaign_id=campaign_id, campaign_company=campaign_company,
            company=company_canonical)
        output_row(map_row, 'campaign_company_map')

        # ratings (may be more than one because of `scope`)
        for rating_row in select_company_ratings(
                campaign_id, campaign_company):
            rating_row['company'] = company_canonical
            output_row(rating_row, 'campaign_company_rating')

    # store company categories
    for cat_row in get_company_categories(
            company_canonical, cd['keys'], category_map):
        output_row(cat_row, 'company_category')

    # store brands
    brand_rows = sorted(brand_to_row.itervalues(), key=lambda r: r['brand'])
    for brand_row in brand_rows:
        brand_row['company'] = company_canonical
        output_row(brand_row, 'brand')

    # store brand map, rating
    campaign_to_company = dict(cd['keys'])
    brand_to_keys = defaultdict(set)

    for (campaign_id, campaign_brand), brand_canonical in brand_map.items():
        campaign_company = campaign_to_company[campaign_id]

        brand_to_keys[brand_canonical].add(
            (campaign_id, campaign_company, campaign_brand))

        # map
        map_row = dict(
            campaign_id=campaign_id, campaign_company=campaign_company,
            campaign_brand=campaign_brand, company=company_canonical,
            brand=brand_canonical)
        output_row(map_row, 'campaign_brand_map')

        # ratings
        for rating_row in select_brand_ratings(
                campaign_id, campaign_company, campaign_brand):
            rating_row['company'] = company_canonical
            rating_row['brand'] = brand_canonical
            output_row(rating_row, 'campaign_brand_rating')

    # store brand categories
    for brand_canonical, keys in sorted(brand_to_keys.items()):
        for cat_row in get_brand_categories(
                company_canonical, brand_canonical, keys, category_map):
            output_row(cat_row, 'brand_category')

def name_company(cd, brands=()):
    """Return short and full names for a company."""
    short_name = pick_short_name_from_variants(cd['display_names'])

    # if matching short name is shorter and matches a brand for one
    # of the companies, use that
    matching_short_name = pick_short_name_from_variants(
        cd['display_names'] | cd['matching_names'])

    if (matching_short_name != short_name and
        matching_short_name in brands):
        short_name = matching_short_name

    full_name = pick_full_name_from_variants(cd['display_names'])

    return short_name, full_name

def match_companies(companies_with_campaign_ids=None, aliases=None):
    """Match up similar company names.

    companies -- sequence of (campaign_id, company). Defaults to
                 get_companies_with_campaign_ids()
    aliases -- sequence of list of matching company names. Defaults to
               COMPANY_ALIASES

    Yields company dicts, which contain:
    keys -- set of tuples of (campaign_id, original_company_name)
    display_names -- set of names appropriate for display
    matching_names -- set of names for matching
    """
    if companies_with_campaign_ids is None:
        companies_with_campaign_ids = select_all_companies()
    if aliases is None:
        aliases = COMPANY_ALIASES

    to_merge = []

    # use variants for matching
    for variants in aliases:
        variants = set(variants)
        to_merge.append({'keys': set(), 'display_names': set(),
                         'matching_names': variants})

    # add in companies
    for campaign_id, company in companies_with_campaign_ids:
        if not company:  # skip blank company names
            continue
        keys = {(campaign_id, company)}
        display, matching = get_company_name_variants(company)

        to_merge.append({'keys': keys, 'display_names': set(display),
                         'matching_names': set(matching)})

    # match up company dicts and merge them together
    def keyfunc(cd):
        normed_variants = set()
        for name in cd['matching_names']:
            normed_variants.update(norm_with_variants(name))
        return normed_variants

    # and merge them together
    for cd_group in group_by_keys(to_merge, keyfunc):
        merged = merge_dicts(cd_group)

        if not merged['keys'] and merged['display_names']:
            # this can happen if hard-coded variants don't match anything
            log.warn('orphaned company dict: {}'.format(repr(merged)))
        else:
            yield merged

def pick_short_name_from_variants(variants):
    # shortest name, ties broken by, not all lower/upper, has accents
    return sorted(variants,
                  key=lambda v: (len(v), v == v.lower(), v == v.upper(),
                                 -len(v.encode('utf8'))))[0]

def pick_full_name_from_variants(variants):
    # longest name, ties broken by, not all lower/upper, has accents
    return sorted(variants,
                  key=lambda v: (-len(v), v == v.lower(), v == v.upper(),
                                 -len(v.encode('utf8'))))[0]

def get_company_name_variants(company):
    """Convert a company name to a set of variants for display,
    and a set of variants for matching only.
    """

    display_variants = set()  # usable as display name
    matching_variants = set()  # usable for matching

    def handle(company):
        company = simplify_whitespace(company)

        company = COMPANY_CORRECTIONS.get(company) or company

        m = COMPANY_TYPE_RE.match(company)
        if m:
            company = m.group('company')
            intl1 = m.group('intl1') or ''
            c_type = m.group('type')
            intl2 = m.group('intl2') or ''
            c_type = COMPANY_TYPE_CORRECTIONS.get(c_type) or c_type
            c_full = company + intl1 + ' ' + c_type + intl2
            display_variants.add(c_full)

            if (c_type in UNSTRIPPABLE_COMPANY_TYPES or
                c_full in UNSTRIPPABLE_COMPANIES):
                return

        display_variants.add(company)

        m = THE_X_CO_RE.match(company)
        if m:
            display_variants.add(m.group('company'))
            return

        m = X_CO_RE.match(company)
        if m:
            display_variants.add(m.group('company'))
            return

        m = THE_X_COMPANY_RE.match(company)
        if m:
            matching_variants.add(m.group('company'))
            return

        m = X_COMPANY_RE.match(company)
        if m:
            matching_variants.add(m.group('company'))
            return

    handle(company)

    matching_variants.update(display_variants)

    # handle slashes in company names
    for mv in list(matching_variants):
        if '/' in mv:
            matching_variants.update(
                filter(None, (part.strip() for part in mv.split('/'))))

    return display_variants, matching_variants

def correct_company_type(c_type):
    if c_type in COMPANY_TYPE_CORRECTIONS:
        return COMPANY_TYPE_CORRECTIONS[c_type]
    elif c_type.endswith('.') and c_type.count('.') == 1:
        return c_type[:-1]
    else:
        return c_type
<|endoftext|>


<commit_before>from raco import algebra
from pipelines import Pipelined
from raco.utility import emit

import logging
LOG = logging.getLogger(__name__)

"""
Apply rules to an expression
If successful, output will
only involve operators in the
target algebra
"""

def optimize_by_rules(expr, rules):
    for rule in rules:
        def recursiverule(expr):
            newexpr = rule(expr)
            LOG.debug("apply rule %s\n--- %s => %s", rule, expr, newexpr)
            newexpr.apply(recursiverule)
            return newexpr
        expr = recursiverule(expr)
    return expr

def optimize(expr, target, source, eliminate_common_subexpressions=False):
    """Fire the rule-based optimizer on an expression.  Fire all rules in the
    source algebra (logical) and the target algebra (physical)"""
    assert isinstance(expr, algebra.Operator)

    def opt(expr):
        so = optimize_by_rules(expr, source.opt_rules())
        newexpr = optimize_by_rules(so, target.opt_rules())
        if eliminate_common_subexpressions:
            newexpr = common_subexpression_elimination(newexpr)
        return newexpr
    return opt(expr)

def compile(expr):
    """Compile physical plan to linearized form for execution"""
    # TODO: Fix this
    algebra.reset()
    exprcode = []

    # TODO, actually use Parallel[Store...]]? Right now assumes it
    if (isinstance(expr, algebra.Sequence)
            or isinstance(expr, algebra.Parallel)):
        assert len(expr.children()) == 1, "expected single expression only"
        store_expr = expr.children()[0]
    else:
        store_expr = expr

    assert isinstance(store_expr, algebra.Store)
    assert len(store_expr.children()) == 1, "expected single expression only"

    lang = store_expr.language

    if isinstance(store_expr, Pipelined):
        body = lang.body(store_expr.compilePipeline())
    else:
        body = lang.body(expr)

    exprcode.append(emit(body))
    return emit(*exprcode)

def search(expr, tofind):
    """yield a sequence of subexpressions equal to tofind"""
    def match(node):
        if node == tofind:
            yield node
    for x in expr.preorder(match):
        yield x

def common_subexpression_elimination(expr):
    """remove redundant subexpressions"""
    def id(expr):
        yield expr
    eqclasses = []
    allfound = []
    for x in expr.preorder(id):
        if x not in allfound:
            found = [x for x in search(expr, x)]
            eqclasses.append((x, found))
            allfound += found

    def replace(expr):
        for witness, ec in eqclasses:
            if expr in ec:
                expr.apply(replace)
                # record the fact that we eliminated the redundant branches
                if witness != expr:
                    # witness.trace("replaces", expr)
                    for k, v in expr.gettrace():
                        witness.trace(k, v)
                return witness

    return expr.apply(replace)
<commit_msg>use disjunction version of isinstance<commit_after>from raco import algebra
from pipelines import Pipelined
from raco.utility import emit

import logging
LOG = logging.getLogger(__name__)

"""
Apply rules to an expression
If successful, output will
only involve operators in the
target algebra
"""

def optimize_by_rules(expr, rules):
    for rule in rules:
        def recursiverule(expr):
            newexpr = rule(expr)
            LOG.debug("apply rule %s\n--- %s => %s", rule, expr, newexpr)
            newexpr.apply(recursiverule)
            return newexpr
        expr = recursiverule(expr)
    return expr

def optimize(expr, target, source, eliminate_common_subexpressions=False):
    """Fire the rule-based optimizer on an expression.  Fire all rules in the
    source algebra (logical) and the target algebra (physical)"""
    assert isinstance(expr, algebra.Operator)

    def opt(expr):
        so = optimize_by_rules(expr, source.opt_rules())
        newexpr = optimize_by_rules(so, target.opt_rules())
        if eliminate_common_subexpressions:
            newexpr = common_subexpression_elimination(newexpr)
        return newexpr
    return opt(expr)

def compile(expr):
    """Compile physical plan to linearized form for execution"""
    # TODO: Fix this
    algebra.reset()
    exprcode = []

    # TODO, actually use Parallel[Store...]]? Right now assumes it
    if (isinstance(expr, algebra.Sequence, algebra.Parallel)):
        assert len(expr.children()) == 1, "expected single expression only"
        store_expr = expr.children()[0]
    else:
        store_expr = expr

    assert isinstance(store_expr, algebra.Store)
    assert len(store_expr.children()) == 1, "expected single expression only"

    lang = store_expr.language

    if isinstance(store_expr, Pipelined):
        body = lang.body(store_expr.compilePipeline())
    else:
        body = lang.body(expr)

    exprcode.append(emit(body))
    return emit(*exprcode)

def search(expr, tofind):
    """yield a sequence of subexpressions equal to tofind"""
    def match(node):
        if node == tofind:
            yield node
    for x in expr.preorder(match):
        yield x

def common_subexpression_elimination(expr):
    """remove redundant subexpressions"""
    def id(expr):
        yield expr
    eqclasses = []
    allfound = []
    for x in expr.preorder(id):
        if x not in allfound:
            found = [x for x in search(expr, x)]
            eqclasses.append((x, found))
            allfound += found

    def replace(expr):
        for witness, ec in eqclasses:
            if expr in ec:
                expr.apply(replace)
                # record the fact that we eliminated the redundant branches
                if witness != expr:
                    # witness.trace("replaces", expr)
                    for k, v in expr.gettrace():
                        witness.trace(k, v)
                return witness

    return expr.apply(replace)
<|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.

from __future__ import absolute_import, division, print_function

import pytest

from cryptography.hazmat.primitives.asymmetric import rsa

from ...utils import load_pkcs1_vectors, load_vectors_from_file

class TestRSA(object):
    @pytest.mark.parametrize(
        "pkcs1_example",
        load_vectors_from_file(
            "asymmetric/RSA/pkcs-1v2-1d2-vec/pss-vect.txt",
            load_pkcs1_vectors
        )
    )
    def test_load_pss_vect_example_keys(self, pkcs1_example):
        secret, public = pkcs1_example

        skey = rsa.RSAPrivateKey(**secret)
        pkey = rsa.RSAPublicKey(**public)
        pkey2 = skey.public_key()

        assert skey and pkey and pkey2

        assert skey.modulus
        assert skey.modulus == pkey.modulus
        assert skey.modulus == skey.n
        assert skey.public_exponent == pkey.public_exponent
        assert skey.public_exponent == skey.e
        assert skey.private_exponent == skey.d

        assert pkey.modulus
        assert pkey.modulus == pkey2.modulus
        assert pkey.modulus == pkey.n
        assert pkey.public_exponent == pkey2.public_exponent
        assert pkey.public_exponent == pkey.e

        assert skey.key_size
        assert skey.key_size == pkey.key_size
        assert skey.key_size == pkey2.key_size

        assert skey.p * skey.q == skey.modulus

    def test_invalid_argument_types(self):
        with pytest.raises(TypeError):
            rsa.RSAPrivateKey(None, None, None, None, None)

        with pytest.raises(TypeError):
            rsa.RSAPublicKey(None, None)

    def test_invalid_argument_values(self):
        # tiny example key
        rsa.RSAPrivateKey(3, 5, 14, 8, 15)

        # modulus too small
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 8, 2)

        # modulus wrong
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 8, 16)

        # p too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(16, 5, 14, 8, 15)

        # q too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 16, 14, 8, 15)

        # private exp too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 16, 8, 15)

        # public exp too low
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 2, 15)

        # public exp too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 16, 15)

        rsa.RSAPublicKey(8, 15)

        # modulus too small
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(8, 2)

        # public exp too low
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(2, 15)

        # public exp too high
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(16, 15)
<commit_msg>Use os.path.join to make paths<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.

from __future__ import absolute_import, division, print_function

import os

import pytest

from cryptography.hazmat.primitives.asymmetric import rsa

from ...utils import load_pkcs1_vectors, load_vectors_from_file

class TestRSA(object):
    @pytest.mark.parametrize(
        "pkcs1_example",
        load_vectors_from_file(
            os.path.join(
                "asymmetric", "RSA", "pkcs-1v2-1d2-vec", "pss-vect.txt"),
            load_pkcs1_vectors
        )
    )
    def test_load_pss_vect_example_keys(self, pkcs1_example):
        secret, public = pkcs1_example

        skey = rsa.RSAPrivateKey(**secret)
        pkey = rsa.RSAPublicKey(**public)
        pkey2 = skey.public_key()

        assert skey and pkey and pkey2

        assert skey.modulus
        assert skey.modulus == pkey.modulus
        assert skey.modulus == skey.n
        assert skey.public_exponent == pkey.public_exponent
        assert skey.public_exponent == skey.e
        assert skey.private_exponent == skey.d

        assert pkey.modulus
        assert pkey.modulus == pkey2.modulus
        assert pkey.modulus == pkey.n
        assert pkey.public_exponent == pkey2.public_exponent
        assert pkey.public_exponent == pkey.e

        assert skey.key_size
        assert skey.key_size == pkey.key_size
        assert skey.key_size == pkey2.key_size

        assert skey.p * skey.q == skey.modulus

    def test_invalid_argument_types(self):
        with pytest.raises(TypeError):
            rsa.RSAPrivateKey(None, None, None, None, None)

        with pytest.raises(TypeError):
            rsa.RSAPublicKey(None, None)

    def test_invalid_argument_values(self):
        # tiny example key
        rsa.RSAPrivateKey(3, 5, 14, 8, 15)

        # modulus too small
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 8, 2)

        # modulus wrong
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 8, 16)

        # p too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(16, 5, 14, 8, 15)

        # q too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 16, 14, 8, 15)

        # private exp too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 16, 8, 15)

        # public exp too low
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 2, 15)

        # public exp too high
        with pytest.raises(ValueError):
            rsa.RSAPrivateKey(3, 5, 14, 16, 15)

        rsa.RSAPublicKey(8, 15)

        # modulus too small
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(8, 2)

        # public exp too low
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(2, 15)

        # public exp too high
        with pytest.raises(ValueError):
            rsa.RSAPublicKey(16, 15)
<|endoftext|>


<commit_before>import re

import click
import gevent
import requests
import structlog
from eth_utils import to_hex
from gevent.event import AsyncResult
from pkg_resources import parse_version
from web3 import Web3

from raiden.utils import gas_reserve, pex
from raiden.utils.runnable import Runnable

CHECK_VERSION_INTERVAL = 3 * 60 * 60
CHECK_GAS_RESERVE_INTERVAL = 5 * 60
LATEST = 'https://api.github.com/repos/raiden-network/raiden/releases/latest'
RELEASE_PAGE = 'https://github.com/raiden-network/raiden/releases'
SECURITY_EXPRESSION = r'\[CRITICAL UPDATE.*?\]'

REMOVE_CALLBACK = object()
log = structlog.get_logger(__name__)  # pylint: disable=invalid-name

def check_version(current_version: str):
    """ Check periodically for a new release """
    app_version = parse_version(current_version)
    while True:
        try:
            content = requests.get(LATEST).json()
            # getting the latest release version
            latest_release = parse_version(content['tag_name'])
            security_message = re.search(SECURITY_EXPRESSION, content['body'])
            if security_message:
                click.secho(security_message.group(0), fg='red')
            # comparing it to the user's application
            if app_version < latest_release:
                msg = "You're running version {}. The latest version is {}".format(
                    app_version,
                    latest_release,
                )
                click.secho(msg, fg='red')
                click.secho("It's time to update! Releases: {}".format(RELEASE_PAGE), fg='red')
        except requests.exceptions.HTTPError as herr:
            click.secho('Error while checking for version', fg='red')
            print(herr)
        except ValueError as verr:
            click.secho('Error while checking the version', fg='red')
            print(verr)
        finally:
            # repeat the process once every 3h
            gevent.sleep(CHECK_VERSION_INTERVAL)

def check_gas_reserve(raiden):
    """ Check periodically for gas reserve in the account """
    while True:
        has_enough_balance, estimated_required_balance = gas_reserve.has_enough_gas_reserve(
            raiden,
            channels_to_open=0,
        )
        estimated_required_balance_eth = Web3.fromWei(estimated_required_balance, 'ether')

        if not has_enough_balance:
            log.info('Missing gas reserve', required_wei=estimated_required_balance)
            click.secho(
                (
                    'WARNING\n'
                    "Your account's balance is below the estimated gas reserve of "
                    f'{estimated_required_balance_eth} eth. This may lead to a loss of '
                    'of funds because your account will be unable to perform on-chain '
                    'transactions. Please add funds to your account as soon as possible.'
                ),
                fg='red',
            )

        gevent.sleep(CHECK_GAS_RESERVE_INTERVAL)

class AlarmTask(Runnable):
    """ Task to notify when a block is mined. """

    def __init__(self, chain):
        super().__init__()

        self.callbacks = list()
        self.chain = chain
        self.chain_id = None
        self.known_block_number = None
        self._stop_event = AsyncResult()

        # TODO: Start with a larger sleep_time and decrease it as the
        # probability of a new block increases.
        self.sleep_time = 0.5

    def start(self):
        log.debug('Alarm task started', node=pex(self.chain.node_address))
        super().start()

    def _run(self):  # pylint: disable=method-hidden
        try:
            self.loop_until_stop()
        finally:
            self.callbacks = list()

    def register_callback(self, callback):
        """ Register a new callback.

        Note:
            The callback will be executed in the AlarmTask context and for
            this reason it should not block, otherwise we can miss block
            changes.
        """
        if not callable(callback):
            raise ValueError('callback is not a callable')

        self.callbacks.append(callback)

    def remove_callback(self, callback):
        """Remove callback from the list of callbacks if it exists"""
        if callback in self.callbacks:
            self.callbacks.remove(callback)

    def loop_until_stop(self):
        # The AlarmTask must have completed its first_run() before starting
        # the background greenlet.
        #
        # This is required because the first run will synchronize the node with
        # the blockchain since the last run.
        assert self.chain_id, 'chain_id not set'
        assert self.known_block_number is not None, 'known_block_number not set'

        chain_id = self.chain_id

        sleep_time = self.sleep_time
        while self._stop_event.wait(sleep_time) is not True:
            latest_block = self.chain.get_block(block_identifier='latest')
            self._maybe_run_callbacks(latest_block)

            if chain_id != self.chain.network_id:
                raise RuntimeError(
                    'Changing the underlying blockchain while the Raiden node is running '
                    'is not supported.',
                )

    def first_run(self, known_block_number):
        """ Blocking call to update the local state, if necessary. """
        assert self.callbacks, 'callbacks not set'

        chain_id = self.chain.network_id
        latest_block = self.chain.get_block(block_identifier='latest')

        log.debug(
            'Alarm task first run',
            known_block_number=known_block_number,
            latest_block_number=latest_block['number'],
            latest_gas_limit=latest_block['gasLimit'],
            latest_block_hash=to_hex(latest_block['hash']),
        )

        self.known_block_number = known_block_number
        self.chain_id = chain_id
        self._maybe_run_callbacks(latest_block)

    def _maybe_run_callbacks(self, latest_block):
        """ Run the callbacks if there is at least one new block.

        The callbacks are executed only if there is a new block, otherwise the
        filters may try to poll for an inexisting block number and the Ethereum
        client can return an JSON-RPC error.
        """
        assert self.known_block_number is not None, 'known_block_number not set'

        latest_block_number = latest_block['number']
        missed_blocks = latest_block_number - self.known_block_number

        if missed_blocks < 0:
            log.critical(
                'Block number decreased',
                chain_id=self.chain_id,
                known_block_number=self.known_block_number,
                old_block_number=latest_block['number'],
                old_gas_limit=latest_block['gasLimit'],
                old_block_hash=to_hex(latest_block['hash']),
            )
        elif missed_blocks > 0:
            log_details = dict(
                known_block_number=self.known_block_number,
                latest_block_number=latest_block_number,
                latest_block_hash=to_hex(latest_block['hash']),
                latest_block_gas_limit=latest_block['gasLimit'],
            )
            if missed_blocks > 1:
                log_details['num_missed_blocks'] = missed_blocks - 1

            log.debug(
                'Received new block',
                **log_details,
            )

            remove = list()
            for callback in self.callbacks:
                result = callback(latest_block)
                if result is REMOVE_CALLBACK:
                    remove.append(callback)

            for callback in remove:
                self.callbacks.remove(callback)

            self.known_block_number = latest_block_number

    def stop(self):
        self._stop_event.set(True)
        log.debug('Alarm task stopped', node=pex(self.chain.node_address))
        return self.join()
<commit_msg>Do not crash check_version if github request returns unexpected data<commit_after>import re

import click
import gevent
import requests
import structlog
from eth_utils import to_hex
from gevent.event import AsyncResult
from pkg_resources import parse_version
from web3 import Web3

from raiden.utils import gas_reserve, pex
from raiden.utils.runnable import Runnable

CHECK_VERSION_INTERVAL = 3 * 60 * 60
CHECK_GAS_RESERVE_INTERVAL = 5 * 60
LATEST = 'https://api.github.com/repos/raiden-network/raiden/releases/latest'
RELEASE_PAGE = 'https://github.com/raiden-network/raiden/releases'
SECURITY_EXPRESSION = r'\[CRITICAL UPDATE.*?\]'

REMOVE_CALLBACK = object()
log = structlog.get_logger(__name__)  # pylint: disable=invalid-name

def _do_check_version(current_version: str):
    content = requests.get(LATEST).json()
    if 'tag_name' not in content:
        # probably API rate limit exceeded
        click.secho(
            'Error while contacting github for latest version. API rate limit exceeded?',
            fg='red',
        )
        return False
    # getting the latest release version
    latest_release = parse_version(content['tag_name'])
    security_message = re.search(SECURITY_EXPRESSION, content['body'])
    if security_message:
        click.secho(security_message.group(0), fg='red')
        # comparing it to the user's application
    if current_version < latest_release:
        msg = "You're running version {}. The latest version is {}".format(
            current_version,
            latest_release,
        )
        click.secho(msg, fg='red')
        click.secho("It's time to update! Releases: {}".format(RELEASE_PAGE), fg='red')
        return False
    return True

def check_version(current_version: str):
    """ Check periodically for a new release """
    app_version = parse_version(current_version)
    while True:
        try:
            _do_check_version(app_version)
        except requests.exceptions.HTTPError as herr:
            click.secho('Error while checking for version', fg='red')
            print(herr)
        except ValueError as verr:
            click.secho('Error while checking the version', fg='red')
            print(verr)
        finally:
            # repeat the process once every 3h
            gevent.sleep(CHECK_VERSION_INTERVAL)

def check_gas_reserve(raiden):
    """ Check periodically for gas reserve in the account """
    while True:
        has_enough_balance, estimated_required_balance = gas_reserve.has_enough_gas_reserve(
            raiden,
            channels_to_open=0,
        )
        estimated_required_balance_eth = Web3.fromWei(estimated_required_balance, 'ether')

        if not has_enough_balance:
            log.info('Missing gas reserve', required_wei=estimated_required_balance)
            click.secho(
                (
                    'WARNING\n'
                    "Your account's balance is below the estimated gas reserve of "
                    f'{estimated_required_balance_eth} eth. This may lead to a loss of '
                    'of funds because your account will be unable to perform on-chain '
                    'transactions. Please add funds to your account as soon as possible.'
                ),
                fg='red',
            )

        gevent.sleep(CHECK_GAS_RESERVE_INTERVAL)

class AlarmTask(Runnable):
    """ Task to notify when a block is mined. """

    def __init__(self, chain):
        super().__init__()

        self.callbacks = list()
        self.chain = chain
        self.chain_id = None
        self.known_block_number = None
        self._stop_event = AsyncResult()

        # TODO: Start with a larger sleep_time and decrease it as the
        # probability of a new block increases.
        self.sleep_time = 0.5

    def start(self):
        log.debug('Alarm task started', node=pex(self.chain.node_address))
        super().start()

    def _run(self):  # pylint: disable=method-hidden
        try:
            self.loop_until_stop()
        finally:
            self.callbacks = list()

    def register_callback(self, callback):
        """ Register a new callback.

        Note:
            The callback will be executed in the AlarmTask context and for
            this reason it should not block, otherwise we can miss block
            changes.
        """
        if not callable(callback):
            raise ValueError('callback is not a callable')

        self.callbacks.append(callback)

    def remove_callback(self, callback):
        """Remove callback from the list of callbacks if it exists"""
        if callback in self.callbacks:
            self.callbacks.remove(callback)

    def loop_until_stop(self):
        # The AlarmTask must have completed its first_run() before starting
        # the background greenlet.
        #
        # This is required because the first run will synchronize the node with
        # the blockchain since the last run.
        assert self.chain_id, 'chain_id not set'
        assert self.known_block_number is not None, 'known_block_number not set'

        chain_id = self.chain_id

        sleep_time = self.sleep_time
        while self._stop_event.wait(sleep_time) is not True:
            latest_block = self.chain.get_block(block_identifier='latest')
            self._maybe_run_callbacks(latest_block)

            if chain_id != self.chain.network_id:
                raise RuntimeError(
                    'Changing the underlying blockchain while the Raiden node is running '
                    'is not supported.',
                )

    def first_run(self, known_block_number):
        """ Blocking call to update the local state, if necessary. """
        assert self.callbacks, 'callbacks not set'

        chain_id = self.chain.network_id
        latest_block = self.chain.get_block(block_identifier='latest')

        log.debug(
            'Alarm task first run',
            known_block_number=known_block_number,
            latest_block_number=latest_block['number'],
            latest_gas_limit=latest_block['gasLimit'],
            latest_block_hash=to_hex(latest_block['hash']),
        )

        self.known_block_number = known_block_number
        self.chain_id = chain_id
        self._maybe_run_callbacks(latest_block)

    def _maybe_run_callbacks(self, latest_block):
        """ Run the callbacks if there is at least one new block.

        The callbacks are executed only if there is a new block, otherwise the
        filters may try to poll for an inexisting block number and the Ethereum
        client can return an JSON-RPC error.
        """
        assert self.known_block_number is not None, 'known_block_number not set'

        latest_block_number = latest_block['number']
        missed_blocks = latest_block_number - self.known_block_number

        if missed_blocks < 0:
            log.critical(
                'Block number decreased',
                chain_id=self.chain_id,
                known_block_number=self.known_block_number,
                old_block_number=latest_block['number'],
                old_gas_limit=latest_block['gasLimit'],
                old_block_hash=to_hex(latest_block['hash']),
            )
        elif missed_blocks > 0:
            log_details = dict(
                known_block_number=self.known_block_number,
                latest_block_number=latest_block_number,
                latest_block_hash=to_hex(latest_block['hash']),
                latest_block_gas_limit=latest_block['gasLimit'],
            )
            if missed_blocks > 1:
                log_details['num_missed_blocks'] = missed_blocks - 1

            log.debug(
                'Received new block',
                **log_details,
            )

            remove = list()
            for callback in self.callbacks:
                result = callback(latest_block)
                if result is REMOVE_CALLBACK:
                    remove.append(callback)

            for callback in remove:
                self.callbacks.remove(callback)

            self.known_block_number = latest_block_number

    def stop(self):
        self._stop_event.set(True)
        log.debug('Alarm task stopped', node=pex(self.chain.node_address))
        return self.join()
<|endoftext|>


<commit_before>def to_dict(obj):
    if hasattr(obj, '_to_dict'):
        return obj._to_dict()
    if isinstance(obj, list):
        return [to_dict(x) for x in obj]
    if isinstance(obj, int):
        return obj
    if isinstance(obj, str):
        return obj
    if isinstance(obj, float):
        return obj
    raise Exception('unhandled type \'{}\''.format(type(obj).__name__))

def from_dict(var, m=globals()):
    if isinstance(var, dict) and '__class__' in var:
        cls = m[var['__class__']]
        if cls:
            return cls._from_dict({k: from_dict(v, m) for k, v in var.items() if not k.startswith('_')})
    if isinstance(var, list):
        return [from_dict(x, m) for x in var]
    return var

def dict_fields(attributes, klass=None):
    def wrapper(klass):
        for a in attributes:
            assert not a.startswith('_')

        def _from_dict(cls, d):
            return cls(**d)

        def _to_dict(self):
            d = {'__class__': self.__class__.__name__}
            for a in attributes:
                v = getattr(self, a)
                if not v is None:
                    d[a] = to_dict(getattr(self, a))
            return d

        klass._from_dict = classmethod(_from_dict)
        klass._to_dict = _to_dict
        return klass

    return wrapper
<commit_msg>Add support for dictionaries<commit_after>def to_dict(obj):
    if hasattr(obj, '_to_dict'):
        return obj._to_dict()
    if isinstance(obj, tuple):
        return [to_dict(x) for x in obj]
    if isinstance(obj, list):
        return [to_dict(x) for x in obj]
    if isinstance(obj, dict):
        return {to_dict(k):to_dict(v) for k,v in obj.items()}
    if isinstance(obj, int):
        return obj
    if isinstance(obj, str):
        return obj
    if isinstance(obj, float):
        return obj
    raise Exception('unhandled type \'{}\''.format(type(obj).__name__))

def from_dict(var, m=globals()):
    if isinstance(var, dict) and '__class__' in var:
        cls = m[var['__class__']]
        if cls:
            return cls._from_dict({k: from_dict(v, m)
                                   for k, v in var.items()
                                   if not k.startswith('_')})
    if isinstance(var, dict):
        return {from_dict(k):from_dict(v) for k,v in var.items()}
    if isinstance(var, list):
        return [from_dict(x, m) for x in var]
    return var

def dict_fields(attributes):
    def wrapper(klass):
        for a in attributes:
            assert not a.startswith('_')

        def _from_dict(cls, d):
            return cls(**d)

        def _to_dict(self):
            d = {'__class__': self.__class__.__name__}
            for a in attributes:
                v = getattr(self, a)
                if not v is None:
                    d[a] = to_dict(getattr(self, a))
            return d

        klass._from_dict = classmethod(_from_dict)
        klass._to_dict = _to_dict
        return klass

    return wrapper
<|endoftext|>


<commit_before>from hashlib import md5
import copy
from utils import get_np_hashable
from store import HDFPickleStore, Store, default_store

__all__ = ['DataContext']

class DataContext(object):
    """
    All Ramp analyses require a DataContext.
    A DataContext represents the environment of the
    analysis. Most importantly this means for a given store
    and pandas index value, Ramp will consider the data immutable --
    it will not check the data again to see if it has changed.
    """
    def __init__(self, store, data=None, train_index=None, prep_index=None):
        """
        **Args**

        store: An instance of `store.Store` or a path. If a path
                Ramp will default to an `HDFPickleStore` at that path
                if PyTables is installed, a `PickleStore` otherwise.
        data: a pandas DataFrame. If all data has been precomputed this
                may not be required.
        train_index: a pandas Index specifying the data instances to be
                used in training. Stored results will be cached against this.
                If not provided, the entire index of `data` will be used.
        prep_index: a pandas Index specifying the data instances to be
                used in prepping ("x" values). Stored results will be cached against this.
                If not provided, the entire index of `data` will be used.
        """
        self.store = store if isinstance(store, Store) else default_store(store)
        self.data = data
        self.train_index = train_index if train_index is not None else self.data.index if self.data is not None else None
        self.prep_index = prep_index if prep_index is not None else self.data.index if self.data is not None else None

    def copy(self):
        return copy.copy(self)

    def create_key(self):
        return md5('%s--%s' % (get_np_hashable(self.train_index),
            get_np_hashable(self.prep_index))).hexdigest()

    def save_context(self, name, config=None):
        """
        Saves this context (specifically it's train and prep indices)
        to it's store with the given
        name, along with the config, if provided.
        """
        ctx = {'train_index':self.train_index,
                'prep_index':self.prep_index,
                'config':config}
        self.store.save('context__%s' % name, ctx)

    def load_context(self, name):
        """
        Loads a previously saved context with given name,
        assigning the stored training and prep indices
        and returning any stored config.
        """
        ctx = self.store.load('context__%s' % name)
        self.train_index = ctx['train_index']
        self.prep_index = ctx['prep_index']
        return ctx['config']
<commit_msg>Context defaults to MemoryStore.<commit_after>from hashlib import md5
import copy
from utils import get_np_hashable
from store import HDFPickleStore, MemoryStore, Store, default_store

__all__ = ['DataContext']

class DataContext(object):
    """
    All Ramp analyses require a DataContext.
    A DataContext represents the environment of the
    analysis. Most importantly this means for a given store
    and pandas index value, Ramp will consider the data immutable --
    it will not check the data again to see if it has changed.
    """
    def __init__(self, store=None, data=None, train_index=None, prep_index=None):
        """
        **Args**

        store: An instance of `store.Store` or a path. If a path
                Ramp will default to an `HDFPickleStore` at that path
                if PyTables is installed, a `PickleStore` otherwise.
                Defaults to MemoryStore.
        data: a pandas DataFrame. If all data has been precomputed this
                may not be required.
        train_index: a pandas Index specifying the data instances to be
                used in training. Stored results will be cached against this.
                If not provided, the entire index of `data` will be used.
        prep_index: a pandas Index specifying the data instances to be
                used in prepping ("x" values). Stored results will be cached against this.
                If not provided, the entire index of `data` will be used.
        """
        if store is None:
            self.store = MemoryStore()
        else:
            self.store = store if isinstance(store, Store) else default_store(store)
        self.data = data
        self.train_index = train_index if train_index is not None else self.data.index if self.data is not None else None
        self.prep_index = prep_index if prep_index is not None else self.data.index if self.data is not None else None

    def copy(self):
        return copy.copy(self)

    def create_key(self):
        return md5('%s--%s' % (get_np_hashable(self.train_index),
            get_np_hashable(self.prep_index))).hexdigest()

    def save_context(self, name, config=None):
        """
        Saves this context (specifically it's train and prep indices)
        to it's store with the given
        name, along with the config, if provided.
        """
        ctx = {'train_index':self.train_index,
                'prep_index':self.prep_index,
                'config':config}
        self.store.save('context__%s' % name, ctx)

    def load_context(self, name):
        """
        Loads a previously saved context with given name,
        assigning the stored training and prep indices
        and returning any stored config.
        """
        ctx = self.store.load('context__%s' % name)
        self.train_index = ctx['train_index']
        self.prep_index = ctx['prep_index']
        return ctx['config']
<|endoftext|>


<commit_before>import log, random
from utils import getRangeDict, chooseFromRange

class Choice(object):
    def __init__(self, restrictions):
        self.restrictions = restrictions
        self.settings = restrictions.settings
        self.log = log.get("Choice")

    # args are return from RandoServices.getPossiblePlacements
    # return itemLoc dict, or None if no possible choice
    def chooseItemLoc(self, itemLocDict, isProg):
        return None

    def getItemList(self, itemLocDict):
        return sorted([wrapper.item for wrapper in itemLocDict.keys()], key=lambda item: item['Type'])
    def getLocList(self, itemLocDict, item):
        return itemLocDict[item['Wrapper']]

# simple random choice (still with early morph check)
class ItemThenLocChoice(Choice):
    def __init__(self, restrictions):
        super(ItemThenLocChoice, self).__init__(restrictions)

    def chooseItemLoc(self, itemLocDict, isProg):
        itemList = self.getItemList(itemLocDict)
        item = self.chooseItem(itemList, isProg)
        if item is None:
            return None
        locList = self.getLocList(itemLocDict, item)
        loc = self.chooseLocation(locList, item, isProg)
        if loc is None:
            return None
        return {
            'Item': item,
            'Location': loc
        }

    def chooseItem(self, itemList, isProg):
        if len(itemList) == 0:
            return None
        if isProg:
            return self.chooseItemProg(itemList)
        else:
            return self.chooseItemRandom(itemList)

    def earlyMorphCheck(self, itemList):
        if not self.restrictions.isEarlyMorph():
            return None
        return next((item for item in itemList if item['Type'] == 'Morph'), None)

    def chooseItemProg(self, itemList):
        ret = self.earlyMorphCheck(itemList)
        if ret is None:
            ret = self.chooseItemRandom(itemList)
        return ret

    def chooseItemRandom(self, itemList):
        return random.choice(itemList)

    def chooseLocation(self, locList, item, isProg):
        if len(locList) == 0:
            return None
        if isProg:
            return self.chooseLocationProg(locList, item)
        else:
            return self.chooseLocationRandom(locList)

    def chooseLocationProg(self, locList, item):
        return self.chooseLocationRandom(locList)

    def chooseLocationRandom(self, locList):
        return random.choice(locList)

class ItemThenLocChoiceProgSpeed(ItemThenLocChoice):
    def __init__(self, restrictions, distanceProp):
        super(ItemThenLocChoiceProgSpeed, self).__init__(restrictions)
        self.distanceProp = distanceProp
        self.chooseItemFuncs = {
            'Random' : self.chooseItemRandom,
            'MinProgression' : self.chooseItemMinProgression,
            'MaxProgression' : self.chooseItemMaxProgression
        }
        self.chooseLocFuncs = {
            'Random' : self.chooseLocationRandom,
            'MinDiff' : self.chooseLocationMinDiff,
            'MaxDiff' : self.chooseLocationMaxDiff
        }

    def chooseItemLoc(self, itemLocDict, isProg, progressionItemLocs):
        self.progressionItemLocs = progressionItemLocs
        super(ItemThenLocChoiceProgSpeed, self).chooseItemLoc(itemLocDict, isProg)

    def determineParameters(self, progSpeed=None, progDiff=None):
        self.chooseLocRanges = getRangeDict(self.getChooseLocs(progDiff))
        self.chooseItemRanges = getRangeDict(self.getChooseItems(progSpeed))
        self.spreadProb = self.settings.getSpreadFactor(speed)

    def getChooseLocs(self, progDiff=None):
        if progDiff is None:
            progDiff = self.settings.progDiff
        return self.getChooseLocDict(self.settings.progDiff)

    def getChooseItems(self, progSpeed):
        if progSpeed is None:
            progSpeed = self.settings.progSpeed
        return self.getChooseItemDict(progSpeed)

    def getChooseLocDict(self, progDiff):
        if progDiff == 'normal':
            return {
                'Random' : 1,
                'MinDiff' : 0,
                'MaxDiff' : 0
            }
        elif progDiff == 'easier':
            return {
                'Random' : 2,
                'MinDiff' : 1,
                'MaxDiff' : 0
            }
        elif progDiff == 'harder':
            return {
                'Random' : 2,
                'MinDiff' : 0,
                'MaxDiff' : 1
            }

    def getChooseItemDict(self, progSpeed):
        if progSpeed == 'slowest':
            return {
                'MinProgression' : 1,
                'Random' : 2,
                'MaxProgression' : 0
            }
        elif progSpeed == 'slow':
            return {
                'MinProgression' : 25,
                'Random' : 75,
                'MaxProgression' : 0
            }
        elif progSpeed == 'medium':
            return {
                'MinProgression' : 0,
                'Random' : 1,
                'MaxProgression' : 0
            }
        elif progSpeed == 'fast':
            return {
                'MinProgression' : 0,
                'Random' : 75,
                'MaxProgression' : 25
            }
        elif progSpeed == 'fastest':
            return {
                'MinProgression' : 0,
                'Random' : 2,
                'MaxProgression' : 1
            }

    def getSpreadFactor(self, progSpeed):
        if progSpeed == 'slowest':
            return 0.9
        elif progSpeed == 'slow':
            return 0.7
        elif progSpeed == 'medium':
            return 0.4
        elif progSpeed == 'fast':
            return 0.1
        return 0

    def chooseItemProg(self, itemList):
        ret = self.earlyMorphCheck(itemList)
        if ret is None:
            ret = self.getChooseFunc(self.chooseItemRanges, self.chooseItemFuncs)(itemList)
        return ret

    def chooseLocationProg(self, locs, item):
        # FIXME find a proper way to add this locations restrictions function. idea: do the same as item pool restriction in item pool container
        # locs = availableLocations
        # if self.isEarlyGame():
        #     # cheat a little bit if non-standard start: place early
        #     # progression away from crateria/blue brin if possible
        #     startAp = getAccessPoint(self.settings.startAP)
        #     if startAp.GraphArea != "Crateria":
        #         locs = [loc for loc in availableLocations if loc['GraphArea'] != 'Crateria']
        #         if len(locs) == 0:
        #             locs = availableLocations
        # isProg = self.isProgItem(item)
        if random.random() < self.spreadProb:
            locs = self.getLocsSpreadProgression(locs)
        random.shuffle(locs)
        self.log.debug("chooseLocation isProg: {}".format(isProg))
        if isProg == True:
            return self.getChooseFunc(self.chooseLocRanges, self.chooseLocFuncs)(locs, item)
        else:
            # choose randomly if non-progression
            return self.chooseLocationRandom(locs, item)

    # get choose function from a weighted dict
    def getChooseFunc(self, rangeDict, funcDict):
        v = chooseFromRange(rangeDict)

        return funcDict[v]

    def chooseItemMinProgression(self, items):
        minNewLocs = 1000
        ret = None

        for item in items:
            newLocs = len(self.currentLocations(item))
            if newLocs < minNewLocs:
                minNewLocs = newLocs
                ret = item
        return ret

    def chooseItemMaxProgression(self, items):
        maxNewLocs = 0
        ret = None

        for item in items:
            newLocs = len(self.currentLocations(item))
            if newLocs > maxNewLocs:
                maxNewLocs = newLocs
                ret = item
        return ret

    def getLocDiff(self, loc):
        # avail difficulty already stored by graph algorithm
        return loc['difficulty']

    def fillLocsDiff(self, locs):
        for loc in locs:
            if 'PostAvailable' in loc:
                loc['difficulty'] = self.smbm.wand(self.getLocDiff(loc), self.smbm.eval(loc['PostAvailable']))

    def chooseLocationMaxDiff(self, availableLocations, item):
        self.log.debug("MAX")
        self.fillLocsDiff(availableLocations)
        self.log.debug("chooseLocationMaxDiff: {}".format([(l['Name'], l['difficulty']) for l in availableLocations]))
        return max(availableLocations, key=lambda loc:loc['difficulty'].difficulty)

    def chooseLocationMinDiff(self, availableLocations, item):
        self.log.debug("MIN")
        self.fillLocsDiff(availableLocations)
        self.log.debug("chooseLocationMinDiff: {}".format([(l['Name'], l['difficulty']) for l in availableLocations]))
        return min(availableLocations, key=lambda loc:loc['difficulty'].difficulty)

    def areaDistance(self, loc, otherLocs):
        areas = [l[self.distanceProp] for l in otherLocs]
        cnt = areas.count(loc[self.distanceProp])
        d = None
        if cnt == 0:
            d = 2
        else:
            d = 1.0/cnt
        return d

    def getLocsSpreadProgression(self, availableLocations):
        progLocs = [il['Location'] for il in self.progressionItemLocs if (self.restriction.split == 'Full' or self.restrictions.split == il['Item']['Class']) and il['Item']['Category'] != "Energy"]
        distances = [self.areaDistance(loc, progLocs) for loc in availableLocations]
        maxDist = max(distances)
        indices = [index for index, d in enumerate(distances) if d == maxDist]
        locs = [availableLocations[i] for i in indices]

        return locs
<commit_msg>algo: improved (or at least changed) spread prog<commit_after>import log, random
from utils import getRangeDict, chooseFromRange

class Choice(object):
    def __init__(self, restrictions):
        self.restrictions = restrictions
        self.settings = restrictions.settings
        self.log = log.get("Choice")

    # args are return from RandoServices.getPossiblePlacements
    # return itemLoc dict, or None if no possible choice
    def chooseItemLoc(self, itemLocDict, isProg):
        return None

    def getItemList(self, itemLocDict):
        return sorted([wrapper.item for wrapper in itemLocDict.keys()], key=lambda item: item['Type'])
    def getLocList(self, itemLocDict, item):
        return itemLocDict[item['Wrapper']]

# simple random choice (still with early morph check)
class ItemThenLocChoice(Choice):
    def __init__(self, restrictions):
        super(ItemThenLocChoice, self).__init__(restrictions)

    def chooseItemLoc(self, itemLocDict, isProg):
        itemList = self.getItemList(itemLocDict)
        item = self.chooseItem(itemList, isProg)
        if item is None:
            return None
        locList = self.getLocList(itemLocDict, item)
        loc = self.chooseLocation(locList, item, isProg)
        if loc is None:
            return None
        return {
            'Item': item,
            'Location': loc
        }

    def chooseItem(self, itemList, isProg):
        if len(itemList) == 0:
            return None
        if isProg:
            return self.chooseItemProg(itemList)
        else:
            return self.chooseItemRandom(itemList)

    def earlyMorphCheck(self, itemList):
        if not self.restrictions.isEarlyMorph():
            return None
        return next((item for item in itemList if item['Type'] == 'Morph'), None)

    def chooseItemProg(self, itemList):
        ret = self.earlyMorphCheck(itemList)
        if ret is None:
            ret = self.chooseItemRandom(itemList)
        return ret

    def chooseItemRandom(self, itemList):
        return random.choice(itemList)

    def chooseLocation(self, locList, item, isProg):
        if len(locList) == 0:
            return None
        if isProg:
            return self.chooseLocationProg(locList, item)
        else:
            return self.chooseLocationRandom(locList)

    def chooseLocationProg(self, locList, item):
        return self.chooseLocationRandom(locList)

    def chooseLocationRandom(self, locList):
        return random.choice(locList)

class ItemThenLocChoiceProgSpeed(ItemThenLocChoice):
    def __init__(self, restrictions, distanceProp):
        super(ItemThenLocChoiceProgSpeed, self).__init__(restrictions)
        self.distanceProp = distanceProp
        self.chooseItemFuncs = {
            'Random' : self.chooseItemRandom,
            'MinProgression' : self.chooseItemMinProgression,
            'MaxProgression' : self.chooseItemMaxProgression
        }
        self.chooseLocFuncs = {
            'Random' : self.chooseLocationRandom,
            'MinDiff' : self.chooseLocationMinDiff,
            'MaxDiff' : self.chooseLocationMaxDiff
        }

    def chooseItemLoc(self, itemLocDict, isProg, progressionItemLocs):
        self.progressionItemLocs = progressionItemLocs
        super(ItemThenLocChoiceProgSpeed, self).chooseItemLoc(itemLocDict, isProg)

    def determineParameters(self, progSpeed=None, progDiff=None):
        self.chooseLocRanges = getRangeDict(self.getChooseLocs(progDiff))
        self.chooseItemRanges = getRangeDict(self.getChooseItems(progSpeed))
        self.spreadProb = self.settings.getSpreadFactor(speed)

    def getChooseLocs(self, progDiff=None):
        if progDiff is None:
            progDiff = self.settings.progDiff
        return self.getChooseLocDict(self.settings.progDiff)

    def getChooseItems(self, progSpeed):
        if progSpeed is None:
            progSpeed = self.settings.progSpeed
        return self.getChooseItemDict(progSpeed)

    def getChooseLocDict(self, progDiff):
        if progDiff == 'normal':
            return {
                'Random' : 1,
                'MinDiff' : 0,
                'MaxDiff' : 0
            }
        elif progDiff == 'easier':
            return {
                'Random' : 2,
                'MinDiff' : 1,
                'MaxDiff' : 0
            }
        elif progDiff == 'harder':
            return {
                'Random' : 2,
                'MinDiff' : 0,
                'MaxDiff' : 1
            }

    def getChooseItemDict(self, progSpeed):
        if progSpeed == 'slowest':
            return {
                'MinProgression' : 1,
                'Random' : 2,
                'MaxProgression' : 0
            }
        elif progSpeed == 'slow':
            return {
                'MinProgression' : 25,
                'Random' : 75,
                'MaxProgression' : 0
            }
        elif progSpeed == 'medium':
            return {
                'MinProgression' : 0,
                'Random' : 1,
                'MaxProgression' : 0
            }
        elif progSpeed == 'fast':
            return {
                'MinProgression' : 0,
                'Random' : 75,
                'MaxProgression' : 25
            }
        elif progSpeed == 'fastest':
            return {
                'MinProgression' : 0,
                'Random' : 2,
                'MaxProgression' : 1
            }

    def getSpreadFactor(self, progSpeed):
        if progSpeed == 'slowest':
            return 0.9
        elif progSpeed == 'slow':
            return 0.7
        elif progSpeed == 'medium':
            return 0.4
        elif progSpeed == 'fast':
            return 0.1
        return 0

    def chooseItemProg(self, itemList):
        ret = self.earlyMorphCheck(itemList)
        if ret is None:
            ret = self.getChooseFunc(self.chooseItemRanges, self.chooseItemFuncs)(itemList)
        return ret

    def chooseLocationProg(self, locs, item):
        # FIXME find a proper way to add this locations restrictions function. idea: do the same as item pool restriction in item pool container
        # locs = availableLocations
        # if self.isEarlyGame():
        #     # cheat a little bit if non-standard start: place early
        #     # progression away from crateria/blue brin if possible
        #     startAp = getAccessPoint(self.settings.startAP)
        #     if startAp.GraphArea != "Crateria":
        #         locs = [loc for loc in availableLocations if loc['GraphArea'] != 'Crateria']
        #         if len(locs) == 0:
        #             locs = availableLocations
        # isProg = self.isProgItem(item)
        locs = self.getLocsSpreadProgression(locs)
        random.shuffle(locs)
        return self.getChooseFunc(self.chooseLocRanges, self.chooseLocFuncs)(locs, item)

    # get choose function from a weighted dict
    def getChooseFunc(self, rangeDict, funcDict):
        v = chooseFromRange(rangeDict)

        return funcDict[v]

    def chooseItemMinProgression(self, items):
        minNewLocs = 1000
        ret = None

        for item in items:
            newLocs = len(self.currentLocations(item))
            if newLocs < minNewLocs:
                minNewLocs = newLocs
                ret = item
        return ret

    def chooseItemMaxProgression(self, items):
        maxNewLocs = 0
        ret = None

        for item in items:
            newLocs = len(self.currentLocations(item))
            if newLocs > maxNewLocs:
                maxNewLocs = newLocs
                ret = item
        return ret

    def getLocDiff(self, loc):
        # avail difficulty already stored by graph algorithm
        return loc['difficulty']

    def fillLocsDiff(self, locs):
        for loc in locs:
            if 'PostAvailable' in loc:
                loc['difficulty'] = self.smbm.wand(self.getLocDiff(loc), self.smbm.eval(loc['PostAvailable']))

    def chooseLocationMaxDiff(self, availableLocations, item):
        self.log.debug("MAX")
        self.fillLocsDiff(availableLocations)
        self.log.debug("chooseLocationMaxDiff: {}".format([(l['Name'], l['difficulty']) for l in availableLocations]))
        return max(availableLocations, key=lambda loc:loc['difficulty'].difficulty)

    def chooseLocationMinDiff(self, availableLocations, item):
        self.log.debug("MIN")
        self.fillLocsDiff(availableLocations)
        self.log.debug("chooseLocationMinDiff: {}".format([(l['Name'], l['difficulty']) for l in availableLocations]))
        return min(availableLocations, key=lambda loc:loc['difficulty'].difficulty)

    def areaDistance(self, loc, otherLocs):
        areas = [l[self.distanceProp] for l in otherLocs]
        cnt = areas.count(loc[self.distanceProp])
        d = None
        if cnt == 0:
            d = 2
        else:
            d = 1.0/cnt
        return d

    def getLocsSpreadProgression(self, availableLocations):
        cond = lambda item: ((self.restriction.split == 'Full' and item['Class'] == 'Major') or self.restriction.split == item['Class']) and item['Category'] != "Energy"
        progLocs = [il['Location'] for il in self.progressionItemLocs if cond(il['Item'])]
        distances = [self.areaDistance(loc, progLocs) for loc in availableLocations]
        maxDist = max(distances)
        locs = []
        for i in range(len(availableLocations)):
            loc = availableLocations[i]
            d = distances[i]
            if d == maxDist or random.random() >= self.spreadProb:
                locs.append(loc)
        return locs
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Import python libs
from __future__ import absolute_import
import re

# Import Salt Testing libs
from salttesting.helpers import ensure_in_syspath
ensure_in_syspath('../../')

# Import salt libs
import integration

# Import 3rd-party libs
import salt.ext.six as six

class SysModuleTest(integration.ModuleCase):
    '''
    Validate the sys module
    '''
    def test_list_functions(self):
        '''
        sys.list_functions
        '''
        # Get all functions
        funcs = self.run_function('sys.list_functions')
        self.assertIn('hosts.list_hosts', funcs)
        self.assertIn('pkg.install', funcs)

        # Just sysctl
        funcs = self.run_function('sys.list_functions', ('sysctl',))
        self.assertNotIn('sys.doc', funcs)
        self.assertIn('sysctl.get', funcs)

        # Just sys
        funcs = self.run_function('sys.list_functions', ('sys.',))
        self.assertNotIn('sysctl.get', funcs)
        self.assertIn('sys.doc', funcs)

        # Staring with sys
        funcs = self.run_function('sys.list_functions', ('sys',))
        self.assertNotIn('sysctl.get', funcs)
        self.assertIn('sys.doc', funcs)

    def test_list_modules(self):
        '''
        sys.list_moduels
        '''
        mods = self.run_function('sys.list_modules')
        self.assertTrue('hosts' in mods)
        self.assertTrue('pkg' in mods)

    def test_valid_docs(self):
        '''
        Make sure no functions are exposed that don't have valid docstrings
        '''
        docs = self.run_function('sys.doc')
        nodoc = set()
        noexample = set()
        allow_failure = (
                'cp.recv',
                'lxc.run_cmd',
                'ipset.long_range',
                'pkg.expand_repo_def',
                'runtests_decorators.depends',
                'runtests_decorators.depends_will_fallback',
                'runtests_decorators.missing_depends',
                'runtests_decorators.missing_depends_will_fallback',
                'swift.head',
                'glance.warn_until',
                'yumpkg.expand_repo_def',
                'yumpkg5.expand_repo_def',
                'container_resource.run',
                'nspawn.stop',
                'nspawn.restart',
                'lowpkg.bin_pkg_info',
                'state.apply',
        )

        for fun in docs:
            if fun.startswith('runtests_helpers'):
                continue
            if fun in allow_failure:
                continue
            if not isinstance(docs[fun], six.string_types):
                nodoc.add(fun)
            elif not re.search(r'([E|e]xample(?:s)?)+(?:.*)::?', docs[fun]):
                noexample.add(fun)

        if not nodoc and not noexample:
            return

        raise AssertionError(
            'There are some functions which do not have a docstring or do not '
            'have an example:\nNo docstring:\n{0}\nNo example:\n{1}\n'.format(
                '\n'.join(['  - {0}'.format(f) for f in sorted(nodoc)]),
                '\n'.join(['  - {0}'.format(f) for f in sorted(noexample)]),
            )
        )

if __name__ == '__main__':
    from integration import run_tests
    run_tests(SysModuleTest)
<commit_msg>Exempt pip.iteritems from test_valid_docs test<commit_after># -*- coding: utf-8 -*-

# Import python libs
from __future__ import absolute_import
import re

# Import Salt Testing libs
from salttesting.helpers import ensure_in_syspath
ensure_in_syspath('../../')

# Import salt libs
import integration

# Import 3rd-party libs
import salt.ext.six as six

class SysModuleTest(integration.ModuleCase):
    '''
    Validate the sys module
    '''
    def test_list_functions(self):
        '''
        sys.list_functions
        '''
        # Get all functions
        funcs = self.run_function('sys.list_functions')
        self.assertIn('hosts.list_hosts', funcs)
        self.assertIn('pkg.install', funcs)

        # Just sysctl
        funcs = self.run_function('sys.list_functions', ('sysctl',))
        self.assertNotIn('sys.doc', funcs)
        self.assertIn('sysctl.get', funcs)

        # Just sys
        funcs = self.run_function('sys.list_functions', ('sys.',))
        self.assertNotIn('sysctl.get', funcs)
        self.assertIn('sys.doc', funcs)

        # Staring with sys
        funcs = self.run_function('sys.list_functions', ('sys',))
        self.assertNotIn('sysctl.get', funcs)
        self.assertIn('sys.doc', funcs)

    def test_list_modules(self):
        '''
        sys.list_moduels
        '''
        mods = self.run_function('sys.list_modules')
        self.assertTrue('hosts' in mods)
        self.assertTrue('pkg' in mods)

    def test_valid_docs(self):
        '''
        Make sure no functions are exposed that don't have valid docstrings
        '''
        docs = self.run_function('sys.doc')
        nodoc = set()
        noexample = set()
        allow_failure = (
                'cp.recv',
                'lxc.run_cmd',
                'ipset.long_range',
                'pkg.expand_repo_def',
                'runtests_decorators.depends',
                'runtests_decorators.depends_will_fallback',
                'runtests_decorators.missing_depends',
                'runtests_decorators.missing_depends_will_fallback',
                'swift.head',
                'glance.warn_until',
                'yumpkg.expand_repo_def',
                'yumpkg5.expand_repo_def',
                'container_resource.run',
                'nspawn.stop',
                'nspawn.restart',
                'lowpkg.bin_pkg_info',
                'state.apply',
                'pip.iteritems',
        )

        for fun in docs:
            if fun.startswith('runtests_helpers'):
                continue
            if fun in allow_failure:
                continue
            if not isinstance(docs[fun], six.string_types):
                nodoc.add(fun)
            elif not re.search(r'([E|e]xample(?:s)?)+(?:.*)::?', docs[fun]):
                noexample.add(fun)

        if not nodoc and not noexample:
            return

        raise AssertionError(
            'There are some functions which do not have a docstring or do not '
            'have an example:\nNo docstring:\n{0}\nNo example:\n{1}\n'.format(
                '\n'.join(['  - {0}'.format(f) for f in sorted(nodoc)]),
                '\n'.join(['  - {0}'.format(f) for f in sorted(noexample)]),
            )
        )

if __name__ == '__main__':
    from integration import run_tests
    run_tests(SysModuleTest)
<|endoftext|>