<commit_before>from fabric.api import task, local, cd, lcd, settings, run, sudo, put, get, abort
from fabric.contrib.files import exists
from fabric.contrib import files
from fabric.contrib.console import confirm
import aws, utils
from decorators import requires_project, requires_aws_stack, echo_output, deffile, setdefault, debugtask, timeit
import os
from os.path import join
from functools import wraps
from slugify import slugify
import buildercore
from buildercore import config, core, cfngen, utils as core_utils, bootstrap, bakery, project
from buildercore.core import stack_conn, stack_pem
from buildercore.utils import first
from buildercore.config import ROOT_USER, DEPLOY_USER, BOOTSTRAP_USER
from buildercore.decorators import osissue, osissuefn

import logging

LOG = logging.getLogger(__name__)

@debugtask
@echo_output
def stack_files(project=None):
    "returns a list of CloudFormation TEMPLATE FILES. accepts optional project name"
    stacks = sorted(core.stack_files())
    if project:
        return filter(lambda stack: stack.startswith("%s-" % project), stacks)
    return stacks

def requires_stack_file(func):
    "test that the stack exists in the STACKS dir"
    @wraps(func)
    def _wrapper(stackname=None, *args, **kwargs):
        flist = stack_files()
        if not flist:
            print 'no stack files exist!'
            return
        if not stackname or stackname not in flist:
            stackname = utils._pick("stack", flist, default_file=deffile('.stack'))
        return func(stackname, *args, **kwargs)
    return _wrapper

#
# tasks
#

@task
def project_list():
    for org, plist in project.org_project_map().items():
        print org
        for p in plist:
            print '  ',p
        print

@debugtask
@echo_output
def aws_detailed_stack_list(project=None):
    region = aws.find_region()
    all_stacks = dict([(i.stack_name, i.__dict__) for i in core.raw_aws_stacks(region)])
    if project:
        return {k: v for k, v in all_stacks.items() if k.startswith("%s-" % project)}
    return all_stacks

@debugtask
@requires_stack_file
def aws_stack_exists(stackname):
    "we may know about the stack on disk, but it might not have been pushed to aws yet..."
    pdata = core.project_data_for_stackname(stackname)
    print pdata
    region = pdata['aws']['region']
    return stackname in core.all_aws_stack_names(region)

@debugtask
@requires_stack_file # @requires_inactive_stack
def delete_stack_file(stackname):
    logging.getLogger('boto').setLevel(logging.CRITICAL)
    files_removed = bootstrap.delete_stack_file
    return files_removed

@task
@requires_aws_stack
def delete_stack(stackname, confirmed):
    try:
        # we want 'confirmed' to equal the boolean True
        #pylint: disable=singleton-comparison
        if confirmed == True:
            return bootstrap.delete_stack(stackname)
    except ValueError:
        # it's possible to delete a stack and for the json stack file to be missing.
        # how? nfi. maybe legacy stuff
        LOG.exception("attempting to delete stack %r and the json template is missing", stackname)

@task
@requires_aws_stack
def aws_delete_stack(stackname):
    "tells aws to delete a stack. this doesn't delete the CloudFormation file from the stacks dir"
    print 'this is a BIG DEAL. you cannot recover from this.'
    print 'type the name of the stack to continue or anything else to quit'
    uin = raw_input('> ')
    if not uin or not uin.strip().lower() == stackname.lower():
        import difflib
        print 'you needed to type "%s" to continue.' % stackname
        print 'got:'
        print '\n'.join(difflib.ndiff([stackname], [uin]))
        exit(1)
    return delete_stack(stackname, confirmed=True)

@task
@requires_aws_stack
def aws_update_stack(stackname):
    """Updates the master and then updates the stack software.
    Only run update commands if salt has successfully been installed.
    Update commands require that the salt deploy user `DEPLOY_USER` exists."""
    return bootstrap.update_stack(stackname)

@task
def update_master():
    return bootstrap.update_stack(core.find_master(aws.find_region()))

@debugtask
@requires_aws_stack
def highstate(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call saltutil.refresh_pillar') # not sure if this even does anything ...
        sudo('salt-call state.highstate')

@debugtask
@requires_aws_stack
def refresh_pillar(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call saltutil.refresh_pillar')

@debugtask
@requires_aws_stack
def pillar(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call pillar.items')

@debugtask
@requires_aws_stack
def aws_update_template(stackname):
    "updates the CloudFormation stack and then updates the environment"
    return bootstrap.update_template(stackname)

@debugtask
@requires_stack_file
def aws_create_update_stack(stackname):
    if not core.stack_is_active(stackname):
        print 'stack does not exist, creating'
        bootstrap.create_stack(stackname)
    print 'updating stack'
    bootstrap.update_stack(stackname)
    return stackname

@task
@echo_output
def aws_stack_list():
    "returns a list of realized stacks. does not include deleted stacks"
    region = aws.find_region()
    return core.all_aws_stack_names(region)

@task
@requires_aws_stack
def ssh(stackname, username=DEPLOY_USER):
    public_ip = core.stack_data(stackname)['instance']['ip_address']
    # -A forwarding of authentication agent connection
    local("ssh %s@%s -A" % (username, public_ip))

@task
@requires_aws_stack
def cmd(stackname, command=None):
    if command is None:
        abort("Please specify a command e.g. ./bldr cmd:%s,ls" % stackname)
    with stack_conn(stackname):
	with settings(abort_on_prompts=True):
            run(command)

@task
@requires_aws_stack
def owner_ssh(stackname):
    "maintainence ssh. uses the pem key and the bootstrap user to login."
    public_ip = core.stack_data(stackname)['instance']['ip_address']
    # -i identify file
    # -A forwarding of authentication agent connection
    local("ssh %s@%s -i %s -A" % (BOOTSTRAP_USER, public_ip, stack_pem(stackname)))

#
# local template management
#

@task
@requires_project
@echo_output
def create_stack(pname):
    """creates a new CloudFormation template for the given project."""
    default_instance_id, cluster_id = core_utils.ymd(), None
    inst_id = utils.uin("instance id", default_instance_id)
    stackname = core.mk_stackname(pname, inst_id, cluster_id)
    more_context = {'instance_id': stackname}

    # prompt user for alternate configurations
    pdata = project.project_data(pname)
    if pdata.has_key('aws-alt'):
        def helpfn(altkey):
            try:
                return pdata['aws-alt'][altkey]['description']
            except KeyError:
                return None
        default = 'skip this step'
        alt_config = [default] + pdata['aws-alt'].keys()
        alt_config = utils._pick('alternative config', alt_config, helpfn=helpfn)
        if alt_config != default:
            more_context['alt-config'] = alt_config

    _, out_fname = cfngen.generate_stack(pname, **more_context)

    print
    print 'CloudFormation template written to:', out_fname
    print 'use `fab cfn.aws_create_update_stack` next'
    print

    return stackname

'''
@debugtask
@requires_project
def print_stack_template(project):
    default_instance_id = core_utils.ymd()
    more_context = dict([
        ('instance_id', slugify(project + "-" + default_instance_id)),
    ])
    context = cfngen.build_context(project, config.PROJECT_FILE, config.PILLAR_DIR, **more_context)
    print cfngen.render_template(context)
'''

@task
@requires_project
@echo_output
def print_project_config(pname):
    return core_utils.remove_ordereddict(project.project_data(pname))

#
#
#

@task
@requires_project
def aws_launch_instance(project):
    try:
        stackname = create_stack(project)
        pdata = core.project_data_for_stackname(stackname)

        print 'attempting to create stack:'
        print '  stackname: ' + stackname
        print '  region:    ' + pdata['aws']['region']
        print '  vpc:       ' + pdata['aws']['vpc-id']
        print '  subnet:    ' + pdata['aws']['subnet-id']
        print
        if not confirm('continue?', default=True):
            exit()

        stackname = aws_create_update_stack(stackname)

        if stackname.startswith('master-server--'):
            print
            print "`master-server` projects must create a deploy key in it's `formula-repo` project"
            print
            print

        if stackname:
            setdefault('.active-stack', stackname)
    except core.NoMasterException, e:
        LOG.warn(e.message)
        print "\n%s\ntry `./bldr master.create`'" % e.message

@task
@requires_aws_stack
def download_file(stackname, path, destination):
    fname = os.path.basename(path)
    utils.mkdirp(destination)
    with stack_conn(stackname):
        get(path, destination, use_sudo=True)

@task
@requires_aws_stack
def upload_file(stackname, local_path, remote_path, overwrite=False):
    with stack_conn(stackname):
        if files.exists(remote_path) and not overwrite:
            print 'remote file exists, not overwriting'
            exit(1)
        put(local_path, remote_path)

@task
@requires_aws_stack
def create_ami(stackname):
    pname = core.project_name_from_stackname(stackname)
    msg = "this will create a new AMI for the project %r. Continue?" % pname
    if not confirm(msg, default=False):
        print 'doing nothing'
        return
    amiid = bakery.create_ami(stackname)
    #amiid = "ami-e9ff3682"
    print 'AWS is now creating AMI with id', amiid
    path = pname + '.aws.ami'
    # wait until ami finished creating?
    #core.update_project_file(pname + ".aws.ami", amiid)
    new_project_file = project.update_project_file(path, amiid)
    output_file = project.write_project_file(new_project_file)
    print '\n' * 4
    print 'wrote', output_file
    print 'updated project file with new ami. these changes must be merged and committed manually'
    print '\n' * 4

#
# rds tests
#

@debugtask
@requires_aws_stack
@echo_output
def aws_rds_snapshots(stackname):
    from boto import rds
    conn = rds.RDSConnection()
    instance = conn.get_all_dbinstances(instance_id=stackname)[0]
    # all snapshots order by creation time
    objdata = conn.get_all_dbsnapshots(instance_id=instance.id)
    data = sorted(map(lambda ss: ss.__dict__, objdata), key=lambda i: i['snapshot_create_time'])
    return data
<commit_msg>replaced a tab with spaces<commit_after>from fabric.api import task, local, cd, lcd, settings, run, sudo, put, get, abort
from fabric.contrib.files import exists
from fabric.contrib import files
from fabric.contrib.console import confirm
import aws, utils
from decorators import requires_project, requires_aws_stack, echo_output, deffile, setdefault, debugtask, timeit
import os
from os.path import join
from functools import wraps
from slugify import slugify
import buildercore
from buildercore import config, core, cfngen, utils as core_utils, bootstrap, bakery, project
from buildercore.core import stack_conn, stack_pem
from buildercore.utils import first
from buildercore.config import ROOT_USER, DEPLOY_USER, BOOTSTRAP_USER
from buildercore.decorators import osissue, osissuefn

import logging

LOG = logging.getLogger(__name__)

@debugtask
@echo_output
def stack_files(project=None):
    "returns a list of CloudFormation TEMPLATE FILES. accepts optional project name"
    stacks = sorted(core.stack_files())
    if project:
        return filter(lambda stack: stack.startswith("%s-" % project), stacks)
    return stacks

def requires_stack_file(func):
    "test that the stack exists in the STACKS dir"
    @wraps(func)
    def _wrapper(stackname=None, *args, **kwargs):
        flist = stack_files()
        if not flist:
            print 'no stack files exist!'
            return
        if not stackname or stackname not in flist:
            stackname = utils._pick("stack", flist, default_file=deffile('.stack'))
        return func(stackname, *args, **kwargs)
    return _wrapper

#
# tasks
#

@task
def project_list():
    for org, plist in project.org_project_map().items():
        print org
        for p in plist:
            print '  ',p
        print

@debugtask
@echo_output
def aws_detailed_stack_list(project=None):
    region = aws.find_region()
    all_stacks = dict([(i.stack_name, i.__dict__) for i in core.raw_aws_stacks(region)])
    if project:
        return {k: v for k, v in all_stacks.items() if k.startswith("%s-" % project)}
    return all_stacks

@debugtask
@requires_stack_file
def aws_stack_exists(stackname):
    "we may know about the stack on disk, but it might not have been pushed to aws yet..."
    pdata = core.project_data_for_stackname(stackname)
    print pdata
    region = pdata['aws']['region']
    return stackname in core.all_aws_stack_names(region)

@debugtask
@requires_stack_file # @requires_inactive_stack
def delete_stack_file(stackname):
    logging.getLogger('boto').setLevel(logging.CRITICAL)
    files_removed = bootstrap.delete_stack_file
    return files_removed

@task
@requires_aws_stack
def delete_stack(stackname, confirmed):
    try:
        # we want 'confirmed' to equal the boolean True
        #pylint: disable=singleton-comparison
        if confirmed == True:
            return bootstrap.delete_stack(stackname)
    except ValueError:
        # it's possible to delete a stack and for the json stack file to be missing.
        # how? nfi. maybe legacy stuff
        LOG.exception("attempting to delete stack %r and the json template is missing", stackname)

@task
@requires_aws_stack
def aws_delete_stack(stackname):
    "tells aws to delete a stack. this doesn't delete the CloudFormation file from the stacks dir"
    print 'this is a BIG DEAL. you cannot recover from this.'
    print 'type the name of the stack to continue or anything else to quit'
    uin = raw_input('> ')
    if not uin or not uin.strip().lower() == stackname.lower():
        import difflib
        print 'you needed to type "%s" to continue.' % stackname
        print 'got:'
        print '\n'.join(difflib.ndiff([stackname], [uin]))
        exit(1)
    return delete_stack(stackname, confirmed=True)

@task
@requires_aws_stack
def aws_update_stack(stackname):
    """Updates the master and then updates the stack software.
    Only run update commands if salt has successfully been installed.
    Update commands require that the salt deploy user `DEPLOY_USER` exists."""
    return bootstrap.update_stack(stackname)

@task
def update_master():
    return bootstrap.update_stack(core.find_master(aws.find_region()))

@debugtask
@requires_aws_stack
def highstate(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call saltutil.refresh_pillar') # not sure if this even does anything ...
        sudo('salt-call state.highstate')

@debugtask
@requires_aws_stack
def refresh_pillar(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call saltutil.refresh_pillar')

@debugtask
@requires_aws_stack
def pillar(stackname):
    with stack_conn(stackname, username=BOOTSTRAP_USER):
        sudo('salt-call pillar.items')

@debugtask
@requires_aws_stack
def aws_update_template(stackname):
    "updates the CloudFormation stack and then updates the environment"
    return bootstrap.update_template(stackname)

@debugtask
@requires_stack_file
def aws_create_update_stack(stackname):
    if not core.stack_is_active(stackname):
        print 'stack does not exist, creating'
        bootstrap.create_stack(stackname)
    print 'updating stack'
    bootstrap.update_stack(stackname)
    return stackname

@task
@echo_output
def aws_stack_list():
    "returns a list of realized stacks. does not include deleted stacks"
    region = aws.find_region()
    return core.all_aws_stack_names(region)

@task
@requires_aws_stack
def ssh(stackname, username=DEPLOY_USER):
    public_ip = core.stack_data(stackname)['instance']['ip_address']
    # -A forwarding of authentication agent connection
    local("ssh %s@%s -A" % (username, public_ip))

@task
@requires_aws_stack
def cmd(stackname, command=None):
    if command is None:
        abort("Please specify a command e.g. ./bldr cmd:%s,ls" % stackname)
    with stack_conn(stackname):
        with settings(abort_on_prompts=True):
            run(command)

@task
@requires_aws_stack
def owner_ssh(stackname):
    "maintainence ssh. uses the pem key and the bootstrap user to login."
    public_ip = core.stack_data(stackname)['instance']['ip_address']
    # -i identify file
    # -A forwarding of authentication agent connection
    local("ssh %s@%s -i %s -A" % (BOOTSTRAP_USER, public_ip, stack_pem(stackname)))

#
# local template management
#

@task
@requires_project
@echo_output
def create_stack(pname):
    """creates a new CloudFormation template for the given project."""
    default_instance_id, cluster_id = core_utils.ymd(), None
    inst_id = utils.uin("instance id", default_instance_id)
    stackname = core.mk_stackname(pname, inst_id, cluster_id)
    more_context = {'instance_id': stackname}

    # prompt user for alternate configurations
    pdata = project.project_data(pname)
    if pdata.has_key('aws-alt'):
        def helpfn(altkey):
            try:
                return pdata['aws-alt'][altkey]['description']
            except KeyError:
                return None
        default = 'skip this step'
        alt_config = [default] + pdata['aws-alt'].keys()
        alt_config = utils._pick('alternative config', alt_config, helpfn=helpfn)
        if alt_config != default:
            more_context['alt-config'] = alt_config

    _, out_fname = cfngen.generate_stack(pname, **more_context)

    print
    print 'CloudFormation template written to:', out_fname
    print 'use `fab cfn.aws_create_update_stack` next'
    print

    return stackname

'''
@debugtask
@requires_project
def print_stack_template(project):
    default_instance_id = core_utils.ymd()
    more_context = dict([
        ('instance_id', slugify(project + "-" + default_instance_id)),
    ])
    context = cfngen.build_context(project, config.PROJECT_FILE, config.PILLAR_DIR, **more_context)
    print cfngen.render_template(context)
'''

@task
@requires_project
@echo_output
def print_project_config(pname):
    return core_utils.remove_ordereddict(project.project_data(pname))

#
#
#

@task
@requires_project
def aws_launch_instance(project):
    try:
        stackname = create_stack(project)
        pdata = core.project_data_for_stackname(stackname)

        print 'attempting to create stack:'
        print '  stackname: ' + stackname
        print '  region:    ' + pdata['aws']['region']
        print '  vpc:       ' + pdata['aws']['vpc-id']
        print '  subnet:    ' + pdata['aws']['subnet-id']
        print
        if not confirm('continue?', default=True):
            exit()

        stackname = aws_create_update_stack(stackname)

        if stackname.startswith('master-server--'):
            print
            print "`master-server` projects must create a deploy key in it's `formula-repo` project"
            print
            print

        if stackname:
            setdefault('.active-stack', stackname)
    except core.NoMasterException, e:
        LOG.warn(e.message)
        print "\n%s\ntry `./bldr master.create`'" % e.message

@task
@requires_aws_stack
def download_file(stackname, path, destination):
    fname = os.path.basename(path)
    utils.mkdirp(destination)
    with stack_conn(stackname):
        get(path, destination, use_sudo=True)

@task
@requires_aws_stack
def upload_file(stackname, local_path, remote_path, overwrite=False):
    with stack_conn(stackname):
        if files.exists(remote_path) and not overwrite:
            print 'remote file exists, not overwriting'
            exit(1)
        put(local_path, remote_path)

@task
@requires_aws_stack
def create_ami(stackname):
    pname = core.project_name_from_stackname(stackname)
    msg = "this will create a new AMI for the project %r. Continue?" % pname
    if not confirm(msg, default=False):
        print 'doing nothing'
        return
    amiid = bakery.create_ami(stackname)
    #amiid = "ami-e9ff3682"
    print 'AWS is now creating AMI with id', amiid
    path = pname + '.aws.ami'
    # wait until ami finished creating?
    #core.update_project_file(pname + ".aws.ami", amiid)
    new_project_file = project.update_project_file(path, amiid)
    output_file = project.write_project_file(new_project_file)
    print '\n' * 4
    print 'wrote', output_file
    print 'updated project file with new ami. these changes must be merged and committed manually'
    print '\n' * 4

#
# rds tests
#

@debugtask
@requires_aws_stack
@echo_output
def aws_rds_snapshots(stackname):
    from boto import rds
    conn = rds.RDSConnection()
    instance = conn.get_all_dbinstances(instance_id=stackname)[0]
    # all snapshots order by creation time
    objdata = conn.get_all_dbsnapshots(instance_id=instance.id)
    data = sorted(map(lambda ss: ss.__dict__, objdata), key=lambda i: i['snapshot_create_time'])
    return data
<|endoftext|>


<commit_before># Copyright 2009 Google Inc. All Rights Reserved.
#
# 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.

"""Tk user interface implementation for namebench."""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import os
import sys
import threading
import tkFont
import traceback
import Queue
from Tkinter import *
import tkMessageBox

import base_ui
import history_parser
import nameserver_list
import util

def closedWindowHandler():
  print 'Au revoir, mes amis!'
  sys.exit(1)

global_message_queue = Queue.Queue()
global_last_message = None

def AddMsg(message, master=None, **kwargs):
  """Add a message to the global queue for output."""
  global global_message_queue
  global global_last_message

  new_message = Message(message, **kwargs)
  if new_message != global_last_message:
    global_message_queue.put(new_message)
    if master:
      try:
        master.event_generate('<<msg>>', when='tail')
      except TclError:
        print "TCL error encountered, not pushing update to UI:"
        traceback.print_exc()
    global_last_message = new_message

class Message(object):
  """Messages to be passed from to the main thread from children.

  Used to avoid thread issues inherent with Tk.
  """
  def __init__(self, message, error=False, count=False, total=False,
               enable_button=None, debug=False):
    self.message = message
    self.error = error
    self.count = count
    self.debug = debug
    self.total = total
    self.enable_button = enable_button

class WorkerThread(threading.Thread, base_ui.BaseUI):
  """Handle benchmarking and preparation in a separate UI thread."""

  def __init__(self, primary, secondary, options, history_parser=None, master=None):
    threading.Thread.__init__(self)
    self.status_callback = self.msg
    self.hparser = history_parser
    self.primary = primary
    self.master = master
    self.secondary = secondary
    self.options = options
    self.resource_dir = os.path.dirname(os.path.dirname(__file__))

  def msg(self, message, **kwargs):
    """Add messages to the main queue."""
    return AddMsg(message, master=self.master, **kwargs)

  def run(self):
    self.msg('Started thread', enable_button=False)
    try:
      self.PrepareBenchmark()
      self.UpdateStatus('Here comes the pain...')
      self.RunBenchmark()
    except nameserver_list.OutgoingUdpInterception:
      (exc_type, exception, tb) = sys.exc_info()
      self.msg('Outgoing requests were intercepted!', error=exception)
    except nameserver_list.TooFewNameservers:
      (exc_type, exception, tb) = sys.exc_info()
      self.msg('Too few nameservers to test', error=exception)
    except:
      (exc_type, exception, tb) = sys.exc_info()
      traceback.print_exc(tb)
      error_msg = '\n'.join(traceback.format_tb(tb)[-4:])
      self.msg(exception, error=error_msg)
    self.msg(None, enable_button=True)

class NameBenchGui(object):
  """The main GUI."""
  def __init__(self, options, supplied_ns, global_ns, regional_ns, version=None):
    self.options = options
    self.supplied_ns = supplied_ns
    self.global_ns = global_ns
    self.regional_ns = regional_ns
    self.version = version

  def Execute(self):
    self.root = Tk()
    app = MainWindow(self.root, self.options, self.supplied_ns, self.global_ns, self.regional_ns, self.version)
    app.DrawWindow()
    self.root.bind('<<msg>>', app.MessageHandler)
    self.root.mainloop()

class MainWindow(Frame, base_ui.BaseUI):
  """The main Tk GUI class."""

  def __init__(self, master, options, supplied_ns, global_ns, regional_ns, version=None):
    """TODO(tstromberg): Remove duplication from NameBenchGui class."""
    Frame.__init__(self)
    self.master = master
    self.options = options
    self.supplied_ns = supplied_ns
    self.global_ns = global_ns
    self.regional_ns = regional_ns
    self.version = version
    self.master.protocol('WM_DELETE_WINDOW', closedWindowHandler)

  def DrawWindow(self):
    """Draws the user interface."""
    self.nameserver_form = StringVar()
    self.status = StringVar()
    self.num_tests = IntVar()
    self.num_runs = IntVar()
    self.data_source = StringVar()
    self.selection_mode = StringVar()
    self.use_global = IntVar()
    self.broken_tk = False
    self.use_regional = IntVar()

    self.master.title("namebench")
    outer_frame = Frame(self.master)
    outer_frame.grid(row=0, padx=16, pady=16)
    inner_frame = Frame(outer_frame, relief=GROOVE, bd=2, padx=12, pady=12)
    inner_frame.grid(row=0, columnspan=2)
    status = Label(outer_frame, text='...', textvariable=self.status)
    status.grid(row=15, sticky=W, column=0)

    bold_font = tkFont.Font(font=status['font'])
    bold_font['weight'] = 'bold'

    ns_label = Label(inner_frame, text="Nameservers")
    ns_label.grid(row=0, columnspan=2, sticky=W)
    ns_label['font'] = bold_font

    nameservers = Entry(inner_frame, bg="white", textvariable=self.nameserver_form, width=65)
    nameservers.grid(row=1, columnspan=2, sticky=W, padx=4, pady=2)
    self.nameserver_form.set(', '.join(util.InternalNameServers()))

    global_button = Checkbutton(inner_frame, text="Include global DNS providers (OpenDNS, UltraDNS)", variable=self.use_global)
    global_button.grid(row=2, columnspan=2, sticky=W)
    global_button.toggle()

    regional_button = Checkbutton(inner_frame, text="Include best available regional DNS services", variable=self.use_regional)
    regional_button.grid(row=3, columnspan=2, sticky=W)
    regional_button.toggle()

    separator = Frame(inner_frame, height=2, width=515, bd=1, relief=SUNKEN)
    separator.grid(row=4, padx=5, pady=5, columnspan=2)

    ds_label = Label(inner_frame, text="Benchmark Data Source")
    ds_label.grid(row=5, column=0, sticky=W)
    ds_label['font'] = bold_font

    numtests_label = Label(inner_frame, text="Number of tests")
    numtests_label.grid(row=5, column=1, sticky=W)
    numtests_label['font'] = bold_font

    self.DiscoverSources()
    source_titles = [history_parser.sourceToTitle(x) for x in self.sources]
    data_source = OptionMenu(inner_frame, self.data_source, *source_titles)
    data_source.configure(width=35)
    data_source.grid(row=6, column=0, sticky=W)
    self.data_source.set(source_titles[0])

    num_tests = Entry(inner_frame, bg="white", textvariable=self.num_tests)
    num_tests.grid(row=6, column=1, sticky=W, padx=4)
    self.num_tests.set(self.options.test_count)

    bds_label = Label(inner_frame, text="Benchmark Data Selection")
    bds_label.grid(row=7, column=0, sticky=W)
    bds_label['font'] = bold_font

    num_runs_label = Label(inner_frame, text="Number of runs")
    num_runs_label.grid(row=7, column=1, sticky=W)
    num_runs_label['font'] = bold_font

    selection_mode = OptionMenu(inner_frame, self.selection_mode, "Weighted", "Random", "Chunk")
    selection_mode.configure(width=35)
    selection_mode.grid(row=8, column=0, sticky=W)
    self.selection_mode.set('Weighted')

    num_runs = Entry(inner_frame, bg="white", textvariable=self.num_runs)
    num_runs.grid(row=8, column=1, sticky=W, padx=4)
    self.num_runs.set(self.options.run_count)

    self.button = Button(outer_frame, command=self.StartJob)
    self.button.grid(row=15, sticky=E, column=1, pady=4, padx=1)
    self.UpdateRunState(running=True)
    self.UpdateRunState(running=False)
    self.UpdateStatus('Ready')

  def MessageHandler(self, event):
    """Pinged when there is a new message in our queue to handle."""
    while global_message_queue.qsize():
      msg = global_message_queue.get()
      if msg.error:
        self.ErrorPopup(msg.message, msg.error)
      elif msg.enable_button == False:
        self.UpdateRunState(running=True)
      elif msg.enable_button == True:
        self.UpdateRunState(running=False)
      self.UpdateStatus(msg.message, count=msg.count, total=msg.total, error=msg.error)

  def UpdateStatus(self, message, count=None, total=None, error=None, debug=False):
    """Update our little status window."""
    if not message:
      return None

    if total:
      state = '%s... [%s/%s]' % (message, count, total)
    elif count:
      state = '%s%s' % (message, '.' * count)
    else:
      state = message

    print "> %s" % state
    self.status.set(state[0:60])

  def ErrorPopup(self, title, message):
    print "Showing popup: %s" % title
    tkMessageBox.showerror(str(title), str(message), master=self.master)

  def UpdateRunState(self, running=True):
    # try/except blocks added to work around broken Tcl/Tk libraries
    # shipped with Fedora 11 (not thread-safe).
    # See http://code.google.com/p/namebench/issues/detail?id=23'
    if self.broken_tk:
      return

    if running:
      try:
        self.button.config(state=DISABLED)
        self.button.config(text='Running')
      except TclError:
        self.broken_tk = True
        self.UpdateStatus('Unable to disable button due to broken Tk library')
      self.UpdateStatus('Running...')
    else:
      try:
        self.button.config(state=NORMAL)
        self.button.config(text='Start Benchmark')
      except TclError:
        pass

  def StartJob(self):
    """Events that get called when the Start button is pressed."""

    self.ProcessForm()
    thread = WorkerThread(self.primary, self.secondary, self.options,
                          history_parser=self.hparser,
                          master=self.master)
    thread.start()

  def ProcessForm(self):
    """Read form and populate instance variables."""
    self.primary = self.supplied_ns + util.ExtractIPTuplesFromString(self.nameserver_form.get())

    if self.use_global.get():
      self.primary += self.global_ns
    if self.use_regional.get():
      self.secondary = self.regional_ns
    else:
      self.secondary = []
    self.options.run_count = self.num_runs.get()
    self.options.test_count = self.num_tests.get()
    self.options.data_source = self.ParseSourceSelection(self.data_source.get())
    self.options.select_mode = self.selection_mode.get().lower()
<commit_msg>Only set global_last_message if no exception occurs<commit_after># Copyright 2009 Google Inc. All Rights Reserved.
#
# 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.

"""Tk user interface implementation for namebench."""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import os
import sys
import threading
import tkFont
import traceback
import Queue
from Tkinter import *
import tkMessageBox

import base_ui
import history_parser
import nameserver_list
import util

def closedWindowHandler():
  print 'Au revoir, mes amis!'
  sys.exit(1)

global_message_queue = Queue.Queue()
global_last_message = None

def AddMsg(message, master=None, **kwargs):
  """Add a message to the global queue for output."""
  global global_message_queue
  global global_last_message

  new_message = Message(message, **kwargs)
  if new_message != global_last_message:
    global_message_queue.put(new_message)
    if master:
      try:
        master.event_generate('<<msg>>', when='tail')
        global_last_message = new_message
      except TclError:
        print "TCL error encountered, not pushing update to UI:"
        traceback.print_exc()

class Message(object):
  """Messages to be passed from to the main thread from children.

  Used to avoid thread issues inherent with Tk.
  """
  def __init__(self, message, error=False, count=False, total=False,
               enable_button=None, debug=False):
    self.message = message
    self.error = error
    self.count = count
    self.debug = debug
    self.total = total
    self.enable_button = enable_button

class WorkerThread(threading.Thread, base_ui.BaseUI):
  """Handle benchmarking and preparation in a separate UI thread."""

  def __init__(self, primary, secondary, options, history_parser=None, master=None):
    threading.Thread.__init__(self)
    self.status_callback = self.msg
    self.hparser = history_parser
    self.primary = primary
    self.master = master
    self.secondary = secondary
    self.options = options
    self.resource_dir = os.path.dirname(os.path.dirname(__file__))

  def msg(self, message, **kwargs):
    """Add messages to the main queue."""
    return AddMsg(message, master=self.master, **kwargs)

  def run(self):
    self.msg('Started thread', enable_button=False)
    try:
      self.PrepareBenchmark()
      self.UpdateStatus('Here comes the pain...')
      self.RunBenchmark()
    except nameserver_list.OutgoingUdpInterception:
      (exc_type, exception, tb) = sys.exc_info()
      self.msg('Outgoing requests were intercepted!', error=exception)
    except nameserver_list.TooFewNameservers:
      (exc_type, exception, tb) = sys.exc_info()
      self.msg('Too few nameservers to test', error=exception)
    except:
      (exc_type, exception, tb) = sys.exc_info()
      traceback.print_exc(tb)
      error_msg = '\n'.join(traceback.format_tb(tb)[-4:])
      self.msg(exception, error=error_msg)
    self.msg(None, enable_button=True)

class NameBenchGui(object):
  """The main GUI."""
  def __init__(self, options, supplied_ns, global_ns, regional_ns, version=None):
    self.options = options
    self.supplied_ns = supplied_ns
    self.global_ns = global_ns
    self.regional_ns = regional_ns
    self.version = version

  def Execute(self):
    self.root = Tk()
    app = MainWindow(self.root, self.options, self.supplied_ns, self.global_ns, self.regional_ns, self.version)
    app.DrawWindow()
    self.root.bind('<<msg>>', app.MessageHandler)
    self.root.mainloop()

class MainWindow(Frame, base_ui.BaseUI):
  """The main Tk GUI class."""

  def __init__(self, master, options, supplied_ns, global_ns, regional_ns, version=None):
    """TODO(tstromberg): Remove duplication from NameBenchGui class."""
    Frame.__init__(self)
    self.master = master
    self.options = options
    self.supplied_ns = supplied_ns
    self.global_ns = global_ns
    self.regional_ns = regional_ns
    self.version = version
    self.master.protocol('WM_DELETE_WINDOW', closedWindowHandler)

  def DrawWindow(self):
    """Draws the user interface."""
    self.nameserver_form = StringVar()
    self.status = StringVar()
    self.num_tests = IntVar()
    self.num_runs = IntVar()
    self.data_source = StringVar()
    self.selection_mode = StringVar()
    self.use_global = IntVar()
    self.broken_tk = False
    self.use_regional = IntVar()

    self.master.title("namebench")
    outer_frame = Frame(self.master)
    outer_frame.grid(row=0, padx=16, pady=16)
    inner_frame = Frame(outer_frame, relief=GROOVE, bd=2, padx=12, pady=12)
    inner_frame.grid(row=0, columnspan=2)
    status = Label(outer_frame, text='...', textvariable=self.status)
    status.grid(row=15, sticky=W, column=0)

    bold_font = tkFont.Font(font=status['font'])
    bold_font['weight'] = 'bold'

    ns_label = Label(inner_frame, text="Nameservers")
    ns_label.grid(row=0, columnspan=2, sticky=W)
    ns_label['font'] = bold_font

    nameservers = Entry(inner_frame, bg="white", textvariable=self.nameserver_form, width=65)
    nameservers.grid(row=1, columnspan=2, sticky=W, padx=4, pady=2)
    self.nameserver_form.set(', '.join(util.InternalNameServers()))

    global_button = Checkbutton(inner_frame, text="Include global DNS providers (OpenDNS, UltraDNS)", variable=self.use_global)
    global_button.grid(row=2, columnspan=2, sticky=W)
    global_button.toggle()

    regional_button = Checkbutton(inner_frame, text="Include best available regional DNS services", variable=self.use_regional)
    regional_button.grid(row=3, columnspan=2, sticky=W)
    regional_button.toggle()

    separator = Frame(inner_frame, height=2, width=515, bd=1, relief=SUNKEN)
    separator.grid(row=4, padx=5, pady=5, columnspan=2)

    ds_label = Label(inner_frame, text="Benchmark Data Source")
    ds_label.grid(row=5, column=0, sticky=W)
    ds_label['font'] = bold_font

    numtests_label = Label(inner_frame, text="Number of tests")
    numtests_label.grid(row=5, column=1, sticky=W)
    numtests_label['font'] = bold_font

    self.DiscoverSources()
    source_titles = [history_parser.sourceToTitle(x) for x in self.sources]
    data_source = OptionMenu(inner_frame, self.data_source, *source_titles)
    data_source.configure(width=35)
    data_source.grid(row=6, column=0, sticky=W)
    self.data_source.set(source_titles[0])

    num_tests = Entry(inner_frame, bg="white", textvariable=self.num_tests)
    num_tests.grid(row=6, column=1, sticky=W, padx=4)
    self.num_tests.set(self.options.test_count)

    bds_label = Label(inner_frame, text="Benchmark Data Selection")
    bds_label.grid(row=7, column=0, sticky=W)
    bds_label['font'] = bold_font

    num_runs_label = Label(inner_frame, text="Number of runs")
    num_runs_label.grid(row=7, column=1, sticky=W)
    num_runs_label['font'] = bold_font

    selection_mode = OptionMenu(inner_frame, self.selection_mode, "Weighted", "Random", "Chunk")
    selection_mode.configure(width=35)
    selection_mode.grid(row=8, column=0, sticky=W)
    self.selection_mode.set('Weighted')

    num_runs = Entry(inner_frame, bg="white", textvariable=self.num_runs)
    num_runs.grid(row=8, column=1, sticky=W, padx=4)
    self.num_runs.set(self.options.run_count)

    self.button = Button(outer_frame, command=self.StartJob)
    self.button.grid(row=15, sticky=E, column=1, pady=4, padx=1)
    self.UpdateRunState(running=True)
    self.UpdateRunState(running=False)
    self.UpdateStatus('Ready')

  def MessageHandler(self, event):
    """Pinged when there is a new message in our queue to handle."""
    while global_message_queue.qsize():
      msg = global_message_queue.get()
      if msg.error:
        self.ErrorPopup(msg.message, msg.error)
      elif msg.enable_button == False:
        self.UpdateRunState(running=True)
      elif msg.enable_button == True:
        self.UpdateRunState(running=False)
      self.UpdateStatus(msg.message, count=msg.count, total=msg.total, error=msg.error)

  def UpdateStatus(self, message, count=None, total=None, error=None, debug=False):
    """Update our little status window."""
    if not message:
      return None

    if total:
      state = '%s... [%s/%s]' % (message, count, total)
    elif count:
      state = '%s%s' % (message, '.' * count)
    else:
      state = message

    print "> %s" % state
    self.status.set(state[0:60])

  def ErrorPopup(self, title, message):
    print "Showing popup: %s" % title
    tkMessageBox.showerror(str(title), str(message), master=self.master)

  def UpdateRunState(self, running=True):
    # try/except blocks added to work around broken Tcl/Tk libraries
    # shipped with Fedora 11 (not thread-safe).
    # See http://code.google.com/p/namebench/issues/detail?id=23'
    if self.broken_tk:
      return

    if running:
      try:
        self.button.config(state=DISABLED)
        self.button.config(text='Running')
      except TclError:
        self.broken_tk = True
        self.UpdateStatus('Unable to disable button due to broken Tk library')
      self.UpdateStatus('Running...')
    else:
      try:
        self.button.config(state=NORMAL)
        self.button.config(text='Start Benchmark')
      except TclError:
        pass

  def StartJob(self):
    """Events that get called when the Start button is pressed."""

    self.ProcessForm()
    thread = WorkerThread(self.primary, self.secondary, self.options,
                          history_parser=self.hparser,
                          master=self.master)
    thread.start()

  def ProcessForm(self):
    """Read form and populate instance variables."""
    self.primary = self.supplied_ns + util.ExtractIPTuplesFromString(self.nameserver_form.get())

    if self.use_global.get():
      self.primary += self.global_ns
    if self.use_regional.get():
      self.secondary = self.regional_ns
    else:
      self.secondary = []
    self.options.run_count = self.num_runs.get()
    self.options.test_count = self.num_tests.get()
    self.options.data_source = self.ParseSourceSelection(self.data_source.get())
    self.options.select_mode = self.selection_mode.get().lower()
<|endoftext|>


<commit_before>import SimpleHTTPServer
import SocketServer
import json
import logging
import os
import urlparse

DEPLOYED_PACKAGE = 'build'
DEVELOPMENT_PACKAGE = 'web'

class KcaaHTTPRequestHandler(SimpleHTTPServer.SimpleHTTPRequestHandler):

    GET_NEW_OBJECTS = '/get_new_objects'
    GET_OBJECT = '/get_object'
    CLIENT_PREFIX = '/client/'

    def do_HEAD(self):
        self.dispatch()

    def do_GET(self):
        self.dispatch()

    def log_message(self, format, *args):
        # Kill verbose HTTP logging.
        return

    def dispatch(self):
        o = urlparse.urlparse(self.path)
        if o.path == KcaaHTTPRequestHandler.GET_NEW_OBJECTS:
            self.handle_get_new_objects(o)
        elif o.path == KcaaHTTPRequestHandler.GET_OBJECT:
            self.handle_get_object(o)
        elif o.path.startswith(KcaaHTTPRequestHandler.CLIENT_PREFIX):
            self.handle_client(o)
        else:
            self.send_error(404, 'File not found: {}'.format(self.path))

    def handle_get_new_objects(self, o):
        if self.command != 'GET':
            self.send_error(501, 'Unknown method: {}'.format(self.command))
            return
        self.send_response(200)
        self.send_header('Content-Type', 'text/json; charset=UTF-8')
        self.end_headers()
        self.wfile.write(json.dumps(list(self.server.new_objects)))

    def handle_get_object(self, o):
        if self.command != 'GET':
            self.send_error(501, 'Unknown method: {}'.format(self.command))
            return
        queries = urlparse.parse_qs(o.query)
        try:
            object_type = queries['type']
            if len(object_type) > 1:
                self.send_error(400, 'Parameter type should have only 1 value')
                return
            object_type = object_type[0]
        except KeyError:
            self.send_error(400, 'Missing parameter: type')
            return
        try:
            data = self.server.objects[object_type]
            self.server.new_objects.remove(object_type)
            self.send_response(200)
            self.send_header('Content-Type', 'text/json; charset=UTF-8')
            self.end_headers()
            self.wfile.write(data)
        except KeyError:
            self.send_error(404)

    def handle_client(self, o):
        self.path = '/' + o.path[len(KcaaHTTPRequestHandler.CLIENT_PREFIX):]
        # Note: HTTP request handlers are not new-style classes.
        # super() cannot be used.
        if self.command == 'HEAD':
            SimpleHTTPServer.SimpleHTTPRequestHandler.do_HEAD(self)
        elif self.command == 'GET':
            SimpleHTTPServer.SimpleHTTPRequestHandler.do_GET(self)
        else:
            self.send_error(501, 'Unknown method: {}'.format(self.command))

def move_to_client_dir():
    # Change directory to client directory so that SimpleHTTPServer can serve
    # client resources.
    client_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                              'client'))
    if not os.path.isdir(client_dir):
        raise IOError('No client directory found: {}'.format(client_dir))
    os.chdir(client_dir)
    # Use 'build' subdirectory when deployed, or 'web' when being developed.
    if os.path.isdir(DEPLOYED_PACKAGE):
        os.chdir(DEPLOYED_PACKAGE)
    elif os.path.isdir(DEVELOPMENT_PACKAGE):
        os.chdir(DEVELOPMENT_PACKAGE)
    else:
        raise IOError('No client package directories found under {}.'.format(
            client_dir))

def setup(args, logger):
    move_to_client_dir()
    httpd = SocketServer.TCPServer(('', args.server_port),
                                   KcaaHTTPRequestHandler)
    _, port = httpd.server_address
    # Don't use query (something like ?key=value). Kancolle widget detects it
    # from referer and rejects to respond.
    root_url = 'http://localhost:{}/client/'.format(port)
    logger.info('KCAA server ready at {}'.format(root_url))
    return httpd, root_url

def handle_server(args, controller_conn, to_exit):
    logger = logging.getLogger('kcaa.server')
    # Wait until the controller reset the proxy.
    if not controller_conn.poll(12.0):
        logger.error('Controller couldn\'t reset the proxy. Shutting down.')
        to_exit.set()
        return
    assert controller_conn.recv()
    httpd, root_url = setup(args, logger)
    httpd.new_objects = set()
    httpd.objects = {}
    controller_conn.send(root_url)
    httpd.timeout = 1.0
    while True:
        httpd.handle_request()
        if to_exit.wait(0.0):
            logger.info('Server got an exit signal. Shutting down.')
            break
        while controller_conn.poll():
            object_type, data = controller_conn.recv()
            httpd.new_objects.add(object_type)
            httpd.objects[object_type] = data
    to_exit.set()
<commit_msg>Clarify the 404 error on unknown handler error.<commit_after>import SimpleHTTPServer
import SocketServer
import json
import logging
import os
import urlparse

DEPLOYED_PACKAGE = 'build'
DEVELOPMENT_PACKAGE = 'web'

class KcaaHTTPRequestHandler(SimpleHTTPServer.SimpleHTTPRequestHandler):

    GET_NEW_OBJECTS = '/get_new_objects'
    GET_OBJECT = '/get_object'
    CLIENT_PREFIX = '/client/'

    def do_HEAD(self):
        self.dispatch()

    def do_GET(self):
        self.dispatch()

    def log_message(self, format, *args):
        # Kill verbose HTTP logging.
        return

    def dispatch(self):
        o = urlparse.urlparse(self.path)
        if o.path == KcaaHTTPRequestHandler.GET_NEW_OBJECTS:
            self.handle_get_new_objects(o)
        elif o.path == KcaaHTTPRequestHandler.GET_OBJECT:
            self.handle_get_object(o)
        elif o.path.startswith(KcaaHTTPRequestHandler.CLIENT_PREFIX):
            self.handle_client(o)
        else:
            self.send_error(404, 'Unknown handler: {}'.format(self.path))

    def handle_get_new_objects(self, o):
        if self.command != 'GET':
            self.send_error(501, 'Unknown method: {}'.format(self.command))
            return
        self.send_response(200)
        self.send_header('Content-Type', 'text/json; charset=UTF-8')
        self.end_headers()
        self.wfile.write(json.dumps(list(self.server.new_objects)))

    def handle_get_object(self, o):
        if self.command != 'GET':
            self.send_error(501, 'Unknown method: {}'.format(self.command))
            return
        queries = urlparse.parse_qs(o.query)
        try:
            object_type = queries['type']
            if len(object_type) > 1:
                self.send_error(400, 'Parameter type should have only 1 value')
                return
            object_type = object_type[0]
        except KeyError:
            self.send_error(400, 'Missing parameter: type')
            return
        try:
            data = self.server.objects[object_type]
            self.server.new_objects.remove(object_type)
            self.send_response(200)
            self.send_header('Content-Type', 'text/json; charset=UTF-8')
            self.end_headers()
            self.wfile.write(data)
        except KeyError:
            self.send_error(404)

    def handle_client(self, o):
        self.path = '/' + o.path[len(KcaaHTTPRequestHandler.CLIENT_PREFIX):]
        # Note: HTTP request handlers are not new-style classes.
        # super() cannot be used.
        if self.command == 'HEAD':
            SimpleHTTPServer.SimpleHTTPRequestHandler.do_HEAD(self)
        elif self.command == 'GET':
            SimpleHTTPServer.SimpleHTTPRequestHandler.do_GET(self)
        else:
            self.send_error(501, 'Unknown method: {}'.format(self.command))

def move_to_client_dir():
    # Change directory to client directory so that SimpleHTTPServer can serve
    # client resources.
    client_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                              'client'))
    if not os.path.isdir(client_dir):
        raise IOError('No client directory found: {}'.format(client_dir))
    os.chdir(client_dir)
    # Use 'build' subdirectory when deployed, or 'web' when being developed.
    if os.path.isdir(DEPLOYED_PACKAGE):
        os.chdir(DEPLOYED_PACKAGE)
    elif os.path.isdir(DEVELOPMENT_PACKAGE):
        os.chdir(DEVELOPMENT_PACKAGE)
    else:
        raise IOError('No client package directories found under {}.'.format(
            client_dir))

def setup(args, logger):
    move_to_client_dir()
    httpd = SocketServer.TCPServer(('', args.server_port),
                                   KcaaHTTPRequestHandler)
    _, port = httpd.server_address
    # Don't use query (something like ?key=value). Kancolle widget detects it
    # from referer and rejects to respond.
    root_url = 'http://localhost:{}/client/'.format(port)
    logger.info('KCAA server ready at {}'.format(root_url))
    return httpd, root_url

def handle_server(args, controller_conn, to_exit):
    logger = logging.getLogger('kcaa.server')
    # Wait until the controller reset the proxy.
    if not controller_conn.poll(12.0):
        logger.error('Controller couldn\'t reset the proxy. Shutting down.')
        to_exit.set()
        return
    assert controller_conn.recv()
    httpd, root_url = setup(args, logger)
    httpd.new_objects = set()
    httpd.objects = {}
    controller_conn.send(root_url)
    httpd.timeout = 1.0
    while True:
        httpd.handle_request()
        if to_exit.wait(0.0):
            logger.info('Server got an exit signal. Shutting down.')
            break
        while controller_conn.poll():
            object_type, data = controller_conn.recv()
            httpd.new_objects.add(object_type)
            httpd.objects[object_type] = data
    to_exit.set()
<|endoftext|>


<commit_before>"""Automation using nox.
"""

# The following comment should be removed at some point in the future.
# mypy: disallow-untyped-defs=False

import glob
import os
import shutil
import sys

import nox

sys.path.append(".")
from tools.automation import release  # isort:skip  # noqa
sys.path.pop()

nox.options.reuse_existing_virtualenvs = True
nox.options.sessions = ["lint"]

LOCATIONS = {
    "common-wheels": "tests/data/common_wheels",
    "protected-pip": "tools/tox_pip.py",
}
REQUIREMENTS = {
    "tests": "tools/requirements/tests.txt",
    "common-wheels": "tools/requirements/tests-common_wheels.txt",
}

AUTHORS_FILE = "AUTHORS.txt"
VERSION_FILE = "src/pip/__init__.py"

def run_with_protected_pip(session, *arguments):
    """Do a session.run("pip", *arguments), using a "protected" pip.

    This invokes a wrapper script, that forwards calls to original virtualenv
    (stable) version, and not the code being tested. This ensures pip being
    used is not the code being tested.
    """
    env = {"VIRTUAL_ENV": session.virtualenv.location}

    command = ("python", LOCATIONS["protected-pip"]) + arguments
    kwargs = {"env": env, "silent": True}
    session.run(*command, **kwargs)

def should_update_common_wheels():
    # If the cache hasn't been created, create it.
    if not os.path.exists(LOCATIONS["common-wheels"]):
        return True

    # If the requirements was updated after cache, we'll repopulate it.
    cache_last_populated_at = os.path.getmtime(LOCATIONS["common-wheels"])
    requirements_updated_at = os.path.getmtime(REQUIREMENTS["common-wheels"])
    need_to_repopulate = requirements_updated_at > cache_last_populated_at

    # Clear the stale cache.
    if need_to_repopulate:
        shutil.rmtree(LOCATIONS["common-wheels"], ignore_errors=True)

    return need_to_repopulate

# -----------------------------------------------------------------------------
# Development Commands
#   These are currently prototypes to evaluate whether we want to switch over
#   completely to nox for all our automation. Contributors should prefer using
#   `tox -e ...` until this note is removed.
# -----------------------------------------------------------------------------
@nox.session(python=["2.7", "3.5", "3.6", "3.7", "3.8", "pypy", "pypy3"])
def test(session):
    # Get the common wheels.
    if should_update_common_wheels():
        run_with_protected_pip(
            session,
            "wheel",
            "-w", LOCATIONS["common-wheels"],
            "-r", REQUIREMENTS["common-wheels"],
        )
    else:
        msg = (
            "Re-using existing common-wheels at {}."
            .format(LOCATIONS["common-wheels"])
        )
        session.log(msg)

    # Build source distribution
    sdist_dir = os.path.join(session.virtualenv.location, "sdist")
    if os.path.exists(sdist_dir):
        shutil.rmtree(sdist_dir, ignore_errors=True)
    session.run(
        "python", "setup.py", "sdist",
        "--formats=zip", "--dist-dir", sdist_dir,
        silent=True,
    )
    generated_files = os.listdir(sdist_dir)
    assert len(generated_files) == 1
    generated_sdist = os.path.join(sdist_dir, generated_files[0])

    # Install source distribution
    run_with_protected_pip(session, "install", generated_sdist)

    # Install test dependencies
    run_with_protected_pip(session, "install", "-r", REQUIREMENTS["tests"])

    # Parallelize tests as much as possible, by default.
    arguments = session.posargs or ["-n", "auto"]

    # Run the tests
    #   LC_CTYPE is set to get UTF-8 output inside of the subprocesses that our
    #   tests use.
    session.run("pytest", *arguments, env={"LC_CTYPE": "en_US.UTF-8"})

@nox.session
def docs(session):
    session.install(".")
    session.install("-r", REQUIREMENTS["docs"])

    def get_sphinx_build_command(kind):
        # Having the conf.py in the docs/html is weird but needed because we
        # can not use a different configuration directory vs source directory
        # on RTD currently. So, we'll pass "-c docs/html" here.
        # See https://github.com/rtfd/readthedocs.org/issues/1543.
        return [
            "sphinx-build",
            "-W",
            "-c", "docs/html",  # see note above
            "-d", "docs/build/doctrees/" + kind,
            "-b", kind,
            "docs/" + kind,
            "docs/build/" + kind,
        ]

    session.run(*get_sphinx_build_command("html"))
    session.run(*get_sphinx_build_command("man"))

@nox.session
def lint(session):
    session.install("pre-commit")

    if session.posargs:
        args = session.posargs + ["--all-files"]
    else:
        args = ["--all-files", "--show-diff-on-failure"]

    session.run("pre-commit", "run", *args)

# -----------------------------------------------------------------------------
# Release Commands
# -----------------------------------------------------------------------------
@nox.session(name="prepare-release")
def prepare_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s prepare-release -- YY.N[.P]")

    session.log("# Ensure nothing is staged")
    if release.modified_files_in_git("--staged"):
        session.error("There are files staged in git")

    session.log(f"# Updating {AUTHORS_FILE}")
    release.generate_authors(AUTHORS_FILE)
    if release.modified_files_in_git():
        release.commit_file(
            session, AUTHORS_FILE, message=f"Update {AUTHORS_FILE}",
        )
    else:
        session.log(f"# No changes to {AUTHORS_FILE}")

    session.log("# Generating NEWS")
    release.generate_news(session, version)

    session.log(f"# Bumping for release {version}")
    release.update_version_file(version, VERSION_FILE)
    release.commit_file(session, VERSION_FILE, message="Bump for release")

    session.log("# Tagging release")
    release.create_git_tag(session, version, message=f"Release {version}")

    session.log("# Bumping for development")
    next_dev_version = release.get_next_development_version(version)
    release.update_version_file(next_dev_version, VERSION_FILE)
    release.commit_file(session, VERSION_FILE, message="Bump for development")

@nox.session(name="build-release")
def build_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s build-release -- YY.N[.P]")

    session.log("# Ensure no files in dist/")
    if release.have_files_in_folder("dist"):
        session.error(
            "There are files in dist/. Remove them and try again. "
            "You can use `git clean -fxdi -- dist` command to do this"
        )

    session.log("# Install dependencies")
    session.install("setuptools", "wheel", "twine")

    session.log("# Checkout the tag")
    session.run("git", "checkout", version, external=True, silent=True)

    session.log("# Cleanup build/ before building the wheel")
    if release.have_files_in_folder("build"):
        shutil.rmtree("build")

    session.log("# Build distributions")
    session.run("python", "setup.py", "sdist", "bdist_wheel", silent=True)

    session.log("# Verify distributions")
    session.run("twine", "check", *glob.glob("dist/*"), silent=True)

    session.log("# Checkout the master branch")
    session.run("git", "checkout", "master", external=True, silent=True)

@nox.session(name="upload-release")
def upload_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s upload-release -- YY.N[.P]")

    session.log("# Install dependencies")
    session.install("twine")

    distribution_files = glob.glob("dist/*")
    session.log(f"# Distribution files: {distribution_files}")

    # Sanity check: Make sure there's 2 distribution files.
    count = len(distribution_files)
    if count != 2:
        session.error(
            f"Expected 2 distribution files for upload, got {count}. "
            f"Remove dist/ and run 'nox -s build-release -- {version}'"
        )
    # Sanity check: Make sure the files are correctly named.
    expected_distribution_files = [
        f"dist/pip-{version}-py2.py3-none-any.whl",
        f"dist/pip-{version}.tar.gz",
    ]
    if sorted(distribution_files) != sorted(expected_distribution_files):
        session.error(
            f"Distribution files do not seem to be for {version} release."
        )

    session.log("# Upload distributions")
    session.run("twine", "upload", *distribution_files)
<commit_msg>🔥 Exterminate files that aren't Git-tracked<commit_after>"""Automation using nox.
"""

# The following comment should be removed at some point in the future.
# mypy: disallow-untyped-defs=False

import glob
import os
import shutil
import sys

import nox

sys.path.append(".")
from tools.automation import release  # isort:skip  # noqa
sys.path.pop()

nox.options.reuse_existing_virtualenvs = True
nox.options.sessions = ["lint"]

LOCATIONS = {
    "common-wheels": "tests/data/common_wheels",
    "protected-pip": "tools/tox_pip.py",
}
REQUIREMENTS = {
    "tests": "tools/requirements/tests.txt",
    "common-wheels": "tools/requirements/tests-common_wheels.txt",
}

AUTHORS_FILE = "AUTHORS.txt"
VERSION_FILE = "src/pip/__init__.py"

def run_with_protected_pip(session, *arguments):
    """Do a session.run("pip", *arguments), using a "protected" pip.

    This invokes a wrapper script, that forwards calls to original virtualenv
    (stable) version, and not the code being tested. This ensures pip being
    used is not the code being tested.
    """
    env = {"VIRTUAL_ENV": session.virtualenv.location}

    command = ("python", LOCATIONS["protected-pip"]) + arguments
    kwargs = {"env": env, "silent": True}
    session.run(*command, **kwargs)

def should_update_common_wheels():
    # If the cache hasn't been created, create it.
    if not os.path.exists(LOCATIONS["common-wheels"]):
        return True

    # If the requirements was updated after cache, we'll repopulate it.
    cache_last_populated_at = os.path.getmtime(LOCATIONS["common-wheels"])
    requirements_updated_at = os.path.getmtime(REQUIREMENTS["common-wheels"])
    need_to_repopulate = requirements_updated_at > cache_last_populated_at

    # Clear the stale cache.
    if need_to_repopulate:
        shutil.rmtree(LOCATIONS["common-wheels"], ignore_errors=True)

    return need_to_repopulate

# -----------------------------------------------------------------------------
# Development Commands
#   These are currently prototypes to evaluate whether we want to switch over
#   completely to nox for all our automation. Contributors should prefer using
#   `tox -e ...` until this note is removed.
# -----------------------------------------------------------------------------
@nox.session(python=["2.7", "3.5", "3.6", "3.7", "3.8", "pypy", "pypy3"])
def test(session):
    # Get the common wheels.
    if should_update_common_wheels():
        run_with_protected_pip(
            session,
            "wheel",
            "-w", LOCATIONS["common-wheels"],
            "-r", REQUIREMENTS["common-wheels"],
        )
    else:
        msg = (
            "Re-using existing common-wheels at {}."
            .format(LOCATIONS["common-wheels"])
        )
        session.log(msg)

    # Build source distribution
    sdist_dir = os.path.join(session.virtualenv.location, "sdist")
    if os.path.exists(sdist_dir):
        shutil.rmtree(sdist_dir, ignore_errors=True)
    session.run(
        "python", "setup.py", "sdist",
        "--formats=zip", "--dist-dir", sdist_dir,
        silent=True,
    )
    generated_files = os.listdir(sdist_dir)
    assert len(generated_files) == 1
    generated_sdist = os.path.join(sdist_dir, generated_files[0])

    # Install source distribution
    run_with_protected_pip(session, "install", generated_sdist)

    # Install test dependencies
    run_with_protected_pip(session, "install", "-r", REQUIREMENTS["tests"])

    # Parallelize tests as much as possible, by default.
    arguments = session.posargs or ["-n", "auto"]

    # Run the tests
    #   LC_CTYPE is set to get UTF-8 output inside of the subprocesses that our
    #   tests use.
    session.run("pytest", *arguments, env={"LC_CTYPE": "en_US.UTF-8"})

@nox.session
def docs(session):
    session.install(".")
    session.install("-r", REQUIREMENTS["docs"])

    def get_sphinx_build_command(kind):
        # Having the conf.py in the docs/html is weird but needed because we
        # can not use a different configuration directory vs source directory
        # on RTD currently. So, we'll pass "-c docs/html" here.
        # See https://github.com/rtfd/readthedocs.org/issues/1543.
        return [
            "sphinx-build",
            "-W",
            "-c", "docs/html",  # see note above
            "-d", "docs/build/doctrees/" + kind,
            "-b", kind,
            "docs/" + kind,
            "docs/build/" + kind,
        ]

    session.run(*get_sphinx_build_command("html"))
    session.run(*get_sphinx_build_command("man"))

@nox.session
def lint(session):
    session.install("pre-commit")

    if session.posargs:
        args = session.posargs + ["--all-files"]
    else:
        args = ["--all-files", "--show-diff-on-failure"]

    session.run("pre-commit", "run", *args)

# -----------------------------------------------------------------------------
# Release Commands
# -----------------------------------------------------------------------------
@nox.session(name="prepare-release")
def prepare_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s prepare-release -- YY.N[.P]")

    session.log("# Ensure nothing is staged")
    if release.modified_files_in_git("--staged"):
        session.error("There are files staged in git")

    session.log(f"# Updating {AUTHORS_FILE}")
    release.generate_authors(AUTHORS_FILE)
    if release.modified_files_in_git():
        release.commit_file(
            session, AUTHORS_FILE, message=f"Update {AUTHORS_FILE}",
        )
    else:
        session.log(f"# No changes to {AUTHORS_FILE}")

    session.log("# Generating NEWS")
    release.generate_news(session, version)

    session.log(f"# Bumping for release {version}")
    release.update_version_file(version, VERSION_FILE)
    release.commit_file(session, VERSION_FILE, message="Bump for release")

    session.log("# Tagging release")
    release.create_git_tag(session, version, message=f"Release {version}")

    session.log("# Bumping for development")
    next_dev_version = release.get_next_development_version(version)
    release.update_version_file(next_dev_version, VERSION_FILE)
    release.commit_file(session, VERSION_FILE, message="Bump for development")

@nox.session(name="build-release")
def build_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s build-release -- YY.N[.P]")

    session.log("# Ensure no files in dist/")
    if release.have_files_in_folder("dist"):
        session.error(
            "There are files in dist/. Remove them and try again. "
            "You can use `git clean -fxdi -- dist` command to do this"
        )

    session.log("# Install dependencies")
    session.install("setuptools", "wheel", "twine")

    session.log("# Checkout the tag")
    session.run("git", "checkout", version, external=True, silent=True)

    session.log("# Wipe Git-untracked files before building the wheel")
    session.run("git", "clean", "-fxd", external=True, silent=True)

    session.log("# Build distributions")
    session.run("python", "setup.py", "sdist", "bdist_wheel", silent=True)

    session.log("# Verify distributions")
    session.run("twine", "check", *glob.glob("dist/*"), silent=True)

    session.log("# Checkout the master branch")
    session.run("git", "checkout", "master", external=True, silent=True)

@nox.session(name="upload-release")
def upload_release(session):
    version = release.get_version_from_arguments(session.posargs)
    if not version:
        session.error("Usage: nox -s upload-release -- YY.N[.P]")

    session.log("# Install dependencies")
    session.install("twine")

    distribution_files = glob.glob("dist/*")
    session.log(f"# Distribution files: {distribution_files}")

    # Sanity check: Make sure there's 2 distribution files.
    count = len(distribution_files)
    if count != 2:
        session.error(
            f"Expected 2 distribution files for upload, got {count}. "
            f"Remove dist/ and run 'nox -s build-release -- {version}'"
        )
    # Sanity check: Make sure the files are correctly named.
    expected_distribution_files = [
        f"dist/pip-{version}-py2.py3-none-any.whl",
        f"dist/pip-{version}.tar.gz",
    ]
    if sorted(distribution_files) != sorted(expected_distribution_files):
        session.error(
            f"Distribution files do not seem to be for {version} release."
        )

    session.log("# Upload distributions")
    session.run("twine", "upload", *distribution_files)
<|endoftext|>


<commit_before>import numpy

def residuals_simple(f, g, s, x, dr, P, Q, parms):

    N = s.shape[0]
    n_s = s.shape[1]

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variable

    res = numpy.zeros_like(x)

    for i_ms in range(n_ms):
        # solving on grid for markov index i_ms
        # m = P[i_ms,:][None,:]
        m = numpy.tile(P[i_ms,:],(N,1))
        xm = x[i_ms,:,:]

        for I_ms in range(n_ms):

            # M = P[I_ms,:][None,:]
            M = numpy.tile(P[I_ms,:], (N,1))
            prob = Q[i_ms, I_ms]

            S = g(m, s, xm, M, parms)
            XM = dr(I_ms, S)
            rr = f(m,s,xm,M,S,XM,parms)
            res[i_ms,:,:] += prob*rr

    return res

def residuals(f, g, s, x, dr, P, Q, parms):

    N = s.shape[0]
    n_s = s.shape[1]
    n_x = x.shape[2]

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variable

    res = numpy.zeros_like(x)

    m = P

    s_ = numpy.tile(s, (n_ms*n_ms, 1))
    x_ = numpy.repeat(x[None,:,:,:], n_ms, axis=0)
    m_ = numpy.repeat( m, N, axis=0 )
    m_ = numpy.tile( m_, (n_ms,1) )

    M_ = numpy.repeat( m, N*n_ms, axis=0 )

    # reshape everything as 2d arrays
    s_ = s_.reshape( (-1, n_s) )
    x_ = x_.reshape( (-1, n_x) )
    m_ = m_.reshape( (-1, n_mv) )
    M_ = M_.reshape( (-1, n_mv) )

    S_ = g(m_, s_, x_, M_, parms)
    X_ = numpy.zeros_like(x_)

    S_r = S_.reshape( (n_ms, n_ms, N, n_s) )
    X_r = X_.reshape( (n_ms, n_ms, N, n_x) )

    for I_ms in range(n_ms):
        SS_ = S_r[I_ms, :, :, :].reshape( (n_ms*N, n_s) )
        X_r[I_ms,:,:,:] = dr(I_ms, SS_).reshape( (n_ms, N, n_x) )

    rr = f(m_, s_, x_, M_, S_, X_, parms)

    rr = rr.reshape( (n_ms, n_ms, N, n_x) )

    for i_ms in range(n_ms):
        res[i_ms,:,:] = 0
        for I_ms in range(n_ms):
            res[i_ms, :, :] += Q[i_ms, I_ms] * rr[I_ms, i_ms, :, :]

    return res

def solve_mfg_model(model, initial_guess=None, with_complementarities=True, verbose=True, orders=None, output_type='dr', maxit=1000, inner_maxit=5, tol=1e-8) :

    assert(model.model_type == 'mfga')

    [P, Q] = model.markov_chain

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variables

    x0 = model.calibration['controls']
    parms = model.calibration['parameters']
    n_x = len(x0)
    n_s = len(model.symbols['states'])

    approx = model.options['approximation_space']
    a = approx['a']
    b = approx['b']
    if orders is None:
        orders = approx['orders']

    from dolo.numeric.decision_rules_markov import MarkovDecisionRule

    mdr = MarkovDecisionRule(n_ms, a, b, orders)

    grid = mdr.grid
    N = grid.shape[0]

#    if isinstance(initial_guess, numpy.ndarray):
#        print("Using initial guess (1)")
#        controls = initial_guess
#    elif isinstance(initial_guess, dict):
#        print("Using initial guess (2)")
#        controls_0 = initial_guess['controls']
#        ap_space = initial_guess['approximation_space']
#        if False in (approx['orders']==orders):
#            print("Interpolating initial guess")
#            old_dr = MarkovDecisionRule(controls_0.shape[0], ap_space['smin'], ap_space['smax'], ap_space['orders'])
#            old_dr.set_values(controls_0)
#            controls_0 = numpy.zeros( (n_ms, N, n_x) )
#            for i in range(n_ms):
#                e = old_dr(i,grid)
#                controls_0[i,:,:] = e
#    else:
#        controls_0 = numpy.zeros((n_ms, N, n_x))

    controls_0 = numpy.zeros((n_ms, N, n_x))

    if initial_guess is None:
        controls_0[:,:,:] = x0[None,None,:]
    else:
        for i_m in range(n_ms):
            m = P[i_m,:][None,:]
            controls_0[i_m,:,:] = initial_guess(i_m, grid)

    ff = model.functions['arbitrage']
    gg = model.functions['transition']
    aa = model.functions['auxiliary']

    if 'arbitrage_lb' in model.functions and with_complementarities==True:
        lb_fun = model.functions['arbitrage_lb']
        ub_fun = model.functions['arbitrage_ub']
        lb = numpy.zeros_like(controls_0)*numpy.nan
        ub = numpy.zeros_like(controls_0)*numpy.nan
        for i_m in range(n_ms):
            m = P[i_m,:][None,:]
            p = parms[None,:]
            m = numpy.repeat(m, N, axis=0)
            p = numpy.repeat(p, N, axis=0)

            lb[i_m,:,:] = lb_fun(m, grid, p)
            ub[i_m,:,:] = ub_fun(m, grid, p)

    else:
        with_complementarities = False

    f = lambda m,s,x,M,S,X,p: ff(m,s,x,aa(m,s,x,p),M,S,X,aa(M,S,X,p),p)
    g = lambda m,s,x,M,p: gg(m,s,x,aa(m,s,x,p),M,p)

    # mdr.set_values(controls)

    sh_c = controls_0.shape

    controls_0 = controls_0.reshape( (-1,n_x) )

    from dolo.numeric.optimize.newton import newton, SerialDifferentiableFunction
    from dolo.numeric.optimize.ncpsolve import ncpsolve

    err = 10
    tol = 1e-8
    it = 0

    if with_complementarities:
        print("Solving WITH complementarities.")
        lb = lb.reshape((-1,n_x))
        ub = ub.reshape((-1,n_x))

    if verbose:
        headline = '|{0:^4} | {1:10} | {2:8} | {3:8} | {4:3} |'.format( 'N',' Error', 'Gain','Time',  'nit' )
        stars = '-'*len(headline)
        print(stars)
        print(headline)
        print(stars)

    import time
    t1 = time.time()

    err_0 = numpy.nan

    verbit = (verbose == 'full')

    while err>tol and it<maxit:

        it += 1

        t_start = time.time()

        mdr.set_values(controls_0.reshape(sh_c))

        fn = lambda x: residuals(f, g, grid, x.reshape(sh_c), mdr, P, Q, parms).reshape((-1,n_x))
        dfn = SerialDifferentiableFunction(fn)

        if with_complementarities:
            [controls,nit] = ncpsolve(dfn, lb, ub, controls_0, verbose=verbit, maxit=inner_maxit)
        else:
            [controls, nit] = newton(dfn, controls_0, verbose=verbit, maxit=inner_maxit)

        err = abs(controls-controls_0).max()

        err_SA = err/err_0
        err_0 = err

        controls_0 = controls

        t_finish = time.time()
        elapsed = t_finish - t_start

        if verbose:
            print('|{0:4} | {1:10.3e} | {2:8.3f} | {3:8.3f} | {4:3} |'.format( it, err, err_SA, elapsed, nit  ))

    controls_0 = controls.reshape(sh_c)

    t2 = time.time()

    if verbose:
        print(stars)
        print("Elapsed: {} seconds.".format(t2-t1))
        print(stars)

    if output_type == 'dr':
        return mdr
    elif output_type == 'controls':
        return controls_0
    else:
        raise Exception("Unsupported ouput type {}.".format(output_type))

if __name__ == '__main__':

    from dolo import *
    model = yaml_import("../../../examples/models/rbc_mfga.yaml")
    print(model.calibration['states'])
    print(model.calibration_dict)
    print(model.markov_chain)

    initial_guess_symbolic = [
        'i = delta*k',
        'n = 0.33'
    ]

    from dolo.compiler.function_compiler_ast import compile_function_ast
    from dolo.compiler.function_compiler import standard_function

    arg_names = [
        ['markov_states',0,'m'],
        ['states',0,'s'],
        ['parameters',0,'p']
    ]

    fun = compile_function_ast(initial_guess_symbolic, model.symbols, arg_names,'initial_guess')
    ff = standard_function(fun, len(initial_guess_symbolic))

    sol = solve_mfg_model(model, initial_guess=ff)
<commit_msg>Changed default tolerance for mfg.time-iteration<commit_after>import numpy

def residuals_simple(f, g, s, x, dr, P, Q, parms):

    N = s.shape[0]
    n_s = s.shape[1]

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variable

    res = numpy.zeros_like(x)

    for i_ms in range(n_ms):
        # solving on grid for markov index i_ms
        # m = P[i_ms,:][None,:]
        m = numpy.tile(P[i_ms,:],(N,1))
        xm = x[i_ms,:,:]

        for I_ms in range(n_ms):

            # M = P[I_ms,:][None,:]
            M = numpy.tile(P[I_ms,:], (N,1))
            prob = Q[i_ms, I_ms]

            S = g(m, s, xm, M, parms)
            XM = dr(I_ms, S)
            rr = f(m,s,xm,M,S,XM,parms)
            res[i_ms,:,:] += prob*rr

    return res

def residuals(f, g, s, x, dr, P, Q, parms):

    N = s.shape[0]
    n_s = s.shape[1]
    n_x = x.shape[2]

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variable

    res = numpy.zeros_like(x)

    m = P

    s_ = numpy.tile(s, (n_ms*n_ms, 1))
    x_ = numpy.repeat(x[None,:,:,:], n_ms, axis=0)
    m_ = numpy.repeat( m, N, axis=0 )
    m_ = numpy.tile( m_, (n_ms,1) )

    M_ = numpy.repeat( m, N*n_ms, axis=0 )

    # reshape everything as 2d arrays
    s_ = s_.reshape( (-1, n_s) )
    x_ = x_.reshape( (-1, n_x) )
    m_ = m_.reshape( (-1, n_mv) )
    M_ = M_.reshape( (-1, n_mv) )

    S_ = g(m_, s_, x_, M_, parms)
    X_ = numpy.zeros_like(x_)

    S_r = S_.reshape( (n_ms, n_ms, N, n_s) )
    X_r = X_.reshape( (n_ms, n_ms, N, n_x) )

    for I_ms in range(n_ms):
        SS_ = S_r[I_ms, :, :, :].reshape( (n_ms*N, n_s) )
        X_r[I_ms,:,:,:] = dr(I_ms, SS_).reshape( (n_ms, N, n_x) )

    rr = f(m_, s_, x_, M_, S_, X_, parms)

    rr = rr.reshape( (n_ms, n_ms, N, n_x) )

    for i_ms in range(n_ms):
        res[i_ms,:,:] = 0
        for I_ms in range(n_ms):
            res[i_ms, :, :] += Q[i_ms, I_ms] * rr[I_ms, i_ms, :, :]

    return res

def solve_mfg_model(model, initial_guess=None, with_complementarities=True, verbose=True, orders=None, output_type='dr', maxit=1000, inner_maxit=10, tol=1e-6) :

    assert(model.model_type == 'mfga')

    [P, Q] = model.markov_chain

    n_ms = P.shape[0]   # number of markov states
    n_mv = P.shape[1] # number of markov variables

    x0 = model.calibration['controls']
    parms = model.calibration['parameters']
    n_x = len(x0)
    n_s = len(model.symbols['states'])

    approx = model.options['approximation_space']
    a = approx['a']
    b = approx['b']
    if orders is None:
        orders = approx['orders']

    from dolo.numeric.decision_rules_markov import MarkovDecisionRule

    mdr = MarkovDecisionRule(n_ms, a, b, orders)

    grid = mdr.grid
    N = grid.shape[0]

#    if isinstance(initial_guess, numpy.ndarray):
#        print("Using initial guess (1)")
#        controls = initial_guess
#    elif isinstance(initial_guess, dict):
#        print("Using initial guess (2)")
#        controls_0 = initial_guess['controls']
#        ap_space = initial_guess['approximation_space']
#        if False in (approx['orders']==orders):
#            print("Interpolating initial guess")
#            old_dr = MarkovDecisionRule(controls_0.shape[0], ap_space['smin'], ap_space['smax'], ap_space['orders'])
#            old_dr.set_values(controls_0)
#            controls_0 = numpy.zeros( (n_ms, N, n_x) )
#            for i in range(n_ms):
#                e = old_dr(i,grid)
#                controls_0[i,:,:] = e
#    else:
#        controls_0 = numpy.zeros((n_ms, N, n_x))

    controls_0 = numpy.zeros((n_ms, N, n_x))

    if initial_guess is None:
        controls_0[:,:,:] = x0[None,None,:]
    else:
        for i_m in range(n_ms):
            m = P[i_m,:][None,:]
            controls_0[i_m,:,:] = initial_guess(i_m, grid)

    ff = model.functions['arbitrage']
    gg = model.functions['transition']
    aa = model.functions['auxiliary']

    if 'arbitrage_lb' in model.functions and with_complementarities==True:
        lb_fun = model.functions['arbitrage_lb']
        ub_fun = model.functions['arbitrage_ub']
        lb = numpy.zeros_like(controls_0)*numpy.nan
        ub = numpy.zeros_like(controls_0)*numpy.nan
        for i_m in range(n_ms):
            m = P[i_m,:][None,:]
            p = parms[None,:]
            m = numpy.repeat(m, N, axis=0)
            p = numpy.repeat(p, N, axis=0)

            lb[i_m,:,:] = lb_fun(m, grid, p)
            ub[i_m,:,:] = ub_fun(m, grid, p)

    else:
        with_complementarities = False

    f = lambda m,s,x,M,S,X,p: ff(m,s,x,aa(m,s,x,p),M,S,X,aa(M,S,X,p),p)
    g = lambda m,s,x,M,p: gg(m,s,x,aa(m,s,x,p),M,p)

    # mdr.set_values(controls)

    sh_c = controls_0.shape

    controls_0 = controls_0.reshape( (-1,n_x) )

    from dolo.numeric.optimize.newton import newton, SerialDifferentiableFunction
    from dolo.numeric.optimize.ncpsolve import ncpsolve

    err = 10
    it = 0

    if with_complementarities:
        print("Solving WITH complementarities.")
        lb = lb.reshape((-1,n_x))
        ub = ub.reshape((-1,n_x))

    if verbose:
        headline = '|{0:^4} | {1:10} | {2:8} | {3:8} | {4:3} |'.format( 'N',' Error', 'Gain','Time',  'nit' )
        stars = '-'*len(headline)
        print(stars)
        print(headline)
        print(stars)

    import time
    t1 = time.time()

    err_0 = numpy.nan

    verbit = (verbose == 'full')

    while err>tol and it<maxit:

        it += 1

        t_start = time.time()

        mdr.set_values(controls_0.reshape(sh_c))

        fn = lambda x: residuals(f, g, grid, x.reshape(sh_c), mdr, P, Q, parms).reshape((-1,n_x))
        dfn = SerialDifferentiableFunction(fn)

        if with_complementarities:
            [controls,nit] = ncpsolve(dfn, lb, ub, controls_0, verbose=verbit, maxit=inner_maxit)
        else:
            [controls, nit] = newton(dfn, controls_0, verbose=verbit, maxit=inner_maxit)

        err = abs(controls-controls_0).max()

        err_SA = err/err_0
        err_0 = err

        controls_0 = controls

        t_finish = time.time()
        elapsed = t_finish - t_start

        if verbose:
            print('|{0:4} | {1:10.3e} | {2:8.3f} | {3:8.3f} | {4:3} |'.format( it, err, err_SA, elapsed, nit  ))

    controls_0 = controls.reshape(sh_c)

    t2 = time.time()

    if verbose:
        print(stars)
        print("Elapsed: {} seconds.".format(t2-t1))
        print(stars)

    if output_type == 'dr':
        return mdr
    elif output_type == 'controls':
        return controls_0
    else:
        raise Exception("Unsupported ouput type {}.".format(output_type))

if __name__ == '__main__':

    from dolo import *
    model = yaml_import("../../../examples/models/rbc_mfga.yaml")
    print(model.calibration['states'])
    print(model.calibration_dict)
    print(model.markov_chain)

    initial_guess_symbolic = [
        'i = delta*k',
        'n = 0.33'
    ]

    from dolo.compiler.function_compiler_ast import compile_function_ast
    from dolo.compiler.function_compiler import standard_function

    arg_names = [
        ['markov_states',0,'m'],
        ['states',0,'s'],
        ['parameters',0,'p']
    ]

    fun = compile_function_ast(initial_guess_symbolic, model.symbols, arg_names,'initial_guess')
    ff = standard_function(fun, len(initial_guess_symbolic))

    sol = solve_mfg_model(model, initial_guess=ff)
<|endoftext|>


<commit_before>#!/usr/bin/python

import argparse
import json
import os, sys
from netboxapi_client import Api, create, show, enum, delete, update, patch
import urllib3

def get_configuration(path="{}/netboxapi.json".format(os.getcwd())):
    """get_configuration

    Returns a dictionnary containing all elements written as json, in the
    configuration file.

    :param path:
    """
    try:
        with open(path) as fd:
            return json.load(fd)
    except Exception:
        print("No configuration file found at {}. Reading environment variables NETBOXAPI_TOKEN and NETBOXAPI_URL.".format(path))
        if 'NETBOXAPI_TOKEN' in os.environ and 'NETBOXAPI_URL' in os.environ:
          config = { "url": os.environ.get('NETBOXAPI_URL'), "token": os.environ.get('NETBOXAPI_TOKEN') }
          return config
        else:
          print("Configuration not properly defined.")
        sys.exit(254)

def main():
    parser = argparse.ArgumentParser()

    ## NOT IMPLEMENTED YET
    #parser.add_argument(
    #    "--url",
    #    help="URL of the target Netbox instance",
    #    dest='url'
    #)
    #parser.add_argument(
    #    "--token",
    #    help="Token that should be used for authentication",
    #    dest='token'
    #)

    config = get_configuration()
    api = Api(
        url=config['url'],
        token=config['token']
        if 'insecure' in config:
            if config['insecure'] is 'True':
                # Disable InsecureRequestWarning
                urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    )

    subparsers = parser.add_subparsers(dest='model')

    mydict = api.get('').json()

    FUNCTION_MAP = {
        'show': show,
        'list': enum,
        'create': create,
        'delete': delete,
        'update': update,
        'patch': patch
    }

    DESCRIPTION_MAP = {
        'show': "Shows an object's data.",
        'list': "Lists objects.",
        'create': "Creates object.",
        'delete': "Deletes object",
        'update': "Updates object. It is required to fill all fields you want to be filled for that object.",
        'patch': "Updates one or some fields of an object. Unlike update method, you can specify only the fiels you want to update."
    }

    ARGUMENTS = {
        'identifier': {
            'option': '-i',
            'type': int,
            'help': "Numerical id of the object",
            'dest': 'id'
        },
        'data_file': {
            'option': '-d',
            'type': str,
            'help': "Path to the file containing json data.",
            'dest': 'data'
        },
        'name': {
            'option': '-n',
            'type': str,
            'help': "Name of the object.",
            'dest': 'name'
        },
        'configuration_file': {
          'option': '-c',
          'type': str,
          'help': "Path to the json configuration file.",
          'dest': 'config'
        }
    }

    for model in mydict.keys():
        model_parser = subparsers.add_parser(
            model,
            help="{0} objects".format(model)
        )
        model_subparser = model_parser.add_subparsers(dest='object')
        for elmt in api.get(model+'/').json():
            elmt_parser = model_subparser.add_parser(
                elmt,
                help="manage {0}".format(elmt)
            )
            action_subparser = elmt_parser.add_subparsers(dest="action")
            for k, v in FUNCTION_MAP.items():
                action_parser = action_subparser.add_parser(
                    k, help="{}".format(DESCRIPTION_MAP[k])
                )
                for k, arg in ARGUMENTS.items():
                    action_parser.add_argument(
                        arg['option'],
                        type=arg['type'],
                        help=arg['help'],
                        dest=arg['dest'],
                    )
    ns = parser.parse_args()
    if 'action' in ns:
        if 'data' in ns and ns.data:
                pprint(
                    FUNCTION_MAP[ns.action](
                        api=api,
                        model=ns.model,
                        obj=ns.object,
                        ident=ns.id,
                        name=ns.name,
                        data=json.loads(ns.data)
                    )
                )
        else:
            FUNCTION_MAP[ns.action](
                api=api,
                model=ns.model,
                obj=ns.object,
                ident=ns.id,
                name=ns.name
            )

if __name__ == "__main__":
    main()
<commit_msg>move insecure to config<commit_after>#!/usr/bin/python

import argparse
import json
import os, sys
from netboxapi_client import Api, create, show, enum, delete, update, patch
import urllib3

def get_configuration(path="{}/netboxapi.json".format(os.getcwd())):
    """get_configuration

    Returns a dictionnary containing all elements written as json, in the
    configuration file.

    :param path:
    """
    try:
        with open(path) as fd:
            return json.load(fd)
    except Exception:
        print("No configuration file found at {}. Reading environment variables NETBOXAPI_TOKEN and NETBOXAPI_URL.".format(path))
        if 'NETBOXAPI_TOKEN' in os.environ and 'NETBOXAPI_URL' in os.environ:
          config = { "url": os.environ.get('NETBOXAPI_URL'), "token": os.environ.get('NETBOXAPI_TOKEN') }
          return config
        else:
          print("Configuration not properly defined.")
        sys.exit(254)

def main():
    parser = argparse.ArgumentParser()

    ## NOT IMPLEMENTED YET
    #parser.add_argument(
    #    "--url",
    #    help="URL of the target Netbox instance",
    #    dest='url'
    #)
    #parser.add_argument(
    #    "--token",
    #    help="Token that should be used for authentication",
    #    dest='token'
    #)

    config = get_configuration()

    # Disable InsecureRequestWarning
    if 'insecure' in config:
        if 'True' in config['insecure']:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    api = Api(
        url=config['url'],
        token=config['token']
    )

    subparsers = parser.add_subparsers(dest='model')

    mydict = api.get('').json()

    FUNCTION_MAP = {
        'show': show,
        'list': enum,
        'create': create,
        'delete': delete,
        'update': update,
        'patch': patch
    }

    DESCRIPTION_MAP = {
        'show': "Shows an object's data.",
        'list': "Lists objects.",
        'create': "Creates object.",
        'delete': "Deletes object",
        'update': "Updates object. It is required to fill all fields you want to be filled for that object.",
        'patch': "Updates one or some fields of an object. Unlike update method, you can specify only the fiels you want to update."
    }

    ARGUMENTS = {
        'identifier': {
            'option': '-i',
            'type': int,
            'help': "Numerical id of the object",
            'dest': 'id'
        },
        'data_file': {
            'option': '-d',
            'type': str,
            'help': "Path to the file containing json data.",
            'dest': 'data'
        },
        'name': {
            'option': '-n',
            'type': str,
            'help': "Name of the object.",
            'dest': 'name'
        },
        'configuration_file': {
          'option': '-c',
          'type': str,
          'help': "Path to the json configuration file.",
          'dest': 'config'
        }
    }

    for model in mydict.keys():
        model_parser = subparsers.add_parser(
            model,
            help="{0} objects".format(model)
        )
        model_subparser = model_parser.add_subparsers(dest='object')
        for elmt in api.get(model+'/').json():
            elmt_parser = model_subparser.add_parser(
                elmt,
                help="manage {0}".format(elmt)
            )
            action_subparser = elmt_parser.add_subparsers(dest="action")
            for k, v in FUNCTION_MAP.items():
                action_parser = action_subparser.add_parser(
                    k, help="{}".format(DESCRIPTION_MAP[k])
                )
                for k, arg in ARGUMENTS.items():
                    action_parser.add_argument(
                        arg['option'],
                        type=arg['type'],
                        help=arg['help'],
                        dest=arg['dest'],
                    )
    ns = parser.parse_args()
    if 'action' in ns:
        if 'data' in ns and ns.data:
                pprint(
                    FUNCTION_MAP[ns.action](
                        api=api,
                        model=ns.model,
                        obj=ns.object,
                        ident=ns.id,
                        name=ns.name,
                        data=json.loads(ns.data)
                    )
                )
        else:
            FUNCTION_MAP[ns.action](
                api=api,
                model=ns.model,
                obj=ns.object,
                ident=ns.id,
                name=ns.name
            )

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