<commit_before>#!/usr/bin/python
# This file is part of Ansible
#
# Ansible 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.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

# This is a DOCUMENTATION stub specific to this module, it extends
# a documentation fragment located in ansible.utils.module_docs_fragments
DOCUMENTATION='''
module: rax_clb_ssl
short_description: Manage SSL termination for a Rackspace Cloud Load Balancer.
description:
- Set up, reconfigure, or remove SSL termination for an existing load balancer.
version_added: "1.9"
options:
  balancer_name:
    description:
    - Name or ID of the load balancer on which to manage SSL termination.
    required: true
  state:
    description:
    - If set to "present", SSL termination will be added to this load balancer.
    - If "absent", SSL termination will be removed instead.
    choices:
      - present
      - absent
    default: present
  enabled:
    description:
    - If set to "false", temporarily disable SSL termination without discarding
    - existing credentials.
    default: true
  private_key:
    description:
    - The private SSL key as a string in PEM format.
  certificate:
    description:
    - The public SSL certificates as a string in PEM format.
  intermediate_certificate:
    description:
    - One or more intermediate certificate authorities as a string in PEM
    - format, concatenated into a single string.
  secure_port:
    description:
    - The port to listen for secure traffic.
    default: 443
  secure_traffic_only:
    description:
    - If "true", the load balancer will *only* accept secure traffic.
    default: false
  https_redirect:
    description:
    - If "true", the load balancer will redirect HTTP traffic to HTTPS.
    - Requires "secure_traffic_only" to be true. Incurs an implicit wait if SSL
    - termination is also applied or removed.
  wait:
    description:
    - Wait for the balancer to be in state "running" before turning.
    default: false
  wait_timeout:
    description:
    - How long before "wait" gives up, in seconds.
    default: 300
author: Ash Wilson
extends_documentation_fragment: rackspace
'''

EXAMPLES = '''
- name: Enable SSL termination on a load balancer
  rax_clb_ssl:
    balancer_name: the_loadbalancer
    state: present
    private_key: "{{ lookup('file', 'credentials/server.key' ) }}"
    certificate: "{{ lookup('file', 'credentials/server.crt' ) }}"
    intermediate_certificate: "{{ lookup('file', 'credentials/trust-chain.crt') }}"
    secure_traffic_only: true
    wait: true

- name: Disable SSL termination
  rax_clb_ssl:
    balancer_name: "{{ registered_lb.balancer.id }}"
    state: absent
    wait: true
'''

from ansible.module_utils.basic import *
from ansible.module_utils.rax import *

try:
    import pyrax
    HAS_PYRAX = True
except ImportError:
    HAS_PYRAX = False

def cloud_load_balancer_ssl(module, balancer_name, state, enabled, private_key,
                            certificate, intermediate_certificate, secure_port,
                            secure_traffic_only, https_redirect,
                            wait, wait_timeout):
    # Validate arguments.

    if not balancer_name:
        module.fail_json(msg='balancer_name is required.')

    if state == 'present':
        if not private_key:
            module.fail_json(msg="private_key must be provided.")
        else:
            private_key = private_key.strip()

        if not certificate:
            module.fail_json(msg="certificate must be provided.")
        else:
            certificate = certificate.strip()

    if state not in ('present', 'absent'):
        module.fail_json(msg="State must be either 'present' or 'absent'.")

    attempts = wait_timeout / 5

    # Locate the load balancer.

    clb = pyrax.cloud_loadbalancers
    if not clb:
        module.fail_json(msg='Failed to instantiate client. This '
                             'typically indicates an invalid region or an '
                             'incorrectly capitalized region name.')

    balancers = []
    for balancer in clb.list():
        if balancer_name == balancer.name or balancer_name == str(balancer.id):
            balancers.append(balancer)

    if not balancers:
        module.fail_json(msg='No load balancers matched your criteria. '
                             'Use rax_clb to create the balancer first.')

    if len(balancers) > 1:
        module.fail_json(msg="%d load balancers were matched your criteria. Try"
                             "using the balancer's id instead." % len(balancers))

    balancer = balancers[0]
    existing_ssl = balancer.get_ssl_termination()

    changed = False

    if state == 'present':
        # Apply or reconfigure SSL termination on the load balancer.
        ssl_attrs = dict(
            securePort=secure_port,
            privatekey=private_key,
            certificate=certificate,
            intermediateCertificate=intermediate_certificate,
            enabled=enabled,
            secureTrafficOnly=secure_traffic_only
        )

        needs_change = False

        if existing_ssl:
            for ssl_attr, value in ssl_attrs.iteritems():
                if ssl_attr == 'privatekey':
                    # The private key is not included in get_ssl_termination's
                    # output (as it shouldn't be). Also, if you're changing the
                    # private key, you'll also be changing the certificate,
                    # so we don't lose anything by not checking it.
                    continue

                if value is not None and existing_ssl.get(ssl_attr) != value:
                    # module.fail_json(msg='Unnecessary change', attr=ssl_attr, value=value, existing=existing_ssl.get(ssl_attr))
                    needs_change = True
        else:
            needs_change = True

        if needs_change:
            balancer.add_ssl_termination(**ssl_attrs)
            changed = True
    elif state == 'absent':
        # Remove SSL termination if it's already configured.
        if existing_ssl:
            balancer.delete_ssl_termination()
            changed = True

    if https_redirect is not None and balancer.httpsRedirect != https_redirect:
        if changed:
            # This wait is unavoidable because load balancers are immutable
            # while the SSL termination changes above are being applied.
            pyrax.utils.wait_for_build(balancer, interval=5, attempts=attempts)

        balancer.update(httpsRedirect=https_redirect)
        changed = True

    if changed and wait:
        pyrax.utils.wait_for_build(balancer, interval=5, attempts=attempts)

    balancer.get()
    new_ssl_termination = balancer.get_ssl_termination()

    # Intentionally omit the private key from the module output, so you don't
    # accidentally echo it with `ansible-playbook -v` or `debug`, and the
    # certificate, which is just long. Convert other attributes to snake_case
    # and include https_redirect at the top-level.
    if new_ssl_termination:
        new_ssl = dict(
            enabled=new_ssl_termination['enabled'],
            secure_port=new_ssl_termination['securePort'],
            secure_traffic_only=new_ssl_termination['secureTrafficOnly']
        )
    else:
        new_ssl = None

    result = dict(
        changed=changed,
        https_redirect=balancer.httpsRedirect,
        ssl_termination=new_ssl
    )
    success = True

    if balancer.status == 'ERROR':
        result['msg'] = '%s failed to build' % balancer.id
        success = False
    elif wait and balancer.status not in ('ACTIVE', 'ERROR'):
        result['msg'] = 'Timeout waiting on %s' % balancer.id
        success = False

    if success:
        module.exit_json(**result)
    else:
        module.fail_json(**result)

def main():
    argument_spec = rax_argument_spec()
    argument_spec.update(dict(
        balancer_name=dict(type='str'),
        state=dict(default='present', choices=['present', 'absent']),
        enabled=dict(type='bool', default=True),
        private_key=dict(),
        certificate=dict(),
        intermediate_certificate=dict(),
        secure_port=dict(type='int', default=443),
        secure_traffic_only=dict(type='bool', default=False),
        https_redirect=dict(type='bool'),
        wait=dict(type='bool', default=False),
        wait_timeout=dict(type='int', default=300)
    ))

    module = AnsibleModule(
        argument_spec=argument_spec,
        required_together=rax_required_together(),
    )

    if not HAS_PYRAX:
        module.fail_json(msg='pyrax is required for this module.')

    balancer_name = module.params.get('balancer_name')
    state = module.params.get('state')
    enabled = module.boolean(module.params.get('enabled'))
    private_key = module.params.get('private_key')
    certificate = module.params.get('certificate')
    intermediate_certificate = module.params.get('intermediate_certificate')
    secure_port = module.params.get('secure_port')
    secure_traffic_only = module.boolean(module.params.get('secure_traffic_only'))
    https_redirect = module.boolean(module.params.get('https_redirect'))
    wait = module.boolean(module.params.get('wait'))
    wait_timeout = module.params.get('wait_timeout')

    setup_rax_module(module, pyrax)

    cloud_load_balancer_ssl(
        module, balancer_name, state, enabled, private_key, certificate,
        intermediate_certificate, secure_port, secure_traffic_only,
        https_redirect, wait, wait_timeout
    )

main()
<commit_msg>Move ansible imports to the module's bottom.<commit_after>#!/usr/bin/python
# This file is part of Ansible
#
# Ansible 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.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

# This is a DOCUMENTATION stub specific to this module, it extends
# a documentation fragment located in ansible.utils.module_docs_fragments
DOCUMENTATION='''
module: rax_clb_ssl
short_description: Manage SSL termination for a Rackspace Cloud Load Balancer.
description:
- Set up, reconfigure, or remove SSL termination for an existing load balancer.
version_added: "1.9"
options:
  balancer_name:
    description:
    - Name or ID of the load balancer on which to manage SSL termination.
    required: true
  state:
    description:
    - If set to "present", SSL termination will be added to this load balancer.
    - If "absent", SSL termination will be removed instead.
    choices:
      - present
      - absent
    default: present
  enabled:
    description:
    - If set to "false", temporarily disable SSL termination without discarding
    - existing credentials.
    default: true
  private_key:
    description:
    - The private SSL key as a string in PEM format.
  certificate:
    description:
    - The public SSL certificates as a string in PEM format.
  intermediate_certificate:
    description:
    - One or more intermediate certificate authorities as a string in PEM
    - format, concatenated into a single string.
  secure_port:
    description:
    - The port to listen for secure traffic.
    default: 443
  secure_traffic_only:
    description:
    - If "true", the load balancer will *only* accept secure traffic.
    default: false
  https_redirect:
    description:
    - If "true", the load balancer will redirect HTTP traffic to HTTPS.
    - Requires "secure_traffic_only" to be true. Incurs an implicit wait if SSL
    - termination is also applied or removed.
  wait:
    description:
    - Wait for the balancer to be in state "running" before turning.
    default: false
  wait_timeout:
    description:
    - How long before "wait" gives up, in seconds.
    default: 300
author: Ash Wilson
extends_documentation_fragment: rackspace
'''

EXAMPLES = '''
- name: Enable SSL termination on a load balancer
  rax_clb_ssl:
    balancer_name: the_loadbalancer
    state: present
    private_key: "{{ lookup('file', 'credentials/server.key' ) }}"
    certificate: "{{ lookup('file', 'credentials/server.crt' ) }}"
    intermediate_certificate: "{{ lookup('file', 'credentials/trust-chain.crt') }}"
    secure_traffic_only: true
    wait: true

- name: Disable SSL termination
  rax_clb_ssl:
    balancer_name: "{{ registered_lb.balancer.id }}"
    state: absent
    wait: true
'''

try:
    import pyrax
    HAS_PYRAX = True
except ImportError:
    HAS_PYRAX = False

def cloud_load_balancer_ssl(module, balancer_name, state, enabled, private_key,
                            certificate, intermediate_certificate, secure_port,
                            secure_traffic_only, https_redirect,
                            wait, wait_timeout):
    # Validate arguments.

    if not balancer_name:
        module.fail_json(msg='balancer_name is required.')

    if state == 'present':
        if not private_key:
            module.fail_json(msg="private_key must be provided.")
        else:
            private_key = private_key.strip()

        if not certificate:
            module.fail_json(msg="certificate must be provided.")
        else:
            certificate = certificate.strip()

    if state not in ('present', 'absent'):
        module.fail_json(msg="State must be either 'present' or 'absent'.")

    attempts = wait_timeout / 5

    # Locate the load balancer.

    clb = pyrax.cloud_loadbalancers
    if not clb:
        module.fail_json(msg='Failed to instantiate client. This '
                             'typically indicates an invalid region or an '
                             'incorrectly capitalized region name.')

    balancers = []
    for balancer in clb.list():
        if balancer_name == balancer.name or balancer_name == str(balancer.id):
            balancers.append(balancer)

    if not balancers:
        module.fail_json(msg='No load balancers matched your criteria. '
                             'Use rax_clb to create the balancer first.')

    if len(balancers) > 1:
        module.fail_json(msg="%d load balancers were matched your criteria. Try"
                             "using the balancer's id instead." % len(balancers))

    balancer = balancers[0]
    existing_ssl = balancer.get_ssl_termination()

    changed = False

    if state == 'present':
        # Apply or reconfigure SSL termination on the load balancer.
        ssl_attrs = dict(
            securePort=secure_port,
            privatekey=private_key,
            certificate=certificate,
            intermediateCertificate=intermediate_certificate,
            enabled=enabled,
            secureTrafficOnly=secure_traffic_only
        )

        needs_change = False

        if existing_ssl:
            for ssl_attr, value in ssl_attrs.iteritems():
                if ssl_attr == 'privatekey':
                    # The private key is not included in get_ssl_termination's
                    # output (as it shouldn't be). Also, if you're changing the
                    # private key, you'll also be changing the certificate,
                    # so we don't lose anything by not checking it.
                    continue

                if value is not None and existing_ssl.get(ssl_attr) != value:
                    # module.fail_json(msg='Unnecessary change', attr=ssl_attr, value=value, existing=existing_ssl.get(ssl_attr))
                    needs_change = True
        else:
            needs_change = True

        if needs_change:
            balancer.add_ssl_termination(**ssl_attrs)
            changed = True
    elif state == 'absent':
        # Remove SSL termination if it's already configured.
        if existing_ssl:
            balancer.delete_ssl_termination()
            changed = True

    if https_redirect is not None and balancer.httpsRedirect != https_redirect:
        if changed:
            # This wait is unavoidable because load balancers are immutable
            # while the SSL termination changes above are being applied.
            pyrax.utils.wait_for_build(balancer, interval=5, attempts=attempts)

        balancer.update(httpsRedirect=https_redirect)
        changed = True

    if changed and wait:
        pyrax.utils.wait_for_build(balancer, interval=5, attempts=attempts)

    balancer.get()
    new_ssl_termination = balancer.get_ssl_termination()

    # Intentionally omit the private key from the module output, so you don't
    # accidentally echo it with `ansible-playbook -v` or `debug`, and the
    # certificate, which is just long. Convert other attributes to snake_case
    # and include https_redirect at the top-level.
    if new_ssl_termination:
        new_ssl = dict(
            enabled=new_ssl_termination['enabled'],
            secure_port=new_ssl_termination['securePort'],
            secure_traffic_only=new_ssl_termination['secureTrafficOnly']
        )
    else:
        new_ssl = None

    result = dict(
        changed=changed,
        https_redirect=balancer.httpsRedirect,
        ssl_termination=new_ssl
    )
    success = True

    if balancer.status == 'ERROR':
        result['msg'] = '%s failed to build' % balancer.id
        success = False
    elif wait and balancer.status not in ('ACTIVE', 'ERROR'):
        result['msg'] = 'Timeout waiting on %s' % balancer.id
        success = False

    if success:
        module.exit_json(**result)
    else:
        module.fail_json(**result)

def main():
    argument_spec = rax_argument_spec()
    argument_spec.update(dict(
        balancer_name=dict(type='str'),
        state=dict(default='present', choices=['present', 'absent']),
        enabled=dict(type='bool', default=True),
        private_key=dict(),
        certificate=dict(),
        intermediate_certificate=dict(),
        secure_port=dict(type='int', default=443),
        secure_traffic_only=dict(type='bool', default=False),
        https_redirect=dict(type='bool'),
        wait=dict(type='bool', default=False),
        wait_timeout=dict(type='int', default=300)
    ))

    module = AnsibleModule(
        argument_spec=argument_spec,
        required_together=rax_required_together(),
    )

    if not HAS_PYRAX:
        module.fail_json(msg='pyrax is required for this module.')

    balancer_name = module.params.get('balancer_name')
    state = module.params.get('state')
    enabled = module.boolean(module.params.get('enabled'))
    private_key = module.params.get('private_key')
    certificate = module.params.get('certificate')
    intermediate_certificate = module.params.get('intermediate_certificate')
    secure_port = module.params.get('secure_port')
    secure_traffic_only = module.boolean(module.params.get('secure_traffic_only'))
    https_redirect = module.boolean(module.params.get('https_redirect'))
    wait = module.boolean(module.params.get('wait'))
    wait_timeout = module.params.get('wait_timeout')

    setup_rax_module(module, pyrax)

    cloud_load_balancer_ssl(
        module, balancer_name, state, enabled, private_key, certificate,
        intermediate_certificate, secure_port, secure_traffic_only,
        https_redirect, wait, wait_timeout
    )

from ansible.module_utils.basic import *
from ansible.module_utils.rax import *

main()
<|endoftext|>


<commit_before># a very minimal script for InMoov
# although this script is very short you can still
# do voice control of a right hand or finger box
# for any command which you say - you will be required to say a confirmation
# e.g. you say -> open hand, InMoov will ask -> "Did you say open hand?", you will need to
# respond with a confirmation ("yes","correct","yeah","ya")

useVirtualArduinos = True

rightPort = "COM8"

i01 = Runtime.start("i01", "InMoov")

if (useVirtualArduinos):
    rightVirtual = Runtime.start("rightVirtual","VirtualArduino")
    leftVirtual  = Runtime.start("leftVirtual","VirtualArduino")
    rightVirtual.connect(rightPort)

# starting parts
i01.startEar()
i01.startMouth()
#to tweak the default voice
##############
i01.startRightHand(rightPort)
# tweaking defaults settings of right hand
#i01.rightHand.thumb.setMinMax(55,135)
#i01.rightHand.index.setMinMax(0,160)
#i01.rightHand.majeure.setMinMax(0,140)
#i01.rightHand.ringFinger.setMinMax(48,145)
#i01.rightHand.pinky.setMinMax(45,146)
#i01.rightHand.thumb.map(0,180,55,135)
#i01.rightHand.index.map(0,180,0,160)
#i01.rightHand.majeure.map(0,180,0,140)
#i01.rightHand.ringFinger.map(0,180,48,145)
#i01.rightHand.pinky.map(0,180,45,146)
#################

# verbal commands
ear = i01.ear

ear.addCommand("attach right hand", "i01.rightHand", "attach")
ear.addCommand("disconnect right hand", "i01.rightHand", "detach")
ear.addCommand("rest", i01.getName(), "rest")
ear.addCommand("open hand", "python", "handopen")
ear.addCommand("close hand", "python", "handclose")
ear.addCommand("capture gesture", ear.getName(), "captureGesture")
ear.addCommand("manual", ear.getName(), "lockOutAllGrammarExcept", "voice control")
ear.addCommand("voice control", ear.getName(), "clearLock")

ear.addComfirmations("yes","correct","yeah","ya")
ear.addNegations("no","wrong","nope","nah")

ear.startListening()

def handopen():
  i01.moveHand("left",0,0,0,0,0)
  i01.moveHand("right",0,0,0,0,0)

def handclose():
  i01.moveHand("left",180,180,180,180,180)
  i01.moveHand("right",180,180,180,180,180)
<commit_msg>refactoring InMoov starter script<commit_after># a very minimal script for InMoov
# although this script is very short you can still
# do voice control of a right hand or finger box
# for any command which you say - you will be required to say a confirmation
# e.g. you say -> open hand, InMoov will ask -> "Did you say open hand?", you will need to
# respond with a confirmation ("yes","correct","yeah","ya")
useRightVirtualArduino = True
useLeftVirtualArduino = True
useVirtualInMoov = True

useVirtualArduinos = True

rightPort = "COM8"
leftPort = "COM10"

i01 = Runtime.start("i01", "InMoov")

if (useRightVirtualArduino):
    rightVirtual = Runtime.start("rightVirtual","VirtualArduino")
    rightVirtual.connect(rightPort)

if (useLeftVirtualArduino):
    leftVirtual  = Runtime.start("leftVirtual","VirtualArduino")
    leftVirtual.connect(leftPort)

# starting parts
i01.startEar()
i01.startMouth()
#to tweak the default voice
##############
i01.startRightHand(rightPort)
# tweaking defaults settings of right hand
#i01.rightHand.thumb.setMinMax(55,135)
#i01.rightHand.index.setMinMax(0,160)
#i01.rightHand.majeure.setMinMax(0,140)
#i01.rightHand.ringFinger.setMinMax(48,145)
#i01.rightHand.pinky.setMinMax(45,146)
#i01.rightHand.thumb.map(0,180,55,135)
#i01.rightHand.index.map(0,180,0,160)
#i01.rightHand.majeure.map(0,180,0,140)
#i01.rightHand.ringFinger.map(0,180,48,145)
#i01.rightHand.pinky.map(0,180,45,146)
#################

# verbal commands
ear = i01.ear

ear.addCommand("attach right hand", "i01.rightHand", "attach")
ear.addCommand("disconnect right hand", "i01.rightHand", "detach")
ear.addCommand("rest", i01.getName(), "rest")
ear.addCommand("open hand", "python", "handopen")
ear.addCommand("close hand", "python", "handclose")
ear.addCommand("capture gesture", ear.getName(), "captureGesture")
ear.addCommand("manual", ear.getName(), "lockOutAllGrammarExcept", "voice control")
ear.addCommand("voice control", ear.getName(), "clearLock")

ear.addComfirmations("yes","correct","yeah","ya")
ear.addNegations("no","wrong","nope","nah")

ear.startListening()

def handopen():
  i01.moveHand("left",0,0,0,0,0)
  i01.moveHand("right",0,0,0,0,0)

def handclose():
  i01.moveHand("left",180,180,180,180,180)
  i01.moveHand("right",180,180,180,180,180)
<|endoftext|>


<commit_before>from django import forms
from django.utils.translation import ugettext_lazy as _
from django_countries import countries

from postal.settings import POSTAL_ADDRESS_LINE1, POSTAL_ADDRESS_LINE2, POSTAL_ADDRESS_CITY, POSTAL_ADDRESS_STATE, POSTAL_ADDRESS_CODE

country_list = (('', '-'*45),) + countries.COUNTRIES
country_dict = {}
for (key, value) in countries.COUNTRIES: country_dict[key] = value

class PostalAddressForm(forms.Form):
    line1 = forms.CharField(label=POSTAL_ADDRESS_LINE1[0], required=POSTAL_ADDRESS_LINE1[1], max_length=100)
    line2 = forms.CharField(label=POSTAL_ADDRESS_LINE2[0], required=POSTAL_ADDRESS_LINE2[1], max_length=100)
    city = forms.CharField(label=POSTAL_ADDRESS_CITY[0], required=POSTAL_ADDRESS_CITY[1], max_length=100)
    state = forms.CharField(label=POSTAL_ADDRESS_STATE[0], required=POSTAL_ADDRESS_STATE[1], max_length=100)
    code = forms.CharField(label=POSTAL_ADDRESS_CODE[0], required=POSTAL_ADDRESS_CODE[1], max_length=100)
    country = forms.ChoiceField(choices=country_list)

    def clean_country(self):
        data = self.cleaned_data['country']
        if data not in country_dict.keys():
            raise forms.ValidationError("You must select a country")
        return data
<commit_msg>Added translatable label to Country field.<commit_after>from django import forms
from django.utils.translation import ugettext_lazy as _
from django_countries import countries

from postal.settings import POSTAL_ADDRESS_LINE1, POSTAL_ADDRESS_LINE2, POSTAL_ADDRESS_CITY, POSTAL_ADDRESS_STATE, POSTAL_ADDRESS_CODE

country_list = (('', '-'*45),) + countries.COUNTRIES
country_dict = {}
for (key, value) in countries.COUNTRIES: country_dict[key] = value

class PostalAddressForm(forms.Form):
    line1 = forms.CharField(label=POSTAL_ADDRESS_LINE1[0], required=POSTAL_ADDRESS_LINE1[1], max_length=100)
    line2 = forms.CharField(label=POSTAL_ADDRESS_LINE2[0], required=POSTAL_ADDRESS_LINE2[1], max_length=100)
    city = forms.CharField(label=POSTAL_ADDRESS_CITY[0], required=POSTAL_ADDRESS_CITY[1], max_length=100)
    state = forms.CharField(label=POSTAL_ADDRESS_STATE[0], required=POSTAL_ADDRESS_STATE[1], max_length=100)
    code = forms.CharField(label=POSTAL_ADDRESS_CODE[0], required=POSTAL_ADDRESS_CODE[1], max_length=100)
    country = forms.ChoiceField(label=_(u"Country"), choices=country_list)

    def clean_country(self):
        data = self.cleaned_data['country']
        if data not in country_dict.keys():
            raise forms.ValidationError("You must select a country")
        return data
<|endoftext|>


<commit_before>import numpy as np

from tools.agent_utils import convert_action_to_int
from tools.game_tree.nodes import BoardCardsNode, ActionNode, TerminalNode
from tools.tree_utils import get_parent_action
from tools.hand_evaluation import get_utility
from tools.io_util import read_strategy_from_file
from tools.utils import is_unique, flatten
from utility_estimation.utils import get_all_board_cards, get_board_cards
from evaluation.player_utility import PlayerUtility

class AivatUtilityEstimator():
    def __init__(self, game, mucking_enabled, **args):
        if game.get_num_players() != 2:
            raise AttributeError(
                'Only games with 2 players are supported')

        self.game = game
        self.mucking_enabled = mucking_enabled

        if mucking_enabled:
            raise AttributeError('Mucking not yet supported')

        if 'equilibirum_strategy_path' not in args:
            raise AttributeError(
                '"equilibirum_strategy_path" argument not provided')

        equilibirum_strategy_path = args['equilibirum_strategy_path']
        equilibirum_strategy, _ = read_strategy_from_file(game, equilibirum_strategy_path)
        utilities_dict = {}

        def callback(nodes, utilities):
            nonlocal utilities_dict

            key = ';'.join(map(lambda m: str(m), nodes))
            if key not in utilities_dict:
                utilities_dict[key] = utilities

        PlayerUtility(game).get_player_utilities(
            [equilibirum_strategy] * 2,
            [],
            [],
            [False] * 2,
            callback)
        self.utilities_dict = utilities_dict

    def get_utility_estimations(self, state, player, sampling_strategy, evaluated_strategies=None):
        if evaluated_strategies is None:
            evaluated_strategies = [sampling_strategy]
        else:
            raise AttributeError('AIVAT does not support importance sampling so far')

        num_players = self.game.get_num_players()
        opponent_player = (player + 1) % 2

        num_evaluated_strategies = len(evaluated_strategies)

        utilities = np.zeros(num_evaluated_strategies)

        any_player_folded = False
        for p in range(num_players):
            any_player_folded = any_player_folded or state.get_player_folded(p)

        all_board_cards = get_all_board_cards(self.game, state)

        opponent_hole_cards = None
        possible_player_hole_cards = None
        if any_player_folded and self.mucking_enabled:
            possible_player_hole_cards = list(filter(
                lambda hole_cards: is_unique(hole_cards, all_board_cards),
                sampling_strategy.children))
        else:
            opponent_hole_cards = [state.get_hole_card(opponent_player, c) for c in range(self.game.get_num_hole_cards())]
            possible_player_hole_cards = list(filter(
                lambda hole_cards: is_unique(hole_cards, opponent_hole_cards, all_board_cards),
                sampling_strategy.children))
        nodes = [
            [
                expert_node.children[hole_cards]
                for hole_cards in possible_player_hole_cards]
            for expert_node in evaluated_strategies]
        sampling_strategy_nodes = [sampling_strategy.children[hole_cards] for hole_cards in possible_player_hole_cards]

        opponent_node = sampling_strategy.children[tuple(sorted(opponent_hole_cards))]

        num_nodes = len(possible_player_hole_cards)
        evaluated_strategies_reach_probabilities = np.ones([num_evaluated_strategies, num_nodes])
        sampling_strategy_reach_probabilities = np.ones(num_nodes)

        def add_terminals_to_utilities(pot_commitment, players_folded, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities):
            nonlocal utilities
            nonlocal player

            sampling_strategy_reach_probability_sum = np.sum(sampling_strategy_reach_probabilities)
            if sampling_strategy_reach_probability_sum == 0:
                return

            for i in range(num_nodes):
                utility = None
                if players_folded[player]:
                    utility = -pot_commitment[player]
                else:
                    hole_cards = [possible_player_hole_cards[i] if p == player else opponent_hole_cards for p in range(num_players)]
                    utility = get_utility(
                        hole_cards,
                        all_board_cards,
                        players_folded,
                        pot_commitment)[player]
                utilities += utility * (evaluated_strategies_reach_probabilities[:, i] / sampling_strategy_reach_probability_sum)

        def update_reach_proabilities(action, sampling_strategy_nodes, nodes, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities):
            for i in range(num_nodes):
                if sampling_strategy_reach_probabilities is not None:
                    sampling_strategy_reach_probabilities[i] *= sampling_strategy_nodes[i].strategy[action]

                for j in range(num_evaluated_strategies):
                    evaluated_strategies_reach_probabilities[j, i] *= nodes[j][i].strategy[action]

        round_index = 0
        action_index = 0
        while True:
            node = nodes[0][0]
            if isinstance(node, BoardCardsNode):
                new_board_cards = get_board_cards(self.game, state, round_index)
                nodes = [[expert_node.children[new_board_cards] for expert_node in expert_nodes] for expert_nodes in nodes]
                sampling_strategy_nodes = [node.children[new_board_cards] for node in sampling_strategy_nodes]
                opponent_node = opponent_node.children[new_board_cards]
            elif isinstance(node, ActionNode):
                action = convert_action_to_int(state.get_action_type(round_index, action_index))
                if node.player == player:
                    # Calculate correction term and add it to utility
                    histories_actions_utilities = {}
                    for h in range(num_nodes):
                        history_actions_utilities = {}
                        histories_actions_utilities[h] = history_actions_utilities
                        for a in node.children:
                            nodes_tmp = [None, None]
                            nodes_tmp[player] = sampling_strategy_nodes[h].children[a]
                            nodes_tmp[opponent_player] = opponent_node.children[a]
                            key = ';'.join(map(lambda m: str(m), nodes_tmp))
                            history_actions_utilities[a] = self.utilities_dict[key][player]

                    history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                    current_history_expected_value = 0
                    for i in range(num_nodes):
                        for a in sampling_strategy_nodes[i].children:
                            current_history_expected_value += histories_actions_utilities[i][a] \
                                * sampling_strategy_reach_probabilities[i] * sampling_strategy_nodes[i].strategy[a]

                    update_reach_proabilities(
                        action,
                        sampling_strategy_nodes,
                        nodes,
                        sampling_strategy_reach_probabilities,
                        evaluated_strategies_reach_probabilities)

                    next_history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                    next_history_expected_value = 0
                    for i in range(num_nodes):
                        next_history_expected_value += histories_actions_utilities[i][action] \
                            * sampling_strategy_reach_probabilities[i]
                    utilities[0] += \
                        (current_history_expected_value / history_sampling_strategy_reach_probabilities_sum) \
                        - (next_history_expected_value / next_history_sampling_strategy_reach_probabilities_sum)

                action_index += 1
                if action_index == state.get_num_actions(round_index):
                    round_index += 1
                    action_index = 0
                nodes = [[expert_node.children[action] for expert_node in expert_nodes] for expert_nodes in nodes]
                sampling_strategy_nodes = [node.children[action] for node in sampling_strategy_nodes]
                opponent_node = opponent_node.children[action]
            elif isinstance(node, TerminalNode):
                players_folded = [state.get_player_folded(p) for p in range(num_players)]
                add_terminals_to_utilities(node.pot_commitment, players_folded, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities)
                break

        return utilities
<commit_msg>Add correction terms for board cards to AIVAT<commit_after>import numpy as np

from tools.agent_utils import convert_action_to_int
from tools.game_tree.nodes import BoardCardsNode, ActionNode, TerminalNode
from tools.tree_utils import get_parent_action
from tools.hand_evaluation import get_utility
from tools.io_util import read_strategy_from_file
from tools.utils import is_unique, flatten
from utility_estimation.utils import get_all_board_cards, get_board_cards
from evaluation.player_utility import PlayerUtility

class AivatUtilityEstimator():
    def __init__(self, game, mucking_enabled, **args):
        if game.get_num_players() != 2:
            raise AttributeError(
                'Only games with 2 players are supported')

        self.game = game
        self.mucking_enabled = mucking_enabled

        if mucking_enabled:
            raise AttributeError('Mucking not yet supported')

        if 'equilibirum_strategy_path' not in args:
            raise AttributeError(
                '"equilibirum_strategy_path" argument not provided')

        equilibirum_strategy_path = args['equilibirum_strategy_path']
        equilibirum_strategy, _ = read_strategy_from_file(game, equilibirum_strategy_path)
        utilities_dict = {}

        def callback(nodes, utilities):
            nonlocal utilities_dict

            key = ';'.join(map(lambda m: str(m), nodes))
            if key not in utilities_dict:
                utilities_dict[key] = utilities

        PlayerUtility(game).get_player_utilities(
            [equilibirum_strategy] * 2,
            [],
            [],
            [False] * 2,
            callback)
        self.utilities_dict = utilities_dict

    def get_utility_estimations(self, state, player, sampling_strategy, evaluated_strategies=None):
        if evaluated_strategies is None:
            evaluated_strategies = [sampling_strategy]
        else:
            raise AttributeError('AIVAT does not support importance sampling so far')

        num_players = self.game.get_num_players()
        opponent_player = (player + 1) % 2

        num_evaluated_strategies = len(evaluated_strategies)

        utilities = np.zeros(num_evaluated_strategies)

        any_player_folded = False
        for p in range(num_players):
            any_player_folded = any_player_folded or state.get_player_folded(p)

        all_board_cards = get_all_board_cards(self.game, state)

        opponent_hole_cards = None
        possible_player_hole_cards = None
        if any_player_folded and self.mucking_enabled:
            possible_player_hole_cards = list(filter(
                lambda hole_cards: is_unique(hole_cards, all_board_cards),
                sampling_strategy.children))
        else:
            opponent_hole_cards = [state.get_hole_card(opponent_player, c) for c in range(self.game.get_num_hole_cards())]
            possible_player_hole_cards = list(filter(
                lambda hole_cards: is_unique(hole_cards, opponent_hole_cards, all_board_cards),
                sampling_strategy.children))
        nodes = [
            [
                expert_node.children[hole_cards]
                for hole_cards in possible_player_hole_cards]
            for expert_node in evaluated_strategies]
        sampling_strategy_nodes = [sampling_strategy.children[hole_cards] for hole_cards in possible_player_hole_cards]

        opponent_node = sampling_strategy.children[tuple(sorted(opponent_hole_cards))]

        num_nodes = len(possible_player_hole_cards)
        evaluated_strategies_reach_probabilities = np.ones([num_evaluated_strategies, num_nodes])
        sampling_strategy_reach_probabilities = np.ones(num_nodes)

        def add_terminals_to_utilities(pot_commitment, players_folded, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities):
            nonlocal utilities
            nonlocal player

            sampling_strategy_reach_probability_sum = np.sum(sampling_strategy_reach_probabilities)
            if sampling_strategy_reach_probability_sum == 0:
                return

            for i in range(num_nodes):
                utility = None
                if players_folded[player]:
                    utility = -pot_commitment[player]
                else:
                    hole_cards = [possible_player_hole_cards[i] if p == player else opponent_hole_cards for p in range(num_players)]
                    utility = get_utility(
                        hole_cards,
                        all_board_cards,
                        players_folded,
                        pot_commitment)[player]
                utilities += utility * (evaluated_strategies_reach_probabilities[:, i] / sampling_strategy_reach_probability_sum)

        def update_reach_proabilities(action, sampling_strategy_nodes, nodes, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities):
            for i in range(num_nodes):
                sampling_strategy_reach_probabilities[i] *= sampling_strategy_nodes[i].strategy[action]
                for j in range(num_evaluated_strategies):
                    evaluated_strategies_reach_probabilities[j, i] *= nodes[j][i].strategy[action]

        round_index = 0
        action_index = 0
        while True:
            node = nodes[0][0]
            if isinstance(node, BoardCardsNode):
                new_board_cards = get_board_cards(self.game, state, round_index)

                num_board_cards = len(opponent_node.children)

                histories_actions_utilities = {}
                for i in range(num_nodes):
                    history_actions_utilities = {}
                    histories_actions_utilities[i] = history_actions_utilities
                    for a in filter(lambda a: a in sampling_strategy_nodes[i].children, opponent_node.children):
                        nodes_tmp = [None, None]
                        nodes_tmp[player] = sampling_strategy_nodes[i].children[a]
                        nodes_tmp[opponent_player] = opponent_node.children[a]
                        key = ';'.join(map(lambda m: str(m), nodes_tmp))
                        history_actions_utilities[a] = self.utilities_dict[key][player]

                history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                current_history_expected_value = 0
                for i in range(num_nodes):
                    for a in filter(lambda a: a in sampling_strategy_nodes[i].children, opponent_node.children):
                        current_history_expected_value += histories_actions_utilities[i][a] \
                            * sampling_strategy_reach_probabilities[i] / num_board_cards

                for i in range(num_nodes):
                    sampling_strategy_reach_probabilities[i] /= num_board_cards
                    for j in range(num_evaluated_strategies):
                        evaluated_strategies_reach_probabilities[j, i] /= num_board_cards

                next_history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                next_history_expected_value = 0
                for i in range(num_nodes):
                    next_history_expected_value += histories_actions_utilities[i][new_board_cards] \
                        * sampling_strategy_reach_probabilities[i]
                utilities[0] += \
                    (current_history_expected_value / history_sampling_strategy_reach_probabilities_sum) \
                    - (next_history_expected_value / next_history_sampling_strategy_reach_probabilities_sum)

                nodes = [[expert_node.children[new_board_cards] for expert_node in expert_nodes] for expert_nodes in nodes]
                sampling_strategy_nodes = [node.children[new_board_cards] for node in sampling_strategy_nodes]
                opponent_node = opponent_node.children[new_board_cards]
            elif isinstance(node, ActionNode):
                action = convert_action_to_int(state.get_action_type(round_index, action_index))
                if node.player == player:
                    # Calculate correction term and add it to utility
                    histories_actions_utilities = {}
                    for i in range(num_nodes):
                        history_actions_utilities = {}
                        histories_actions_utilities[i] = history_actions_utilities
                        for a in node.children:
                            nodes_tmp = [None, None]
                            nodes_tmp[player] = sampling_strategy_nodes[i].children[a]
                            nodes_tmp[opponent_player] = opponent_node.children[a]
                            key = ';'.join(map(lambda m: str(m), nodes_tmp))
                            history_actions_utilities[a] = self.utilities_dict[key][player]

                    history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                    current_history_expected_value = 0
                    for i in range(num_nodes):
                        for a in sampling_strategy_nodes[i].children:
                            current_history_expected_value += histories_actions_utilities[i][a] \
                                * sampling_strategy_reach_probabilities[i] * sampling_strategy_nodes[i].strategy[a]

                    update_reach_proabilities(
                        action,
                        sampling_strategy_nodes,
                        nodes,
                        sampling_strategy_reach_probabilities,
                        evaluated_strategies_reach_probabilities)

                    next_history_sampling_strategy_reach_probabilities_sum = np.sum(sampling_strategy_reach_probabilities)
                    next_history_expected_value = 0
                    for i in range(num_nodes):
                        next_history_expected_value += histories_actions_utilities[i][action] \
                            * sampling_strategy_reach_probabilities[i]
                    utilities[0] += \
                        (current_history_expected_value / history_sampling_strategy_reach_probabilities_sum) \
                        - (next_history_expected_value / next_history_sampling_strategy_reach_probabilities_sum)

                action_index += 1
                if action_index == state.get_num_actions(round_index):
                    round_index += 1
                    action_index = 0
                nodes = [[expert_node.children[action] for expert_node in expert_nodes] for expert_nodes in nodes]
                sampling_strategy_nodes = [node.children[action] for node in sampling_strategy_nodes]
                opponent_node = opponent_node.children[action]
            elif isinstance(node, TerminalNode):
                players_folded = [state.get_player_folded(p) for p in range(num_players)]
                add_terminals_to_utilities(node.pot_commitment, players_folded, sampling_strategy_reach_probabilities, evaluated_strategies_reach_probabilities)
                break

        return utilities
<|endoftext|>


<commit_before># Released under The MIT License (MIT)
# http://opensource.org/licenses/MIT
# Copyright (c) 2013 SCoT Development Team

""" Object oriented API to SCoT """

import numpy as np
from .varica import mvarica
from .datatools import dot_special
from .connectivity import Connectivity
from . import var
from eegtopo.topoplot import Topoplot

try:
    import matplotlib.pyplot as plt
    _have_pyplot = True
except ImportError:
    _have_pyplot = False

class SCoT:

    def __init__(self, var_order, var_delta=None, locations=None, reducedim=0.99, nfft=512, fs=2, backend=None):
        self.data_ = None
        self.cl_ = None
        self.fs_ = fs
        self.unmixing_ = None
        self.mixing_ = None
        self.activations_ = None
        self.var_model_ = None
        self.var_cov_ = None
        self.var_order_ = var_order
        self.var_delta_ = var_delta
        self.connectivity_ = None
        self.locations_ = locations
        self.reducedim_ = reducedim
        self.nfft_ = nfft
        self.backend_ = backend

        self.topo_ = None
        self.mixmaps_ = []
        self.unmixmaps_ = []

    def setData(self, data, cl=None):
        self.data_ = np.atleast_3d(data)
        self.cl_ = cl
        self.var_model_ = None
        self.var_cov_ = None
        self.connectivity_ = None

        if self.unmixing_ != None:
            self.activations_ = dot_special(self.data_, self.unmixing_)

    def doMVARICA(self):
        if self.data_ == None:
            raise RuntimeError("MVARICA requires data to be set")
        if self.reducedim_ < 1:
            rv = self.reducedim_
            nc = None
        else:
            rv = None
            nc = self.reducedim_
        result = mvarica(X=self.data_, P=self.var_order_, retain_variance=rv, numcomp=nc, delta=self.var_delta_, backend=self.backend_)
        self.mixing_ = result.mixing
        self.unmixing_ = result.unmixing
        self.var_model_ = result.B
        self.var_cov_ = result.C
        self.var_delta_ = result.delta
        self.connectivity_ = Connectivity(self.var_model_, self.var_cov_, self.nfft_)

    def fitVAR(self):
        if self.activations_ == None:
            raise RuntimeError("VAR fitting requires activations (call setData after doMVARICA)")
        if self.cl_ == None:
            self.var_model_, self.var_cov_ = var.fit(data=self.activations_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
            self.connectivity_ = Connectivity(self.var_model_, self.var_cov_, self.nfft_)
        else:
            self.var_model_, self.var_cov_ = var.fit_multiclass(data=self.activations_, cl=self.cl_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
            self.connectivity_ = {}
            for c in np.unique(self.cl_):
                self.connectivity_[c] = Connectivity(self.var_model_[c], self.var_cov_[c], self.nfft_)

    def getConnectivity(self, measure):
        if self.connectivity_ == None:
            raise RuntimeError("Connectivity requires a VAR model (run doMVARICA or fitVAR first)")
        if isinstance(self.connectivity_, dict):
            result = {}
            for c in np.unique(self.cl_):
                result[c] = getattr(self.connectivity_[c], measure)()
            return result
        else:
            return getattr(self.connectivity_, measure)()

    def getTFConnectivity(self, measure, winlen, winstep):
        if self.activations_ == None:
            raise RuntimeError("Time/Frequency Connectivity requires activations (call setData after doMVARICA)")
        [N,M,T] = self.activations_.shape

        Nstep = (N-winlen)//winstep

        if self.cl_ == None:
            result = np.zeros((M, M, Nstep, self.nfft_), np.complex64)
            i = 0
            for n in range(0, N-winlen, winstep):
                win = np.arange(winlen) + n
                data = self.activations_[win,:,:]
                B, C = var.fit(data, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
                con = Connectivity(B, C, self.nfft_)
                result[:,:,i,:] = getattr(con, measure)()
                i += 1

        else:
            result = {}
            for c in np.unique(self.cl_):
                result[c] = np.zeros((M, M, Nstep, self.nfft_), np.complex64)
            i = 0
            for n in range(0, N-winlen, winstep):
                win = np.arange(winlen) + n
                data = self.activations_[win,:,:]
                B, C = var.fit_multiclass(data, cl=self.cl_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
                for c in np.unique(self.cl_):
                    con = Connectivity(B[c], C[c], self.nfft_)
                    result[c][:,:,i,:] = getattr(con, measure)()
                i += 1
        return result

    def preparePlots(self, mixing=False, unmixing=False):
        if self.locations_ == None:
            raise RuntimeError("Need sensor locations for plotting")

        if self.topo_ == None:
            self.topo_ = Topoplot( )
            self.topo_.set_locations(self.locations_)

        if mixing and not self.mixmaps_:
            for i in range(self.mixing_.shape[0]):
                self.topo_.set_values(self.mixing_[i,:])
                self.topo_.create_map()
                self.mixmaps_.append(self.topo_.get_map())

        if unmixing and not self.unmixmaps_:
            for i in range(self.unmixing_.shape[1]):
                self.topo_.set_values(self.unmixing_[:,i])
                self.topo_.create_map()
                self.unmixmaps_.append(self.topo_.get_map())

    def showPlots(self):
        plt.show()

    def plotComponents(self, global_scale=None):
        """ global_scale:
               None - scales each topo individually
               1-99 - percentile of maximum of all plots
        """
        if not _have_pyplot:
            raise ImportError("matplotlib.pyplot is required for plotting")
        if self.unmixing_ == None and self.mixing_ == None:
            raise RuntimeError("No components available (run doMVARICA first)")
        self.preparePlots(True, True)

        M = self.mixing_.shape[0]

        urange, mrange = None, None

        if global_scale:
            tmp = np.asarray(self.unmixmaps_)
            tmp = tmp[np.logical_not(np.isnan(tmp))]
            umax = np.percentile(np.abs(tmp), global_scale)
            umin = -umax
            urange = [umin,umax]

            tmp = np.asarray(self.mixmaps_)
            tmp = tmp[np.logical_not(np.isnan(tmp))]
            mmax = np.percentile(np.abs(tmp), global_scale)
            mmin = -mmax
            mrange = [umin,umax]

        Y = np.floor(np.sqrt(M*3/4))
        X = np.ceil(M/Y)

        fig = plt.figure()

        axes = []
        for m in range(M):
            axes.append(fig.add_subplot(2*Y, X, m+1))
            h1 = self._plotUnmixing(axes[-1], m, crange=urange)

            axes.append(fig.add_subplot(2*Y, X, M+m+1))
            h1 = self._plotMixing(axes[-1], m, crange=mrange)

        for a in axes:
            a.set_yticks([])
            a.set_xticks([])
            a.set_frame_on(False)

        axes[0].set_ylabel('Unmixing weights')
        axes[1].set_ylabel('Scalp projections')

        #plt.colorbar(h1, plt.subplot(2, M+1, M+1))
        #plt.colorbar(h2, plt.subplot(2, M+1, 0))

    def plotConnectivity(self, measure):
        if not _have_pyplot:
            raise ImportError("matplotlib.pyplot is required for plotting")
        self.preparePlots(True, False)
        fig = plt.figure()
        if isinstance(self.connectivity_, dict):
            for c in np.unique(self.cl_):
                cm = getattr(self.connectivity_[c], measure)()
                self._plotSpectral(fig, cm)
        else:
            cm = getattr(self.connectivity_, measure)()
            self._plotSpectral(fig, cm)

    def _plotSpectral(self, fig, A):
        [N,M,F] = A.shape
        freq = np.linspace(0, self.fs_/2, F)

        lowest, highest = np.inf, -np.inf

        axes = []
        for n in range(N):
            arow = []
            for m in range(M):
                ax = fig.add_subplot(N, M, m+n*M+1)
                arow.append(ax)

                if n == m:
                    self._plotMixing(ax, m)
                    ax.set_yticks([])
                    ax.set_xticks([])
                    ax.set_frame_on(False)
                else:
                    ax.plot(freq, A[n,m,:])
                    lowest = min(lowest, np.min(A[n,m,:]))
                    highest = max(highest, np.max(A[n,m,:]))
                    ax.set_xlim(0, self.fs_/2)
            axes.append(arow)

        for n in range(N):
            for m in range(M):
                if n == m:
                    pass
                else:
                    axes[n][m].set_ylim(lowest, highest)
                    if 0 < n < N-1:
                        axes[n][m].set_xticks([])
                    if 0 < m < M-1:
                        axes[n][m].set_yticks([])
            axes[n][0].yaxis.tick_left()
            axes[n][-1].yaxis.tick_right()

        for m in range(M):
            axes[0][m].xaxis.tick_top()
            axes[-1][m].xaxis.tick_bottom()

        fig.text(0.5, 0.05, 'frequency', horizontalalignment='center')
        fig.text(0.05, 0.5, 'magnitude', horizontalalignment='center', rotation='vertical')

    def _plotMixing(self, axis, idx, crange=None):
        self.topo_.set_map(self.mixmaps_[idx])
        h = self.topo_.plot_map(axis, crange=crange)
        self.topo_.plot_locations(axis)
        self.topo_.plot_head(axis)
        return h

    def _plotUnmixing(self, axis, idx, crange=None):
        self.topo_.set_map(self.unmixmaps_[idx])
        h = self.topo_.plot_map(axis, crange=crange)
        self.topo_.plot_locations(axis)
        self.topo_.plot_head(axis)
        return h<commit_msg>Added __str__ member to SCoT<commit_after># Released under The MIT License (MIT)
# http://opensource.org/licenses/MIT
# Copyright (c) 2013 SCoT Development Team

""" Object oriented API to SCoT """

import numpy as np
from .varica import mvarica
from .datatools import dot_special
from .connectivity import Connectivity
from . import var
from eegtopo.topoplot import Topoplot

try:
    import matplotlib.pyplot as plt
    _have_pyplot = True
except ImportError:
    _have_pyplot = False

class SCoT:

    def __init__(self, var_order, var_delta=None, locations=None, reducedim=0.99, nfft=512, fs=2, backend=None):
        self.data_ = None
        self.cl_ = None
        self.fs_ = fs
        self.unmixing_ = None
        self.mixing_ = None
        self.activations_ = None
        self.var_model_ = None
        self.var_cov_ = None
        self.var_order_ = var_order
        self.var_delta_ = var_delta
        self.connectivity_ = None
        self.locations_ = locations
        self.reducedim_ = reducedim
        self.nfft_ = nfft
        self.backend_ = backend

        self.topo_ = None
        self.mixmaps_ = []
        self.unmixmaps_ = []

    def __str__(self):

        if self.data_ is not None:
            data = '%d samples, %d channels, %d trials'%self.data_.shape
        else:
            data = 'None'

        cl = str(np.unique(self.cl_))

        if self.unmixing_ is not None:
            sources = str(self.unmixing_.shape[1])
        else:
            sources = 'None'

        if self.var_model_ is None:
            var = 'None'
        elif isinstance(self.var_model_, dict):
            var = str(len(self.var_model_))
        else:
            var = '1'

        s = 'SCoT(var_order = %d):\n'%self.var_order_
        s += '  Data      : ' + data + '\n'
        s += '  Classes   : ' + cl + '\n'
        s += '  Sources   : ' + sources + '\n'
        s += '  VAR models: ' + var + '\n'

        return s

    def setData(self, data, cl=None):
        self.data_ = np.atleast_3d(data)
        self.cl_ = cl
        self.var_model_ = None
        self.var_cov_ = None
        self.connectivity_ = None

        if self.unmixing_ != None:
            self.activations_ = dot_special(self.data_, self.unmixing_)

    def doMVARICA(self):
        if self.data_ == None:
            raise RuntimeError("MVARICA requires data to be set")
        if self.reducedim_ < 1:
            rv = self.reducedim_
            nc = None
        else:
            rv = None
            nc = self.reducedim_
        result = mvarica(X=self.data_, P=self.var_order_, retain_variance=rv, numcomp=nc, delta=self.var_delta_, backend=self.backend_)
        self.mixing_ = result.mixing
        self.unmixing_ = result.unmixing
        self.var_model_ = result.B
        self.var_cov_ = result.C
        self.var_delta_ = result.delta
        self.connectivity_ = Connectivity(self.var_model_, self.var_cov_, self.nfft_)

    def fitVAR(self):
        if self.activations_ == None:
            raise RuntimeError("VAR fitting requires source activations (run doMVARICA first)")
        if self.cl_ == None:
            self.var_model_, self.var_cov_ = var.fit(data=self.activations_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
            self.connectivity_ = Connectivity(self.var_model_, self.var_cov_, self.nfft_)
        else:
            self.var_model_, self.var_cov_ = var.fit_multiclass(data=self.activations_, cl=self.cl_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
            self.connectivity_ = {}
            for c in np.unique(self.cl_):
                self.connectivity_[c] = Connectivity(self.var_model_[c], self.var_cov_[c], self.nfft_)

    def getConnectivity(self, measure):
        if self.connectivity_ == None:
            raise RuntimeError("Connectivity requires a VAR model (run doMVARICA or fitVAR first)")
        if isinstance(self.connectivity_, dict):
            result = {}
            for c in np.unique(self.cl_):
                result[c] = getattr(self.connectivity_[c], measure)()
            return result
        else:
            return getattr(self.connectivity_, measure)()

    def getTFConnectivity(self, measure, winlen, winstep):
        if self.activations_ == None:
            raise RuntimeError("Time/Frequency Connectivity requires activations (call setData after doMVARICA)")
        [N,M,T] = self.activations_.shape

        Nstep = (N-winlen)//winstep

        if self.cl_ == None:
            result = np.zeros((M, M, Nstep, self.nfft_), np.complex64)
            i = 0
            for n in range(0, N-winlen, winstep):
                win = np.arange(winlen) + n
                data = self.activations_[win,:,:]
                B, C = var.fit(data, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
                con = Connectivity(B, C, self.nfft_)
                result[:,:,i,:] = getattr(con, measure)()
                i += 1

        else:
            result = {}
            for c in np.unique(self.cl_):
                result[c] = np.zeros((M, M, Nstep, self.nfft_), np.complex64)
            i = 0
            for n in range(0, N-winlen, winstep):
                win = np.arange(winlen) + n
                data = self.activations_[win,:,:]
                B, C = var.fit_multiclass(data, cl=self.cl_, P=self.var_order_, delta=self.var_delta_, return_covariance=True)
                for c in np.unique(self.cl_):
                    con = Connectivity(B[c], C[c], self.nfft_)
                    result[c][:,:,i,:] = getattr(con, measure)()
                i += 1
        return result

    def preparePlots(self, mixing=False, unmixing=False):
        if self.locations_ == None:
            raise RuntimeError("Need sensor locations for plotting")

        if self.topo_ == None:
            self.topo_ = Topoplot( )
            self.topo_.set_locations(self.locations_)

        if mixing and not self.mixmaps_:
            for i in range(self.mixing_.shape[0]):
                self.topo_.set_values(self.mixing_[i,:])
                self.topo_.create_map()
                self.mixmaps_.append(self.topo_.get_map())

        if unmixing and not self.unmixmaps_:
            for i in range(self.unmixing_.shape[1]):
                self.topo_.set_values(self.unmixing_[:,i])
                self.topo_.create_map()
                self.unmixmaps_.append(self.topo_.get_map())

    def showPlots(self):
        plt.show()

    def plotComponents(self, global_scale=None):
        """ global_scale:
               None - scales each topo individually
               1-99 - percentile of maximum of all plots
        """
        if not _have_pyplot:
            raise ImportError("matplotlib.pyplot is required for plotting")
        if self.unmixing_ == None and self.mixing_ == None:
            raise RuntimeError("No components available (run doMVARICA first)")
        self.preparePlots(True, True)

        M = self.mixing_.shape[0]

        urange, mrange = None, None

        if global_scale:
            tmp = np.asarray(self.unmixmaps_)
            tmp = tmp[np.logical_not(np.isnan(tmp))]
            umax = np.percentile(np.abs(tmp), global_scale)
            umin = -umax
            urange = [umin,umax]

            tmp = np.asarray(self.mixmaps_)
            tmp = tmp[np.logical_not(np.isnan(tmp))]
            mmax = np.percentile(np.abs(tmp), global_scale)
            mmin = -mmax
            mrange = [umin,umax]

        Y = np.floor(np.sqrt(M*3/4))
        X = np.ceil(M/Y)

        fig = plt.figure()

        axes = []
        for m in range(M):
            axes.append(fig.add_subplot(2*Y, X, m+1))
            h1 = self._plotUnmixing(axes[-1], m, crange=urange)

            axes.append(fig.add_subplot(2*Y, X, M+m+1))
            h1 = self._plotMixing(axes[-1], m, crange=mrange)

        for a in axes:
            a.set_yticks([])
            a.set_xticks([])
            a.set_frame_on(False)

        axes[0].set_ylabel('Unmixing weights')
        axes[1].set_ylabel('Scalp projections')

        #plt.colorbar(h1, plt.subplot(2, M+1, M+1))
        #plt.colorbar(h2, plt.subplot(2, M+1, 0))

    def plotConnectivity(self, measure):
        if not _have_pyplot:
            raise ImportError("matplotlib.pyplot is required for plotting")
        self.preparePlots(True, False)
        fig = plt.figure()
        if isinstance(self.connectivity_, dict):
            for c in np.unique(self.cl_):
                cm = getattr(self.connectivity_[c], measure)()
                self._plotSpectral(fig, cm)
        else:
            cm = getattr(self.connectivity_, measure)()
            self._plotSpectral(fig, cm)

    def _plotSpectral(self, fig, A):
        [N,M,F] = A.shape
        freq = np.linspace(0, self.fs_/2, F)

        lowest, highest = np.inf, -np.inf

        axes = []
        for n in range(N):
            arow = []
            for m in range(M):
                ax = fig.add_subplot(N, M, m+n*M+1)
                arow.append(ax)

                if n == m:
                    self._plotMixing(ax, m)
                    ax.set_yticks([])
                    ax.set_xticks([])
                    ax.set_frame_on(False)
                else:
                    ax.plot(freq, A[n,m,:])
                    lowest = min(lowest, np.min(A[n,m,:]))
                    highest = max(highest, np.max(A[n,m,:]))
                    ax.set_xlim(0, self.fs_/2)
            axes.append(arow)

        for n in range(N):
            for m in range(M):
                if n == m:
                    pass
                else:
                    axes[n][m].set_ylim(lowest, highest)
                    if 0 < n < N-1:
                        axes[n][m].set_xticks([])
                    if 0 < m < M-1:
                        axes[n][m].set_yticks([])
            axes[n][0].yaxis.tick_left()
            axes[n][-1].yaxis.tick_right()

        for m in range(M):
            axes[0][m].xaxis.tick_top()
            axes[-1][m].xaxis.tick_bottom()

        fig.text(0.5, 0.05, 'frequency', horizontalalignment='center')
        fig.text(0.05, 0.5, 'magnitude', horizontalalignment='center', rotation='vertical')

    def _plotMixing(self, axis, idx, crange=None):
        self.topo_.set_map(self.mixmaps_[idx])
        h = self.topo_.plot_map(axis, crange=crange)
        self.topo_.plot_locations(axis)
        self.topo_.plot_head(axis)
        return h

    def _plotUnmixing(self, axis, idx, crange=None):
        self.topo_.set_map(self.unmixmaps_[idx])
        h = self.topo_.plot_map(axis, crange=crange)
        self.topo_.plot_locations(axis)
        self.topo_.plot_head(axis)
        return h<|endoftext|>


<commit_before>#!/ms/dist/python/PROJ/core/2.5.2-1/bin/python
""" to refresh network data from dsdb """

import os
import sys
import logging
import optparse

DIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.realpath(os.path.join(DIR, '..', '..', '..')))
import aquilon.aqdb.depends

from sqlalchemy.exceptions      import DatabaseError, IntegrityError

#TODO: fix all the damn location imports with a shortcut in loc/__init__
from aquilon.aqdb.loc.hub       import Hub
from aquilon.aqdb.loc.city      import City
from aquilon.aqdb.loc.company   import Company
from aquilon.aqdb.loc.country   import Country
from aquilon.aqdb.loc.building  import Building
from aquilon.aqdb.loc.continent import Continent
from aquilon.aqdb.loc.location  import Location

#TODO: likewise, all the imports for system subtypes
from aquilon.aqdb.sy.system     import System
from aquilon.aqdb.sy.tor_switch import TorSwitch

from aquilon.aqdb.db_factory    import db_factory
from aquilon.aqdb.dsdb          import DsdbConnection
from aquilon.aqdb.net.network   import Network, _mask_to_cidr, get_bcast
from aquilon.aqdb.data_sync     import NetRecord
from aquilon.aqdb.data_sync     import RefreshReport
from aquilon.aqdb.utils.shutils import ipshell

class NetRefresher(object):
    """ Class to encapsulate what's needed to replicate networks from AQDB
        to AQDB"""
    __shared_state = {}

    #Dependency injection: allows us to supply our own *fake* dsdb connection
    def __init__(self, dsdb_cnxn, aqdb, *args, **kw):
        self.__dict__ = self.__shared_state
        #TODO: accept args for log level, etc

        if not getattr(self, 'dsdb', None):
            self.dsdb = dsdb_cnxn
            assert self.dsdb

        if not getattr(self, 'aqdb', None):
            self.aqdb = aqdb
            assert self.aqdb

        if not getattr(self, 'location', None):
            q = self.aqdb.s.query(Building)
            self.location = q.filter_by(name=kw['bldg']).one()
            assert type(self.location) is Building

        if not getattr(self, 'report', None):
            self.report = RefreshReport()
            assert self.report

        if not getattr(self, 'log', None):
            self.log = kw['log']
            assert self.log

        self.commit    = kw['commit']

    def _pull_dsdb_data(self, *args, **kw):
        """ loc argument is a sysloc string (DSDB stores this instead) """
        d = {}
        #query returns in order name, ip, mask, type_id, bldg, side
        for (name, ip, mask, type,
             bldg, side, dsdb_id) in self.dsdb.get_network_by_sysloc(
                                                    self.location.sysloc()):
            d[ip] = NetRecord(ip=ip, name=name, net_type=type,
                              mask=mask, bldg=bldg, side=side, dsdb_id=dsdb_id)
        return d

    def _pull_aqdb_data(self, *args, **kw):
        """ loc argument is a BUILDING CODE, different from dsdb """
        d = {}

        q = self.aqdb.s.query(Network)
        for n in q.filter_by(location=self.location).all():
            d[n.ip] = n
        return d

    def _rollback(self, msg):
        """ a DRY method for error handling """
        self.log.error(msg)
        self.report.errs.append(msg)
        self.aqdb.s.rollback()

    def refresh(self, *args, **kw):
        """ compares and refreshes the network table from dsdb to aqdb using
            sets makes computing union and delta of keys simple and succinct """

        ds = self._pull_dsdb_data(*args, **kw)
        dset = set(ds.keys())

        aq = self._pull_aqdb_data(*args, **kw)
        aset = set(aq.keys())

        deletes  = aset - dset
        if deletes:
            self._do_deletes(deletes, aq, *args, **kw)
            aq = self._pull_aqdb_data(*args, **kw)
            aset = set(aq.keys())

        adds = dset - aset
        if adds:
            self._do_adds(adds, ds, *args, **kw)
            aq = self._pull_aqdb_data(*args, **kw)
            aset = set(aq.keys())

        compares = aset & dset
        if compares:
            self._do_updates(aset & dset , aq, ds, *args, **kw)

    def _do_deletes(self, k, aq, *args, **kw):
        """ Deletes networks in aqdb and not in dsdb. It logs and handles
            associated reporting messages.

            arguments: list of keys and dict of aqdb networks to delete
            returns:   True/False on total success/a single failure """
            #TODO: revisit this ^^^
        for i in k:
            self.log.debug('deleting %s'%(aq[i]))
            self.report.dels.append(aq[i])

            if self.commit:
                try:
                    self.aqdb.s.delete(aq[i])
                    self.aqdb.s.commit()
                    self.aqdb.s.flush()
                except IntegrityError,e:
#                    """ TODO: get records that connects back to the network and
#                            do something about it: delete it, or send message.
#                            enhancement: start marking objects for deletion and
#                            delete them after a week or two of no activity """
                    self._rollback(e)
                except Exception, e:
                    self._rollback(e)
                    continue

    def _do_adds(self, k, ds, *args, **kw):
        """ Adds networks in dsdb and not in aqdb. Handles the logging and
            job reporting along the way.

            arguments: list of keys (ip addresses) and dsdb NetRecords

            returns:   # of networks added or False based on total success or
                       a single failure """ #TODO: revisit this
        for i in k:
            try:
                c = _mask_to_cidr[ds[i].mask]
                net = Network(name         = ds[i].name,
                              ip           = ds[i].ip,
                              mask         = ds[i].mask,
                              cidr         = c,
                              bcast        = get_bcast(ds[i].ip, c),
                              network_type = ds[i].net_type,
                              side         = ds[i].side,
                              location     = self.location,
                              dsdb_id      = ds[i].dsdb_id)
                net.comments = getattr(ds[i], 'comments', None)
    #TODO: use a memoized query:
    #self.aqdb.s.query(Building).filter_by(name=ds[i].bldg).one()
            except Exception, e:
                self.report.errs.append(e)
                self.log.error(e)
                continue

            #log here to get more detailed/uniform output
            self.log.debug('adding %s'%(net))
            self.report.adds.append(net)

            if self.commit:
                self.report.adds.append(net)
                try:
                    self.aqdb.s.add(net)
                    self.aqdb.s.commit()
                    self.aqdb.s.flush()
                except Exception, e:
                    self._rollback(e)

    def _do_updates(self, k, aq, ds, *args, **kw):
        """ Makes changes to networks which have differences, logs/reports. We
            do this *VERY* cautiously with seperate try/except for everything

            arguments: list of keys (ip addresses) for networks present in both
                       data sources, plus a hash of the relevent data to compare

            returns:   # of networks successfully updated or False if there's a
                       single failure  """
        for i in k:
            if ds[i] != aq[i]:
                #get an updated version of the aqdb network
                try:
                    aq[i] = ds[i].update_aq_net(aq[i], self.log, self.report)
                except ValueError, e:
                    self.report.errs.append(e)
                    self.log.error(e)

                if self.commit:
                    try:
                        self.log.debug('trying to commit the update\n')
                        self.aqdb.s.update(aq[i])
                        self.aqdb.s.commit()
                        self.aqdb.s.flush()
                        self.report.upds.append(ds[i].name)
                    except Exception, e:
                        self._rollback(e)

def main(*args, **kw):
    usage = """ usage: %prog [options]
    refreshes location data from dsdb to aqdb """

    desc = 'Synchronizes networks from DSDB to AQDB'

    p = optparse.OptionParser(usage=usage, prog=sys.argv[0], version='0.1',
                              description=desc)

    p.add_option('-v',
                 action = 'count',
                 dest   = 'verbose',
                 default = 0,
                 help   = 'increase verbosity by adding more (vv), etc.')

    p.add_option('-b', '--building',
                 action = 'store',
                 dest   = 'building',
                 default = 'dd' )

    p.add_option('-n', '--dry-run',
                      action  = 'store_true',
                      dest    = 'dry_run',
                      default = False,
                      help    = "no commit (for testing, default=False)")
    opts, args = p.parse_args()

    dsdb = DsdbConnection()
    aqdb = db_factory()

    log = logging.getLogger('net_refresh')

    if opts.verbose > 1:
        log_level = logging.DEBUG
    elif opts.verbose > 0:
        log_level = logging.INFO
    else:
        log_level = logging.WARN

    #TODO: call this in main
    logging.basicConfig(level=log_level,
                    format='%(asctime)s %(levelname)-6s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S')
    assert log

    nr = NetRefresher(dsdb, aqdb,
                      log=log,
                      bldg=opts.building,
                      #keep logic stated in postive language (readability)
                      commit = not(opts.dry_run))
    nr.refresh()

    if opts.verbose < 2:
        #really verbose runs don't need reporting
        nr.report.display()

if __name__ == '__main__':
    main(sys.argv)

# Copyright (C) 2008 Morgan Stanley
# This module is part of Aquilon

# ex: set expandtab softtabstop=4 shiftwidth=4: -*- cpy-indent-level: 4; indent-tabs-mode: nil -*-
<commit_msg>Fix session and log handling in net_refresh class.<commit_after>#!/ms/dist/python/PROJ/core/2.5.2-1/bin/python
""" to refresh network data from dsdb """

import os
import sys
import logging
import optparse

DIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.realpath(os.path.join(DIR, '..', '..', '..')))
import aquilon.aqdb.depends

from sqlalchemy.exceptions      import DatabaseError, IntegrityError

#TODO: fix all the damn location imports with a shortcut in loc/__init__
from aquilon.aqdb.loc.hub       import Hub
from aquilon.aqdb.loc.city      import City
from aquilon.aqdb.loc.company   import Company
from aquilon.aqdb.loc.country   import Country
from aquilon.aqdb.loc.building  import Building
from aquilon.aqdb.loc.continent import Continent
from aquilon.aqdb.loc.location  import Location

#TODO: likewise, all the imports for system subtypes
from aquilon.aqdb.sy.system     import System
from aquilon.aqdb.sy.tor_switch import TorSwitch

from aquilon.aqdb.db_factory    import db_factory
from aquilon.aqdb.dsdb          import DsdbConnection
from aquilon.aqdb.net.network   import Network, _mask_to_cidr, get_bcast
from aquilon.aqdb.data_sync     import NetRecord
from aquilon.aqdb.data_sync     import RefreshReport
from aquilon.aqdb.utils.shutils import ipshell

class NetRefresher(object):
    """ Class to encapsulate what's needed to replicate networks from AQDB
        to AQDB"""
    __shared_state = {}

    #Dependency injection: allows us to supply our own *fake* dsdb connection
    def __init__(self, dsdb_cnxn, session, *args, **kw):

        self.dsdb = dsdb_cnxn
        assert self.dsdb

        self.session = session
        assert self.session

        q = self.session.query(Building)
        self.location = q.filter_by(name=kw['bldg']).one()
        assert type(self.location) is Building

        self.report = RefreshReport()
        assert self.report

        self.log = logging.getLogger('net_refresh')
        assert self.log

        self.commit    = kw['commit']

    def _pull_dsdb_data(self, *args, **kw):
        """ loc argument is a sysloc string (DSDB stores this instead) """
        d = {}
        #query returns in order name, ip, mask, type_id, bldg, side
        for (name, ip, mask, type,
             bldg, side, dsdb_id) in self.dsdb.get_network_by_sysloc(
                                                    self.location.sysloc()):
            d[ip] = NetRecord(ip=ip, name=name, net_type=type,
                              mask=mask, bldg=bldg, side=side, dsdb_id=dsdb_id)
        return d

    def _pull_aqdb_data(self, *args, **kw):
        """ loc argument is a BUILDING CODE, different from dsdb """
        d = {}

        q = self.session.query(Network)
        for n in q.filter_by(location=self.location).all():
            d[n.ip] = n
        return d

    def _rollback(self, msg):
        """ a DRY method for error handling """
        self.log.error(msg)
        self.report.errs.append(msg)
        self.session.rollback()

    def refresh(self, *args, **kw):
        """ compares and refreshes the network table from dsdb to aqdb using
            sets makes computing union and delta of keys simple and succinct """

        self.log.debug('Starting network refresh')

        ds = self._pull_dsdb_data(*args, **kw)
        dset = set(ds.keys())

        aq = self._pull_aqdb_data(*args, **kw)
        aset = set(aq.keys())

        deletes  = aset - dset
        if deletes:
            self._do_deletes(deletes, aq, *args, **kw)
            aq = self._pull_aqdb_data(*args, **kw)
            aset = set(aq.keys())

        adds = dset - aset
        if adds:
            self._do_adds(adds, ds, *args, **kw)
            aq = self._pull_aqdb_data(*args, **kw)
            aset = set(aq.keys())

        compares = aset & dset
        if compares:
            self._do_updates(aset & dset , aq, ds, *args, **kw)

        self.log.debug('Finished network refresh logic')

    def _do_deletes(self, k, aq, *args, **kw):
        """ Deletes networks in aqdb and not in dsdb. It logs and handles
            associated reporting messages.

            arguments: list of keys and dict of aqdb networks to delete
            returns:   True/False on total success/a single failure """
            #TODO: revisit this ^^^
        for i in k:
            self.log.debug('deleting %s'%(aq[i]))
            self.report.dels.append(aq[i])

            if self.commit:
                try:
                    self.session.delete(aq[i])
                    self.session.commit()
                except IntegrityError,e:
#                    """ TODO: get records that connects back to the network and
#                            do something about it: delete it, or send message.
#                            enhancement: start marking objects for deletion and
#                            delete them after a week or two of no activity """
                    self._rollback(e)
                except Exception, e:
                    self._rollback(e)
                    continue

    def _do_adds(self, k, ds, *args, **kw):
        """ Adds networks in dsdb and not in aqdb. Handles the logging and
            job reporting along the way.

            arguments: list of keys (ip addresses) and dsdb NetRecords

            returns:   # of networks added or False based on total success or
                       a single failure """ #TODO: revisit this
        for i in k:
            try:
                c = _mask_to_cidr[ds[i].mask]
                net = Network(name         = ds[i].name,
                              ip           = ds[i].ip,
                              mask         = ds[i].mask,
                              cidr         = c,
                              bcast        = get_bcast(ds[i].ip, c),
                              network_type = ds[i].net_type,
                              side         = ds[i].side,
                              location     = self.location,
                              dsdb_id      = ds[i].dsdb_id)
                net.comments = getattr(ds[i], 'comments', None)
    #TODO: use a memoized query:
    #self.session.query(Building).filter_by(name=ds[i].bldg).one()
            except Exception, e:
                self.report.errs.append(e)
                self.log.error(e)
                continue

            #log here to get more detailed/uniform output
            self.log.debug('adding %s'%(net))
            self.report.adds.append(net)

            if self.commit:
                self.report.adds.append(net)
                try:
                    self.session.add(net)
                    self.session.commit()
                except Exception, e:
                    self._rollback(e)

    def _do_updates(self, k, aq, ds, *args, **kw):
        """ Makes changes to networks which have differences, logs/reports. We
            do this *VERY* cautiously with seperate try/except for everything

            arguments: list of keys (ip addresses) for networks present in both
                       data sources, plus a hash of the relevent data to compare

            returns:   # of networks successfully updated or False if there's a
                       single failure  """
        for i in k:
            if ds[i] != aq[i]:
                #get an updated version of the aqdb network
                try:
                    aq[i] = ds[i].update_aq_net(aq[i], self.log, self.report)
                except ValueError, e:
                    self.report.errs.append(e)
                    self.log.error(e)

                if self.commit:
                    try:
                        self.log.debug('trying to commit the update\n')
                        self.session.update(aq[i])
                        self.session.commit()
                        self.report.upds.append(ds[i].name)
                    except Exception, e:
                        self._rollback(e)

def main(*args, **kw):
    usage = """ usage: %prog [options]
    refreshes location data from dsdb to aqdb """

    desc = 'Synchronizes networks from DSDB to AQDB'

    p = optparse.OptionParser(usage=usage, prog=sys.argv[0], version='0.1',
                              description=desc)

    p.add_option('-v',
                 action = 'count',
                 dest   = 'verbose',
                 default = 0,
                 help   = 'increase verbosity by adding more (vv), etc.')

    p.add_option('-b', '--building',
                 action = 'store',
                 dest   = 'building',
                 default = 'dd' )

    p.add_option('-n', '--dry-run',
                      action  = 'store_true',
                      dest    = 'dry_run',
                      default = False,
                      help    = "no commit (for testing, default=False)")
    opts, args = p.parse_args()

    dsdb = DsdbConnection()
    aqdb = db_factory()

    if opts.verbose > 1:
        log_level = logging.DEBUG
    elif opts.verbose > 0:
        log_level = logging.INFO
    else:
        log_level = logging.WARN

    logging.basicConfig(level=log_level,
                    format='%(asctime)s %(levelname)-6s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S')

    nr = NetRefresher(dsdb, aqdb.session(),
                      bldg=opts.building,
                      #keep logic stated in postive language (readability)
                      commit = not(opts.dry_run))
    nr.refresh()

    if opts.verbose < 2:
        #really verbose runs don't need reporting
        nr.report.display()

if __name__ == '__main__':
    main(sys.argv)

# Copyright (C) 2008 Morgan Stanley
# This module is part of Aquilon

# ex: set expandtab softtabstop=4 shiftwidth=4: -*- cpy-indent-level: 4; indent-tabs-mode: nil -*-
<|endoftext|>