<commit_before># This file is part of Moksha.
# Copyright (C) 2008-2014  Red Hat, Inc.
#
# 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.

"""
:mod:`moksha.hub.api.consumer` - The Moksha Consumer API
========================================================
Moksha provides a simple API for creating "consumers" of message topics.

This means that your consumer is instantiated when the MokshaHub is initially
loaded, and receives each message for the specified topic through the
:meth:`Consumer.consume` method.

.. moduleauthor:: Luke Macken <lmacken@redhat.com>
"""

import json
import logging
log = logging.getLogger('moksha.hub')

from kitchen.iterutils import iterate
from moksha.common.lib.helpers import create_app_engine

class Consumer(object):
    """ A message consumer """
    topic = ''

    # Automatically decode JSON data
    jsonify = True

    # Internal use only
    _initialized = False
    _exception_count = 0

    def __init__(self, hub):
        self.hub = hub
        self.log = log

        callback = self._consume
        if self.jsonify:
            callback = self._consume_json

        for topic in iterate(self.topic):
            log.debug('Subscribing to consumer topic %s' % topic)
            self.hub.subscribe(topic, callback)

        # If the consumer specifies an 'app', then setup `self.engine` to
        # be a SQLAlchemy engine, along with a configured DBSession
        app = getattr(self, 'app', None)
        self.engine = self.DBSession = None
        if app:
            log.debug("Setting up individual engine for consumer")
            from sqlalchemy.orm import sessionmaker
            self.engine = create_app_engine(app, hub.config)
            self.DBSession = sessionmaker(bind=self.engine)()

        self._initialized = True

    def __json__(self):
        return {
            "name": type(self).__name__,
            "module": type(self).__module__,
            "topic": self.topic,
            "initialized": self._initialized,
            "exceptions": self._exception_count,
            "jsonify": self.jsonify,
        }

    def _consume_json(self, message):
        """ Convert our AMQP messages into a consistent dictionary format.

        This method exists because our STOMP & AMQP message brokers consume
        messages in different formats.  This causes our messaging abstraction
        to leak into the consumers themselves.

        :Note: We do not pass the message headers to the consumer (in this AMQP consumer)
        because the current AMQP.js bindings do not allow the client to change them.
        Thus, we need to throw any topic/queue details into the JSON body itself.
        """
        try:
            body = json.loads(message.body)
        except:
            log.debug("Unable to decode message body to JSON: %r" % message.body)
            body = message.body
        topic = None

        # Try some stuff for AMQP:
        try:
            topic = message.headers[0].routing_key
        except TypeError:
            # We didn't get a JSON dictionary
            pass
        except AttributeError:
            # We didn't get headers or a routing key?
            pass

        # If that didn't work, it might be zeromq
        if not topic:
            try:
                topic = message.topic
            except AttributeError:
                # Weird.  I have no idea...
                pass

        try:
            message_as_dict = {'body': body, 'topic': topic}
            self._consume(message_as_dict)
            self._exception_count = 0  # Reset if everything went swimmingly
        except Exception:
            # Otherwise, keep track of how many exceptions we've hit in a row
            self._exception_count = self._exception_count + 1
            # And then re-raise the exception to be logged
            raise

    def _consume(self, message):
        try:
            self.validate(message)
        except Exception, e:
            log.warn("Received invalid message %r" % e)
            return

        self.consume(message)

    def validate(self, message):
        """ Override to implement your own validation scheme. """
        pass

    def consume(self, message):
        raise NotImplementedError

    def send_message(self, topic, message):
        try:
            self.hub.send_message(topic, message)
        except Exception, e:
            log.error('Cannot send message: %s' % e)

    def stop(self):
        if hasattr(self, 'hub'):
            self.hub.close()

        if getattr(self, 'DBSession', None):
            self.DBSession.close()
<commit_msg>Threaded consumers.<commit_after># This file is part of Moksha.
# Copyright (C) 2008-2014  Red Hat, Inc.
#
# 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.

"""
:mod:`moksha.hub.api.consumer` - The Moksha Consumer API
========================================================
Moksha provides a simple API for creating "consumers" of message topics.

This means that your consumer is instantiated when the MokshaHub is initially
loaded, and receives each message for the specified topic through the
:meth:`Consumer.consume` method.

.. moduleauthor:: Luke Macken <lmacken@redhat.com>
"""

import json
import threading
import Queue as queue
import logging
log = logging.getLogger('moksha.hub')

from kitchen.iterutils import iterate
from moksha.common.lib.helpers import create_app_engine
import moksha.hub.reactor

class Consumer(object):
    """ A message consumer """
    topic = ''

    # Automatically decode JSON data
    jsonify = True

    # Internal use only
    _initialized = False
    _exception_count = 0

    def __init__(self, hub):
        self.hub = hub
        self.log = log

        # Set up a queue to communicate between the main twisted thread
        # receiving raw messages, and a worker thread that pulls items off
        # the queue to do "consume" work.
        # TODO -- someday, have this be monitored by the monitoring socket
        self.incoming = queue.Queue()

        callback = self._consume
        if self.jsonify:
            callback = self._consume_json

        for topic in iterate(self.topic):
            log.debug('Subscribing to consumer topic %s' % topic)
            self.hub.subscribe(topic, callback)

        # If the consumer specifies an 'app', then setup `self.engine` to
        # be a SQLAlchemy engine, along with a configured DBSession
        app = getattr(self, 'app', None)
        self.engine = self.DBSession = None
        if app:
            log.debug("Setting up individual engine for consumer")
            from sqlalchemy.orm import sessionmaker
            self.engine = create_app_engine(app, hub.config)
            self.DBSession = sessionmaker(bind=self.engine)()

        self.N = int(self.hub.config.get('moksha.workers_per_consumer', 0))
        for i in range(self.N):
            moksha.hub.reactor.reactor.callInThread(self._work)

        self._initialized = True

    def __json__(self):
        return {
            "name": type(self).__name__,
            "module": type(self).__module__,
            "topic": self.topic,
            "initialized": self._initialized,
            "exceptions": self._exception_count,
            "jsonify": self.jsonify,
        }

    def debug(self, message):
        idx = threading.current_thread().ident
        log.debug("%r thread %r | %s" % (type(self).__name__, idx, message))

    def _consume_json(self, message):
        """ Convert our AMQP messages into a consistent dictionary format.

        This method exists because our STOMP & AMQP message brokers consume
        messages in different formats.  This causes our messaging abstraction
        to leak into the consumers themselves.

        :Note: We do not pass the message headers to the consumer (in this AMQP consumer)
        because the current AMQP.js bindings do not allow the client to change them.
        Thus, we need to throw any topic/queue details into the JSON body itself.
        """
        try:
            body = json.loads(message.body)
        except:
            log.debug("Unable to decode message body to JSON: %r" % message.body)
            body = message.body
        topic = None

        # Try some stuff for AMQP:
        try:
            topic = message.headers[0].routing_key
        except TypeError:
            # We didn't get a JSON dictionary
            pass
        except AttributeError:
            # We didn't get headers or a routing key?
            pass

        # If that didn't work, it might be zeromq
        if not topic:
            try:
                topic = message.topic
            except AttributeError:
                # Weird.  I have no idea...
                pass

        try:
            message_as_dict = {'body': body, 'topic': topic}
            self._consume(message_as_dict)
            self._exception_count = 0  # Reset if everything went swimmingly
        except Exception:
            # Otherwise, keep track of how many exceptions we've hit in a row
            self._exception_count = self._exception_count + 1
            # And then re-raise the exception to be logged
            raise

    def _consume(self, message):
        self.incoming.put(message)

    def _work(self):
        while True:
            # This is a blocking call.  It waits until a message is available.
            message = self.incoming.get()

            # Then we are being asked to quit
            if message is StopIteration:
                break

            self.debug("Worker thread picking up %r" % message)
            try:
                self.validate(message)
            except Exception, e:
                log.warn("Received invalid message %r" % e)
                return

            self.consume(message)
            self.debug("Going back to waiting on the incoming queue.")

        self.debug("Worker thread exiting.")

    def validate(self, message):
        """ Override to implement your own validation scheme. """
        pass

    def consume(self, message):
        raise NotImplementedError

    def send_message(self, topic, message):
        try:
            self.hub.send_message(topic, message)
        except Exception, e:
            log.error('Cannot send message: %s' % e)

    def stop(self):
        for i in range(getattr(self, 'N', 0)):
            self.incoming.put(StopIteration)

        if hasattr(self, 'hub'):
            self.hub.close()

        if getattr(self, 'DBSession', None):
            self.DBSession.close()
<|endoftext|>


<commit_before># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import contextlib
import httplib
import json
import logging
import pprint
import re
import socket
import sys
import urllib2

from telemetry import decorators
from telemetry.core import exceptions
from telemetry.core import forwarders
from telemetry.core import user_agent
from telemetry.core import util
from telemetry.core import web_contents
from telemetry.core import wpr_modes
from telemetry.core import wpr_server
from telemetry.core.backends import browser_backend
from telemetry.core.backends.chrome import extension_backend
from telemetry.core.backends.chrome import system_info_backend
from telemetry.core.backends.chrome import tab_list_backend
from telemetry.core.backends.chrome import tracing_backend
from telemetry.timeline import tracing_timeline_data
from telemetry.unittest import options_for_unittests

class ChromeBrowserBackend(browser_backend.BrowserBackend):
  """An abstract class for chrome browser backends. Provides basic functionality
  once a remote-debugger port has been established."""
  # It is OK to have abstract methods. pylint: disable=W0223

  def __init__(self, supports_tab_control, supports_extensions, browser_options,
               output_profile_path, extensions_to_load):
    super(ChromeBrowserBackend, self).__init__(
        supports_extensions=supports_extensions,
        browser_options=browser_options,
        tab_list_backend=tab_list_backend.TabListBackend)
    self._port = None

    self._supports_tab_control = supports_tab_control
    self._tracing_backend = None
    self._system_info_backend = None

    self._output_profile_path = output_profile_path
    self._extensions_to_load = extensions_to_load

    if browser_options.netsim:
      self.wpr_port_pairs = forwarders.PortPairs(
          http=forwarders.PortPair(80, 80),
          https=forwarders.PortPair(443, 443),
          dns=forwarders.PortPair(53, 53))
    else:
      self.wpr_port_pairs = forwarders.PortPairs(
          http=forwarders.PortPair(0, 0),
          https=forwarders.PortPair(0, 0),
          dns=None)

    if (self.browser_options.dont_override_profile and
        not options_for_unittests.AreSet()):
      sys.stderr.write('Warning: Not overriding profile. This can cause '
                       'unexpected effects due to profile-specific settings, '
                       'such as about:flags settings, cookies, and '
                       'extensions.\n')

  def AddReplayServerOptions(self, extra_wpr_args):
    if self.browser_options.netsim:
      extra_wpr_args.append('--net=%s' % self.browser_options.netsim)
    else:
      extra_wpr_args.append('--no-dns_forwarding')

  @property
  @decorators.Cache
  def extension_backend(self):
    if not self.supports_extensions:
      return None
    return extension_backend.ExtensionBackendDict(self)

  def GetBrowserStartupArgs(self):
    args = []
    args.extend(self.browser_options.extra_browser_args)
    args.append('--enable-net-benchmarking')
    args.append('--metrics-recording-only')
    args.append('--no-default-browser-check')
    args.append('--no-first-run')

    # Turn on GPU benchmarking extension for all runs. The only side effect of
    # the extension being on is that render stats are tracked. This is believed
    # to be effectively free. And, by doing so here, it avoids us having to
    # programmatically inspect a pageset's actions in order to determine if it
    # might eventually scroll.
    args.append('--enable-gpu-benchmarking')

    # Set --no-proxy-server to work around some XP issues unless
    # some other flag indicates a proxy is needed.
    if not '--enable-spdy-proxy-auth' in args:
      args.append('--no-proxy-server')

    if self.browser_options.disable_background_networking:
      args.append('--disable-background-networking')

    if self.browser_options.netsim:
      args.append('--ignore-certificate-errors')
    elif self.browser_options.wpr_mode != wpr_modes.WPR_OFF:
      args.extend(wpr_server.GetChromeFlags(self.forwarder_factory.host_ip,
                                            self.wpr_port_pairs))
    args.extend(user_agent.GetChromeUserAgentArgumentFromType(
        self.browser_options.browser_user_agent_type))

    extensions = [extension.local_path
                  for extension in self._extensions_to_load
                  if not extension.is_component]
    extension_str = ','.join(extensions)
    if len(extensions) > 0:
      args.append('--load-extension=%s' % extension_str)

    component_extensions = [extension.local_path
                            for extension in self._extensions_to_load
                            if extension.is_component]
    component_extension_str = ','.join(component_extensions)
    if len(component_extensions) > 0:
      args.append('--load-component-extension=%s' % component_extension_str)

    if self.browser_options.no_proxy_server:
      args.append('--no-proxy-server')

    if self.browser_options.disable_component_extensions_with_background_pages:
      args.append('--disable-component-extensions-with-background-pages')

    return args

  def HasBrowserFinishedLaunching(self):
    try:
      self.Request('', timeout=.1)
    except (exceptions.BrowserGoneException,
            exceptions.BrowserConnectionGoneException):
      return False
    else:
      return True

  def _WaitForBrowserToComeUp(self, wait_for_extensions=True):
    try:
      util.WaitFor(self.HasBrowserFinishedLaunching, timeout=30)
    except (util.TimeoutException, exceptions.ProcessGoneException) as e:
      if not self.IsBrowserRunning():
        raise exceptions.BrowserGoneException(self.browser, e)
      raise exceptions.BrowserConnectionGoneException(self.browser, e)

    def AllExtensionsLoaded():
      # Extension pages are loaded from an about:blank page,
      # so we need to check that the document URL is the extension
      # page in addition to the ready state.
      extension_ready_js = """
          document.URL.lastIndexOf('chrome-extension://%s/', 0) == 0 &&
          (document.readyState == 'complete' ||
           document.readyState == 'interactive')
      """
      for e in self._extensions_to_load:
        try:
          extension_objects = self.extension_backend[e.extension_id]
        except KeyError:
          return False
        for extension_object in extension_objects:
          try:
            res = extension_object.EvaluateJavaScript(
                extension_ready_js % e.extension_id)
          except exceptions.EvaluateException:
            # If the inspected page is not ready, we will get an error
            # when we evaluate a JS expression, but we can just keep polling
            # until the page is ready (crbug.com/251913).
            res = None

          # TODO(tengs): We don't have full support for getting the Chrome
          # version before launch, so for now we use a generic workaround to
          # check for an extension binding bug in old versions of Chrome.
          # See crbug.com/263162 for details.
          if res and extension_object.EvaluateJavaScript(
              'chrome.runtime == null'):
            extension_object.Reload()
          if not res:
            return False
      return True

    if wait_for_extensions and self._supports_extensions:
      try:
        util.WaitFor(AllExtensionsLoaded, timeout=60)
      except util.TimeoutException:
        logging.error('ExtensionsToLoad: ' +
            repr([e.extension_id for e in self._extensions_to_load]))
        logging.error('Extension list: ' +
            pprint.pformat(self.extension_backend, indent=4))
        raise

  def ListInspectableContexts(self):
    return json.loads(self.Request(''))

  def Request(self, path, timeout=30, throw_network_exception=False):
    url = 'http://127.0.0.1:%i/json' % self._port
    if path:
      url += '/' + path
    try:
      proxy_handler = urllib2.ProxyHandler({})  # Bypass any system proxy.
      opener = urllib2.build_opener(proxy_handler)
      with contextlib.closing(opener.open(url, timeout=timeout)) as req:
        return req.read()
    except (socket.error, httplib.BadStatusLine, urllib2.URLError) as e:
      if throw_network_exception:
        raise e
      if not self.IsBrowserRunning():
        raise exceptions.BrowserGoneException(self.browser, e)
      raise exceptions.BrowserConnectionGoneException(self.browser, e)

  @property
  def browser_directory(self):
    raise NotImplementedError()

  @property
  def profile_directory(self):
    raise NotImplementedError()

  @property
  @decorators.Cache
  def chrome_branch_number(self):
    # Detect version information.
    data = self.Request('version')
    resp = json.loads(data)
    if 'Protocol-Version' in resp:
      if 'Browser' in resp:
        branch_number_match = re.search('Chrome/\d+\.\d+\.(\d+)\.\d+',
                                        resp['Browser'])
      else:
        branch_number_match = re.search(
            'Chrome/\d+\.\d+\.(\d+)\.\d+ (Mobile )?Safari',
            resp['User-Agent'])

      if branch_number_match:
        branch_number = int(branch_number_match.group(1))
        if branch_number:
          return branch_number

    # Branch number can't be determined, so fail any branch number checks.
    return 0

  @property
  def supports_tab_control(self):
    return self._supports_tab_control

  @property
  def supports_tracing(self):
    return True

  def StartTracing(self, trace_options, custom_categories=None,
                   timeout=web_contents.DEFAULT_WEB_CONTENTS_TIMEOUT):
    """
    Args:
        trace_options: An tracing_options.TracingOptions instance.
        custom_categories: An optional string containing a list of
                         comma separated categories that will be traced
                         instead of the default category set.  Example: use
                         "webkit,cc,disabled-by-default-cc.debug" to trace only
                         those three event categories.
    """
    assert trace_options and trace_options.enable_chrome_trace
    if self._tracing_backend is None:
      self._tracing_backend = tracing_backend.TracingBackend(self._port, self)
    return self._tracing_backend.StartTracing(
        trace_options, custom_categories, timeout)

  @property
  def is_tracing_running(self):
    if not self._tracing_backend:
      return None
    return self._tracing_backend.is_tracing_running

  def StopTracing(self):
    """ Stops tracing and returns the result as TimelineData object. """
    tab_ids_list = []
    for (i, _) in enumerate(self._browser.tabs):
      tab = self.tab_list_backend.Get(i, None)
      if tab:
        success = tab.EvaluateJavaScript(
            "console.time('" + tab.id + "');" +
            "console.timeEnd('" + tab.id + "');" +
            "console.time.toString().indexOf('[native code]') != -1;")
        if not success:
          raise Exception('Page stomped on console.time')
        tab_ids_list.append(tab.id)
    trace_events = self._tracing_backend.StopTracing()
    # Augment tab_ids data to trace events.
    event_data = {'traceEvents' : trace_events, 'tabIds': tab_ids_list}
    return tracing_timeline_data.TracingTimelineData(event_data)

  def GetProcessName(self, cmd_line):
    """Returns a user-friendly name for the process of the given |cmd_line|."""
    if not cmd_line:
      # TODO(tonyg): Eventually we should make all of these known and add an
      # assertion.
      return 'unknown'
    if 'nacl_helper_bootstrap' in cmd_line:
      return 'nacl_helper_bootstrap'
    if ':sandboxed_process' in cmd_line:
      return 'renderer'
    m = re.match(r'.* --type=([^\s]*) .*', cmd_line)
    if not m:
      return 'browser'
    return m.group(1)

  def Close(self):
    if self._tracing_backend:
      self._tracing_backend.Close()
      self._tracing_backend = None

  @property
  def supports_system_info(self):
    return self.GetSystemInfo() != None

  def GetSystemInfo(self):
    if self._system_info_backend is None:
      self._system_info_backend = system_info_backend.SystemInfoBackend(
          self._port)
    return self._system_info_backend.GetSystemInfo()
<commit_msg>[Telemetry] Increase timeout for HasBrowserFinishedLaunching<commit_after># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import contextlib
import httplib
import json
import logging
import pprint
import re
import socket
import sys
import urllib2

from telemetry import decorators
from telemetry.core import exceptions
from telemetry.core import forwarders
from telemetry.core import user_agent
from telemetry.core import util
from telemetry.core import web_contents
from telemetry.core import wpr_modes
from telemetry.core import wpr_server
from telemetry.core.backends import browser_backend
from telemetry.core.backends.chrome import extension_backend
from telemetry.core.backends.chrome import system_info_backend
from telemetry.core.backends.chrome import tab_list_backend
from telemetry.core.backends.chrome import tracing_backend
from telemetry.timeline import tracing_timeline_data
from telemetry.unittest import options_for_unittests

class ChromeBrowserBackend(browser_backend.BrowserBackend):
  """An abstract class for chrome browser backends. Provides basic functionality
  once a remote-debugger port has been established."""
  # It is OK to have abstract methods. pylint: disable=W0223

  def __init__(self, supports_tab_control, supports_extensions, browser_options,
               output_profile_path, extensions_to_load):
    super(ChromeBrowserBackend, self).__init__(
        supports_extensions=supports_extensions,
        browser_options=browser_options,
        tab_list_backend=tab_list_backend.TabListBackend)
    self._port = None

    self._supports_tab_control = supports_tab_control
    self._tracing_backend = None
    self._system_info_backend = None

    self._output_profile_path = output_profile_path
    self._extensions_to_load = extensions_to_load

    if browser_options.netsim:
      self.wpr_port_pairs = forwarders.PortPairs(
          http=forwarders.PortPair(80, 80),
          https=forwarders.PortPair(443, 443),
          dns=forwarders.PortPair(53, 53))
    else:
      self.wpr_port_pairs = forwarders.PortPairs(
          http=forwarders.PortPair(0, 0),
          https=forwarders.PortPair(0, 0),
          dns=None)

    if (self.browser_options.dont_override_profile and
        not options_for_unittests.AreSet()):
      sys.stderr.write('Warning: Not overriding profile. This can cause '
                       'unexpected effects due to profile-specific settings, '
                       'such as about:flags settings, cookies, and '
                       'extensions.\n')

  def AddReplayServerOptions(self, extra_wpr_args):
    if self.browser_options.netsim:
      extra_wpr_args.append('--net=%s' % self.browser_options.netsim)
    else:
      extra_wpr_args.append('--no-dns_forwarding')

  @property
  @decorators.Cache
  def extension_backend(self):
    if not self.supports_extensions:
      return None
    return extension_backend.ExtensionBackendDict(self)

  def GetBrowserStartupArgs(self):
    args = []
    args.extend(self.browser_options.extra_browser_args)
    args.append('--enable-net-benchmarking')
    args.append('--metrics-recording-only')
    args.append('--no-default-browser-check')
    args.append('--no-first-run')

    # Turn on GPU benchmarking extension for all runs. The only side effect of
    # the extension being on is that render stats are tracked. This is believed
    # to be effectively free. And, by doing so here, it avoids us having to
    # programmatically inspect a pageset's actions in order to determine if it
    # might eventually scroll.
    args.append('--enable-gpu-benchmarking')

    # Set --no-proxy-server to work around some XP issues unless
    # some other flag indicates a proxy is needed.
    if not '--enable-spdy-proxy-auth' in args:
      args.append('--no-proxy-server')

    if self.browser_options.disable_background_networking:
      args.append('--disable-background-networking')

    if self.browser_options.netsim:
      args.append('--ignore-certificate-errors')
    elif self.browser_options.wpr_mode != wpr_modes.WPR_OFF:
      args.extend(wpr_server.GetChromeFlags(self.forwarder_factory.host_ip,
                                            self.wpr_port_pairs))
    args.extend(user_agent.GetChromeUserAgentArgumentFromType(
        self.browser_options.browser_user_agent_type))

    extensions = [extension.local_path
                  for extension in self._extensions_to_load
                  if not extension.is_component]
    extension_str = ','.join(extensions)
    if len(extensions) > 0:
      args.append('--load-extension=%s' % extension_str)

    component_extensions = [extension.local_path
                            for extension in self._extensions_to_load
                            if extension.is_component]
    component_extension_str = ','.join(component_extensions)
    if len(component_extensions) > 0:
      args.append('--load-component-extension=%s' % component_extension_str)

    if self.browser_options.no_proxy_server:
      args.append('--no-proxy-server')

    if self.browser_options.disable_component_extensions_with_background_pages:
      args.append('--disable-component-extensions-with-background-pages')

    return args

  def HasBrowserFinishedLaunching(self):
    try:
      self.Request('', timeout=.1)
    except (exceptions.BrowserGoneException,
            exceptions.BrowserConnectionGoneException):
      return False
    else:
      return True

  def _WaitForBrowserToComeUp(self, wait_for_extensions=True):
    try:
      util.WaitFor(self.HasBrowserFinishedLaunching, timeout=50)
    except (util.TimeoutException, exceptions.ProcessGoneException) as e:
      if not self.IsBrowserRunning():
        raise exceptions.BrowserGoneException(self.browser, e)
      raise exceptions.BrowserConnectionGoneException(self.browser, e)

    def AllExtensionsLoaded():
      # Extension pages are loaded from an about:blank page,
      # so we need to check that the document URL is the extension
      # page in addition to the ready state.
      extension_ready_js = """
          document.URL.lastIndexOf('chrome-extension://%s/', 0) == 0 &&
          (document.readyState == 'complete' ||
           document.readyState == 'interactive')
      """
      for e in self._extensions_to_load:
        try:
          extension_objects = self.extension_backend[e.extension_id]
        except KeyError:
          return False
        for extension_object in extension_objects:
          try:
            res = extension_object.EvaluateJavaScript(
                extension_ready_js % e.extension_id)
          except exceptions.EvaluateException:
            # If the inspected page is not ready, we will get an error
            # when we evaluate a JS expression, but we can just keep polling
            # until the page is ready (crbug.com/251913).
            res = None

          # TODO(tengs): We don't have full support for getting the Chrome
          # version before launch, so for now we use a generic workaround to
          # check for an extension binding bug in old versions of Chrome.
          # See crbug.com/263162 for details.
          if res and extension_object.EvaluateJavaScript(
              'chrome.runtime == null'):
            extension_object.Reload()
          if not res:
            return False
      return True

    if wait_for_extensions and self._supports_extensions:
      try:
        util.WaitFor(AllExtensionsLoaded, timeout=60)
      except util.TimeoutException:
        logging.error('ExtensionsToLoad: ' +
            repr([e.extension_id for e in self._extensions_to_load]))
        logging.error('Extension list: ' +
            pprint.pformat(self.extension_backend, indent=4))
        raise

  def ListInspectableContexts(self):
    return json.loads(self.Request(''))

  def Request(self, path, timeout=30, throw_network_exception=False):
    url = 'http://127.0.0.1:%i/json' % self._port
    if path:
      url += '/' + path
    try:
      proxy_handler = urllib2.ProxyHandler({})  # Bypass any system proxy.
      opener = urllib2.build_opener(proxy_handler)
      with contextlib.closing(opener.open(url, timeout=timeout)) as req:
        return req.read()
    except (socket.error, httplib.BadStatusLine, urllib2.URLError) as e:
      if throw_network_exception:
        raise e
      if not self.IsBrowserRunning():
        raise exceptions.BrowserGoneException(self.browser, e)
      raise exceptions.BrowserConnectionGoneException(self.browser, e)

  @property
  def browser_directory(self):
    raise NotImplementedError()

  @property
  def profile_directory(self):
    raise NotImplementedError()

  @property
  @decorators.Cache
  def chrome_branch_number(self):
    # Detect version information.
    data = self.Request('version')
    resp = json.loads(data)
    if 'Protocol-Version' in resp:
      if 'Browser' in resp:
        branch_number_match = re.search('Chrome/\d+\.\d+\.(\d+)\.\d+',
                                        resp['Browser'])
      else:
        branch_number_match = re.search(
            'Chrome/\d+\.\d+\.(\d+)\.\d+ (Mobile )?Safari',
            resp['User-Agent'])

      if branch_number_match:
        branch_number = int(branch_number_match.group(1))
        if branch_number:
          return branch_number

    # Branch number can't be determined, so fail any branch number checks.
    return 0

  @property
  def supports_tab_control(self):
    return self._supports_tab_control

  @property
  def supports_tracing(self):
    return True

  def StartTracing(self, trace_options, custom_categories=None,
                   timeout=web_contents.DEFAULT_WEB_CONTENTS_TIMEOUT):
    """
    Args:
        trace_options: An tracing_options.TracingOptions instance.
        custom_categories: An optional string containing a list of
                         comma separated categories that will be traced
                         instead of the default category set.  Example: use
                         "webkit,cc,disabled-by-default-cc.debug" to trace only
                         those three event categories.
    """
    assert trace_options and trace_options.enable_chrome_trace
    if self._tracing_backend is None:
      self._tracing_backend = tracing_backend.TracingBackend(self._port, self)
    return self._tracing_backend.StartTracing(
        trace_options, custom_categories, timeout)

  @property
  def is_tracing_running(self):
    if not self._tracing_backend:
      return None
    return self._tracing_backend.is_tracing_running

  def StopTracing(self):
    """ Stops tracing and returns the result as TimelineData object. """
    tab_ids_list = []
    for (i, _) in enumerate(self._browser.tabs):
      tab = self.tab_list_backend.Get(i, None)
      if tab:
        success = tab.EvaluateJavaScript(
            "console.time('" + tab.id + "');" +
            "console.timeEnd('" + tab.id + "');" +
            "console.time.toString().indexOf('[native code]') != -1;")
        if not success:
          raise Exception('Page stomped on console.time')
        tab_ids_list.append(tab.id)
    trace_events = self._tracing_backend.StopTracing()
    # Augment tab_ids data to trace events.
    event_data = {'traceEvents' : trace_events, 'tabIds': tab_ids_list}
    return tracing_timeline_data.TracingTimelineData(event_data)

  def GetProcessName(self, cmd_line):
    """Returns a user-friendly name for the process of the given |cmd_line|."""
    if not cmd_line:
      # TODO(tonyg): Eventually we should make all of these known and add an
      # assertion.
      return 'unknown'
    if 'nacl_helper_bootstrap' in cmd_line:
      return 'nacl_helper_bootstrap'
    if ':sandboxed_process' in cmd_line:
      return 'renderer'
    m = re.match(r'.* --type=([^\s]*) .*', cmd_line)
    if not m:
      return 'browser'
    return m.group(1)

  def Close(self):
    if self._tracing_backend:
      self._tracing_backend.Close()
      self._tracing_backend = None

  @property
  def supports_system_info(self):
    return self.GetSystemInfo() != None

  def GetSystemInfo(self):
    if self._system_info_backend is None:
      self._system_info_backend = system_info_backend.SystemInfoBackend(
          self._port)
    return self._system_info_backend.GetSystemInfo()
<|endoftext|>


<commit_before>import APIFunctions
import argparse
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from Classes import Base

parser = argparse.ArgumentParser()
parser.add_argument("routes", help="The routes you want to capture in MBTA v2 format-- example: Red,Green-D,Orange")
args = vars(parser.parse_args())
routes = args['routes'].split(',')

# Make a database connection
db = create_engine('sqlite:///mbta.db', echo=False)
Base.metadata.create_all(db)
Session = sessionmaker(bind=db)
session = Session()

routes = APIFunctions.get_routes()
for route in routes:
    session.add(route)

# Get the station objects and add them to the database
stations = APIFunctions.get_stations(session)
for station in stations:
    session.add(station)

session.commit()<commit_msg>Initial setup no longer takes any arguments<commit_after>import APIFunctions
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from Classes import Base

# Make a database connection
db = create_engine('sqlite:///mbta.db', echo=False)
Base.metadata.create_all(db)
Session = sessionmaker(bind=db)
session = Session()

routes = APIFunctions.get_routes()
for route in routes:
    session.add(route)

# Get the station objects and add them to the database
stations = APIFunctions.get_stations(session)
for station in stations:
    session.add(station)

session.commit()<|endoftext|>


<commit_before># Copyright 2020 The TensorFlow Authors. 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.
# ==============================================================================
"""Analyzes the latency and numerics information of sample model inference."""

import itertools
import json
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import numpy as np

from tensorflow.python.compiler.tensorrt.model_tests import model_handler
import tensorflow.python.compiler.tensorrt.trt_convert as trt
from tensorflow.python.platform import tf_logging as logging

# pylint: disable=bad-whitespace
class DataFrame:
  """Lightweight immutable Dataframe similar to Pandas Dataframe."""

  def __init__(self,
               column_names: Sequence[str],
               rows: Sequence[Sequence[Any]] = None,
               columns: Sequence[Sequence[Any]] = None):
    self._column_names = column_names
    if not rows and not columns:
      raise ValueError("Cannot initialize with empty data!")
    self._rows = rows
    self._columns = columns

  @property
  def n_rows(self) -> int:
    return len(self._rows) if self._rows else len(self._columns[0])

  @property
  def n_columns(self) -> int:
    return len(self._columns) if self._columns else len(self._rows[0])

  @property
  def column_names(self) -> Sequence[str]:
    return self._column_names

  @property
  def rows(self) -> Sequence[Sequence[Any]]:
    return self._rows if self._rows else [
        [c[i] for c in self._columns] for i in range(len(self._columns[0]))
    ]

  @property
  def columns(self) -> Sequence[Sequence[Any]]:
    return self._columns if self._columns else [
        [r[i] for r in self._rows] for i in range(len(self._rows[0]))
    ]

  def __add__(self, other: "DataFrame") -> "DataFrame":
    if (not set(self.column_names).intersection(other.column_names) and
        len(self.rows) == len(other.rows)):
      return DataFrame(
          column_names=list(
              itertools.chain(self.column_names, other.column_names)),
          columns=list(itertools.chain(self.columns, other.columns)))
    if self.column_names == other.column_names:
      return DataFrame(
          column_names=self.column_names,
          rows=list(itertools.chain(self.rows, other.rows)))
    raise ValueError("Cannot combine two DataFrame")

  def __iadd__(self, other: "DataFrame") -> "DataFrame":
    tmp = self + other
    self._column_names = tmp._column_names
    self._rows, self._columns = tmp._rows, tmp._columns
    return self

  def __call__(self, r: int, c: Optional[Union[int, str]] = None) -> Any:
    if c is None:
      return dict(zip(self.column_names, self.rows[r]))
    c = self._column_names.index(c) if isinstance(c, str) else c
    return self._rows[r][c] if self._rows else self._columns[c][r]

  def __str__(self) -> str:
    return ",".join(self.column_names) + "\n" + "\n".join(",".join(
        "N/A" if v is None else str(v) for v in row) for row in self.rows)

  def to_csv(self, path: str):
    with open(path, "w") as file:
      file.write(str(self))

  def to_json(self, path: str):
    with open(path, "w") as file:
      json.dump([dict(zip(self.column_names, r)) for r in self.rows], file)

def extract_test_info(
    test_results: model_handler.TestResultCollection) -> DataFrame:
  """Extracts the test infomation."""
  column_names = list(
      itertools.chain(model_handler.ModelConfig._fields,
                      ["enable_gpu", "trt_model"],
                      trt.TrtConversionParams._fields))
  rows = []
  for result in test_results.results:
    r = list(result.model_config) + [result.enable_gpu]
    if result.trt_convert_params is not None:
      r += [True] + list(result.trt_convert_params)
    else:
      r += [False] + [None for _ in trt.TrtConversionParams._fields]
    rows.append(r)
  return DataFrame(column_names=column_names, rows=rows)

def analyze_test_latency(test_results: model_handler.TestResultCollection,
                         use_cpu_baseline: bool) -> DataFrame:
  """Analyzes test latency."""
  base_result = (
      test_results.cpu_base_result
      if use_cpu_baseline else test_results.gpu_base_result)
  if base_result is None:
    raise ValueError(
        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
  base_mean_time = np.asscalar(np.mean(base_result.model_latency))
  column_names = ["time(ms)", "speedup"]
  rows = []
  for result in test_results.results:
    mean_time = np.asscalar(np.mean(result.model_latency))
    rows.append([mean_time * 1000.0, base_mean_time / mean_time])
  return DataFrame(column_names=column_names, rows=rows)

def analyze_test_numerics(test_results: model_handler.TestResultCollection,
                          use_cpu_baseline: bool) -> DataFrame:
  """Analyzes test numerics."""
  preprocess_funcs = {
      "diff": lambda x, y: np.fabs(x - y),
      # Ensures dividends are not zero to avoid exceptions/NaNs.
      "rel_diff": lambda x, y: np.fabs(x - y) / np.fmax(np.fabs(y), 1.0e-6)
  }
  postprocess_funcs = {"mean": np.mean, "std": np.std}
  column_names = []
  columns = []
  base_result = (
      test_results.cpu_base_result
      if use_cpu_baseline else test_results.gpu_base_result)
  if base_result is None:
    raise ValueError(
        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
  for fn0, fn1 in itertools.product(preprocess_funcs, postprocess_funcs):
    func0, func1 = preprocess_funcs[fn0], postprocess_funcs[fn1]
    column_names.append("{}_{}".format(fn0, fn1))
    columns.append([])
    for result in test_results.results:
      columns[-1].append(dict())
      for idx, tensor in enumerate(result.output_tensors):
        name = base_result.output_names[idx]
        cpu_tensor = base_result.output_tensors[idx]
        metric_value = np.asscalar(func1(func0(tensor, cpu_tensor)))
        columns[-1][-1][name] = metric_value
  return DataFrame(column_names=column_names, columns=columns)

def check_column(df: DataFrame, name: str, fn: Callable[[float], bool]) -> bool:
  """Checks the values of a column using a custom function and logs abnormals.

  The check is only performed on TensorRT models, not native CPU/GPU models.

  Args:
    df: The DataFrame to be checked.
    name: The name of the column to be checked.
    fn: The function that takes a value of at the specified column and returns
      if the value statisfies the check.

  Returns:
    Whether all the values of the specified column satisfies the provided check.
  """
  is_ok = True
  for r in range(df.n_rows):
    if df(r, "trt_model"):
      if not fn(df(r, name)):
        logging.error("Unsatisfied %s found at: %s", name, df(r))
        is_ok = False
  return is_ok

class ResultAnalyzer:
  """Analyzes ModelHandlerManager results."""

  def __init__(
      self,
      use_cpu_latency_baseline: bool,
      use_cpu_numerics_baseline: bool,
      checkers: Sequence[Callable[[DataFrame], bool]],
  ):
    self._use_cpu_latency_baseline = use_cpu_latency_baseline
    self._use_cpu_numerics_baseline = use_cpu_numerics_baseline
    self._checkers = checkers

  def analysis(
      self, test_results: model_handler.TestResultCollection
  ) -> Tuple[DataFrame, Sequence[bool]]:
    df = extract_test_info(test_results)
    df += analyze_test_latency(test_results, self._use_cpu_latency_baseline)
    df += analyze_test_numerics(test_results, self._use_cpu_numerics_baseline)
    checks = [c(df) for c in self._checkers]
    return df, checks
<commit_msg>[NumPy] Replace numpy.asscalar(x) with x.item() in preparation for upgrading NumPy to 1.23.<commit_after># Copyright 2020 The TensorFlow Authors. 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.
# ==============================================================================
"""Analyzes the latency and numerics information of sample model inference."""

import itertools
import json
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import numpy as np

from tensorflow.python.compiler.tensorrt.model_tests import model_handler
import tensorflow.python.compiler.tensorrt.trt_convert as trt
from tensorflow.python.platform import tf_logging as logging

# pylint: disable=bad-whitespace
class DataFrame:
  """Lightweight immutable Dataframe similar to Pandas Dataframe."""

  def __init__(self,
               column_names: Sequence[str],
               rows: Sequence[Sequence[Any]] = None,
               columns: Sequence[Sequence[Any]] = None):
    self._column_names = column_names
    if not rows and not columns:
      raise ValueError("Cannot initialize with empty data!")
    self._rows = rows
    self._columns = columns

  @property
  def n_rows(self) -> int:
    return len(self._rows) if self._rows else len(self._columns[0])

  @property
  def n_columns(self) -> int:
    return len(self._columns) if self._columns else len(self._rows[0])

  @property
  def column_names(self) -> Sequence[str]:
    return self._column_names

  @property
  def rows(self) -> Sequence[Sequence[Any]]:
    return self._rows if self._rows else [
        [c[i] for c in self._columns] for i in range(len(self._columns[0]))
    ]

  @property
  def columns(self) -> Sequence[Sequence[Any]]:
    return self._columns if self._columns else [
        [r[i] for r in self._rows] for i in range(len(self._rows[0]))
    ]

  def __add__(self, other: "DataFrame") -> "DataFrame":
    if (not set(self.column_names).intersection(other.column_names) and
        len(self.rows) == len(other.rows)):
      return DataFrame(
          column_names=list(
              itertools.chain(self.column_names, other.column_names)),
          columns=list(itertools.chain(self.columns, other.columns)))
    if self.column_names == other.column_names:
      return DataFrame(
          column_names=self.column_names,
          rows=list(itertools.chain(self.rows, other.rows)))
    raise ValueError("Cannot combine two DataFrame")

  def __iadd__(self, other: "DataFrame") -> "DataFrame":
    tmp = self + other
    self._column_names = tmp._column_names
    self._rows, self._columns = tmp._rows, tmp._columns
    return self

  def __call__(self, r: int, c: Optional[Union[int, str]] = None) -> Any:
    if c is None:
      return dict(zip(self.column_names, self.rows[r]))
    c = self._column_names.index(c) if isinstance(c, str) else c
    return self._rows[r][c] if self._rows else self._columns[c][r]

  def __str__(self) -> str:
    return ",".join(self.column_names) + "\n" + "\n".join(",".join(
        "N/A" if v is None else str(v) for v in row) for row in self.rows)

  def to_csv(self, path: str):
    with open(path, "w") as file:
      file.write(str(self))

  def to_json(self, path: str):
    with open(path, "w") as file:
      json.dump([dict(zip(self.column_names, r)) for r in self.rows], file)

def extract_test_info(
    test_results: model_handler.TestResultCollection) -> DataFrame:
  """Extracts the test infomation."""
  column_names = list(
      itertools.chain(model_handler.ModelConfig._fields,
                      ["enable_gpu", "trt_model"],
                      trt.TrtConversionParams._fields))
  rows = []
  for result in test_results.results:
    r = list(result.model_config) + [result.enable_gpu]
    if result.trt_convert_params is not None:
      r += [True] + list(result.trt_convert_params)
    else:
      r += [False] + [None for _ in trt.TrtConversionParams._fields]
    rows.append(r)
  return DataFrame(column_names=column_names, rows=rows)

def analyze_test_latency(test_results: model_handler.TestResultCollection,
                         use_cpu_baseline: bool) -> DataFrame:
  """Analyzes test latency."""
  base_result = (
      test_results.cpu_base_result
      if use_cpu_baseline else test_results.gpu_base_result)
  if base_result is None:
    raise ValueError(
        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
  base_mean_time = np.mean(base_result.model_latency).item()
  column_names = ["time(ms)", "speedup"]
  rows = []
  for result in test_results.results:
    mean_time = np.mean(result.model_latency).item()
    rows.append([mean_time * 1000.0, base_mean_time / mean_time])
  return DataFrame(column_names=column_names, rows=rows)

def analyze_test_numerics(test_results: model_handler.TestResultCollection,
                          use_cpu_baseline: bool) -> DataFrame:
  """Analyzes test numerics."""
  preprocess_funcs = {
      "diff": lambda x, y: np.fabs(x - y),
      # Ensures dividends are not zero to avoid exceptions/NaNs.
      "rel_diff": lambda x, y: np.fabs(x - y) / np.fmax(np.fabs(y), 1.0e-6)
  }
  postprocess_funcs = {"mean": np.mean, "std": np.std}
  column_names = []
  columns = []
  base_result = (
      test_results.cpu_base_result
      if use_cpu_baseline else test_results.gpu_base_result)
  if base_result is None:
    raise ValueError(
        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
  for fn0, fn1 in itertools.product(preprocess_funcs, postprocess_funcs):
    func0, func1 = preprocess_funcs[fn0], postprocess_funcs[fn1]
    column_names.append("{}_{}".format(fn0, fn1))
    columns.append([])
    for result in test_results.results:
      columns[-1].append(dict())
      for idx, tensor in enumerate(result.output_tensors):
        name = base_result.output_names[idx]
        cpu_tensor = base_result.output_tensors[idx]
        metric_value = func1(func0(tensor, cpu_tensor)).item()
        columns[-1][-1][name] = metric_value
  return DataFrame(column_names=column_names, columns=columns)

def check_column(df: DataFrame, name: str, fn: Callable[[float], bool]) -> bool:
  """Checks the values of a column using a custom function and logs abnormals.

  The check is only performed on TensorRT models, not native CPU/GPU models.

  Args:
    df: The DataFrame to be checked.
    name: The name of the column to be checked.
    fn: The function that takes a value of at the specified column and returns
      if the value statisfies the check.

  Returns:
    Whether all the values of the specified column satisfies the provided check.
  """
  is_ok = True
  for r in range(df.n_rows):
    if df(r, "trt_model"):
      if not fn(df(r, name)):
        logging.error("Unsatisfied %s found at: %s", name, df(r))
        is_ok = False
  return is_ok

class ResultAnalyzer:
  """Analyzes ModelHandlerManager results."""

  def __init__(
      self,
      use_cpu_latency_baseline: bool,
      use_cpu_numerics_baseline: bool,
      checkers: Sequence[Callable[[DataFrame], bool]],
  ):
    self._use_cpu_latency_baseline = use_cpu_latency_baseline
    self._use_cpu_numerics_baseline = use_cpu_numerics_baseline
    self._checkers = checkers

  def analysis(
      self, test_results: model_handler.TestResultCollection
  ) -> Tuple[DataFrame, Sequence[bool]]:
    df = extract_test_info(test_results)
    df += analyze_test_latency(test_results, self._use_cpu_latency_baseline)
    df += analyze_test_numerics(test_results, self._use_cpu_numerics_baseline)
    checks = [c(df) for c in self._checkers]
    return df, checks
<|endoftext|>


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

import os
import ConfigParser
import time
import subprocess

import alfred

def b2d_exec(command, args=None, binary_path='/usr/local/bin/boot2docker'):
    if not args:
        args = []
    cmdline = ' '.join([binary_path, command] + args)
    ret = subprocess.check_output('{}; exit 0'.format(cmdline), shell=True)
    return ret

class Boot2dockerWorkflow(alfred.AlfredWorkflow):
    _reserved_words = []

    def __init__(self, max_results=20):
        self.max_results = max_results

    def command_autocomplete_iter(self, query):
        status = self.status
        if status == 'unknown':
            yield self.item(title='Boot2Docker is in an unknown state',
                            description='Try to run boot2docker in your '
                            'terminal to investigate', ignore=True)
        elif status == 'notexists':
            yield self.item(title='init', description='Initialize Boot2Docker',
                            autocomplete=True, arg='init', match=query)
        elif status == 'running':
            yield self.item(title='stop', description='Stop Boot2Docker',
                            autocomplete=True, arg='stop', match=query)
            yield self.item(title='suspend', description='Suspend Boot2Docker',
                            autocomplete=True, arg='suspend', match=query)
            yield self.item(title='restart', description='Restart Boot2Docker',
                            autocomplete=True, arg='restart', match=query)
        elif status in ('stopped', 'aborted', 'paused', 'suspended'):
            yield self.item(title='start', description='Start Boot2Docker',
                            autocomplete=True, arg='start', match=query)

    def do_command_autocomplete(self, query):
        self.write_items(self.command_autocomplete_iter(query))

    def do_command_notify(self, query=None):
        if query == 'start':
            self.write_text('Starting... you will get notified when done')
        elif query == 'stop':
            self.write_text('Stopping... you will get notified when done')
        elif query == 'suspend':
            self.write_text('Suspending... you will get notified when done')
        elif query == 'restart':
            self.write_text('Restarting... you will get notified when done')

    def do_command_run(self, query=None):
        command = query
        return self.route_action(command, '')

    def do_start(self, query=None):
        if self.status == 'running':
            return
        ret = b2d_exec('start')
        if ret.find('Started.') > -1:
            self.write_text('Boot2Docker is started.')
        else:
            self.write_text('An error has occured when starting.\n'
                            'please run boot2docker from command line')

    def do_restart(self, query=None):
        if self.status == 'restart':
            return
        ret = b2d_exec('restart')
        if ret.find('Started.') > -1:
            self.write_text('Boot2Docker is started.')
        else:
            self.write_text('An error has occured when restarting.\n'
                            'please run boot2docker from command line')

    def do_stop(self, query=None):
        if self.status != 'running':
            return
        ret = b2d_exec('stop')
        if ret.find('Shutting down') > -1:
            self.write('Boot2Docker is stopped')
        elif ret.find('is not running.') > -1:
            self.write('Boot2Docker was not running')
        else:
            self.write('Error while trying to stop instance')

    def do_suspend(self, query=None):
        if self.status != 'running':
            return
        ret = b2d_exec('suspend')
        if ret.find('100%') > -1:
            self.write('Boot2Docker is suspended')

    def do_init(self, query=None):
        if self.status != 'notexists':
            return
        ret = b2d_exec('init')
        if ret.find('You can now type boot2docker up') > -1:
            self.write('Boot2Docker is initialized')

    @property
    def status(self):
        ret = b2d_exec('status')
        if ret.find('running') > -1:
            return 'running'
        elif ret.find('stopped') > -1:
            return 'stopped'
        elif ret.find('aborted') > -1:
            return 'aborted'
        elif ret.find('suspended') > -1:
            return 'suspended'
        elif ret.find('paused') > -1:
            return 'paused'
        elif ret.find('does not exist') > -1:
            return 'notexists'
        else:
            return 'unknown'

    def do_status(self, query=None):
        status = self.status
        if status == 'running':
            self.write_text('Boot2Docker is running')
        elif status == 'stopped':
            self.write_text('Boot2Docker is stopped')
        elif status == 'aborted':
            self.write_text('Boot2Docker is aborted')
        elif status == 'suspended':
            self.write_text('Boot2Docker is suspended')
        elif status == 'paused':
            self.write_text('Boot2Docker is paused')
        elif status == 'notexists':
            self.write_text('Boot2Docker VM does not exist')
        else:
            self.write_text('Boot2Docker is in an unknown state')

def main(action, query):
    boot2docker = Boot2dockerWorkflow()
    boot2docker.route_action(action, query)

if __name__ == "__main__":
    main(action=alfred.args()[0], query=alfred.args()[1])
<commit_msg>updated notification messages<commit_after># -*- coding: utf-8 -*-

import os
import ConfigParser
import time
import subprocess

import alfred

def b2d_exec(command, args=None, binary_path='/usr/local/bin/boot2docker'):
    if not args:
        args = []
    cmdline = ' '.join([binary_path, command] + args)
    ret = subprocess.check_output('{}; exit 0'.format(cmdline), shell=True)
    return ret

class Boot2dockerWorkflow(alfred.AlfredWorkflow):
    _reserved_words = []

    def __init__(self, max_results=20):
        self.max_results = max_results

    def command_autocomplete_iter(self, query):
        status = self.status
        if status == 'unknown':
            yield self.item(title='Boot2Docker is in an unknown state',
                            description='Try to run boot2docker in your '
                            'terminal to investigate', ignore=True)
        elif status == 'notexists':
            yield self.item(title='init', description='Initialize Boot2Docker',
                            autocomplete=True, arg='init', match=query)
        elif status == 'running':
            yield self.item(title='stop', description='Stop Boot2Docker',
                            autocomplete=True, arg='stop', match=query)
            yield self.item(title='suspend', description='Suspend Boot2Docker',
                            autocomplete=True, arg='suspend', match=query)
            yield self.item(title='restart', description='Restart Boot2Docker',
                            autocomplete=True, arg='restart', match=query)
        elif status in ('stopped', 'aborted', 'paused', 'suspended'):
            yield self.item(title='start', description='Start Boot2Docker',
                            autocomplete=True, arg='start', match=query)

    def do_command_autocomplete(self, query):
        self.write_items(self.command_autocomplete_iter(query))

    def do_command_notify(self, query=None):
        if query == 'start':
            self.write_text('Starting...\nyou will get notified when done.')
        elif query == 'stop':
            self.write_text('Stopping...\nyou will get notified when done.')
        elif query == 'suspend':
            self.write_text('Suspending...\nyou will get notified when done.')
        elif query == 'restart':
            self.write_text('Restarting...\nyou will get notified when done.')
        elif query == 'init':
            self.write_text('Initializing...\nyou will get notified when done.')

    def do_command_run(self, query=None):
        command = query
        return self.route_action(command, '')

    def do_start(self, query=None):
        if self.status == 'running':
            return
        ret = b2d_exec('start')
        if ret.find('Started.') > -1:
            self.write_text('Boot2Docker is started.')
        else:
            self.write_text('An error has occured when starting.\n'
                            'please run boot2docker from command line')

    def do_restart(self, query=None):
        if self.status == 'restart':
            return
        ret = b2d_exec('restart')
        if ret.find('Started.') > -1:
            self.write_text('Boot2Docker is started.')
        else:
            self.write_text('An error has occured when restarting.\n'
                            'please run boot2docker from command line')

    def do_stop(self, query=None):
        if self.status != 'running':
            return
        ret = b2d_exec('stop')
        if ret.find('Shutting down') > -1:
            self.write('Boot2Docker is stopped')
        elif ret.find('is not running.') > -1:
            self.write('Boot2Docker was not running')
        else:
            self.write('Error while trying to stop instance')

    def do_suspend(self, query=None):
        if self.status != 'running':
            return
        ret = b2d_exec('suspend')
        if ret.find('100%') > -1:
            self.write('Boot2Docker is suspended')

    def do_init(self, query=None):
        if self.status != 'notexists':
            return
        ret = b2d_exec('init')
        if ret.find('You can now type boot2docker up') > -1:
            self.write('Boot2Docker is initialized')

    @property
    def status(self):
        ret = b2d_exec('status')
        if ret.find('running') > -1:
            return 'running'
        elif ret.find('stopped') > -1:
            return 'stopped'
        elif ret.find('aborted') > -1:
            return 'aborted'
        elif ret.find('suspended') > -1:
            return 'suspended'
        elif ret.find('paused') > -1:
            return 'paused'
        elif ret.find('does not exist') > -1:
            return 'notexists'
        else:
            return 'unknown'

    def do_status(self, query=None):
        status = self.status
        if status == 'running':
            self.write_text('Boot2Docker is running')
        elif status == 'stopped':
            self.write_text('Boot2Docker is stopped')
        elif status == 'aborted':
            self.write_text('Boot2Docker is aborted')
        elif status == 'suspended':
            self.write_text('Boot2Docker is suspended')
        elif status == 'paused':
            self.write_text('Boot2Docker is paused')
        elif status == 'notexists':
            self.write_text('Boot2Docker VM does not exist')
        else:
            self.write_text('Boot2Docker is in an unknown state')

def main(action, query):
    boot2docker = Boot2dockerWorkflow()
    boot2docker.route_action(action, query)

if __name__ == "__main__":
    main(action=alfred.args()[0], query=alfred.args()[1])
<|endoftext|>


<commit_before>import logging
import re

from django.contrib.auth import get_user_model
from django.core.management.base import BaseCommand

from chat.models import ChatMessage
from notifications import sms
from utils.datastructures import flatten

logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
# NOTE: This will pretty aggressively send out notifications at all times of
#       of the day or night if there are ChatMessages. To combat that, we
#       should schedule this command during certain hours only, e.g.
#
#       Something like:
#
#           Every 15 minutes between the hours of 8am and 18 (6pm) CST.
#           http://corntab.com/?c=*/15_8-18_*_*_*
#
# -----------------------------------------------------------------------------

class Command(BaseCommand):
    help = (
        "Sends a notification (currently SMS) to a user who's received "
        "'recent' chat messages. Schedule this to run frequently."
    )

    def handle(self, *args, **options):
        User = get_user_model()

        # Get all the users that probably have messages.
        # (we don't save the recipient's ID in the message so this is a hack)
        user_ids = flatten(
            room.split('-')[1:] for room in
            ChatMessage.objects.recent().values_list('room', flat=True)
        )

        # filter out those who don't have a listed phone number
        users = User.objects.filter(id__in=user_ids, userprofile__phone__gt='')

        # Now make sure we have them in a numbers-only format.
        numbers = users.values_list('userprofile__phone', flat=True)
        message = (
            "You have new messages!\n"
            "View them at http://officehours.tndata.org"
        )
        sms.mass_send(numbers, message, topic_name="ChatNotifications")
<commit_msg>only send notifications about _unread_ messages<commit_after>import logging
import re

from django.contrib.auth import get_user_model
from django.core.management.base import BaseCommand

from chat.models import ChatMessage
from notifications import sms
from utils.datastructures import flatten

logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
# NOTE: This will pretty aggressively send out notifications at all times of
#       of the day or night if there are ChatMessages. To combat that, we
#       should schedule this command during certain hours only, e.g.
#
#       Something like:
#
#           Every 15 minutes between the hours of 8am and 18 (6pm) CST.
#           http://corntab.com/?c=*/15_8-18_*_*_*
#
# -----------------------------------------------------------------------------

class Command(BaseCommand):
    help = (
        "Sends a notification (currently SMS) to a user who's received "
        "'recent' chat messages. Schedule this to run frequently."
    )

    def handle(self, *args, **options):
        User = get_user_model()

        # Get all the users that probably have messages.
        # (we don't save the recipient's ID in the message so this is a hack)
        messages = ChatMessage.objects.recent().filter(read=False)
        messages = messages.values_list('room', flat=True)
        user_ids = flatten(room.split('-')[1:] for room in messages)

        # filter out those who don't have a listed phone number
        users = User.objects.filter(id__in=user_ids, userprofile__phone__gt='')

        # Now make sure we have them in a numbers-only format.
        numbers = users.values_list('userprofile__phone', flat=True)
        message = (
            "You have new messages!\n"
            "View them at http://officehours.tndata.org"
        )
        sms.mass_send(numbers, message, topic_name="ChatNotifications")
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
from __future__ import unicode_literals

import json
from datetime import date, datetime, timedelta

from django.db.models import Count
from django.http import HttpResponse, HttpResponseNotAllowed
from django.shortcuts import get_object_or_404
from django.utils.datastructures import SortedDict
from django.views.generic import DetailView, TemplateView, ListView

from .forms import PeriodForm
from .models import (Section, Student, Corporation, CorpContact, Period,
    Training, Referent, Availability)

def school_year_start():
    """ Return first official day of current school year """
    current_year = date.today().year
    if date(current_year, 8, 1) > date.today():
        return date(current_year-1, 8, 1)
    else:
        return date(current_year, 8, 1)

class CorporationListView(ListView):
    model = Corporation
    template_name = 'corporations.html'

class CorporationView(DetailView):
    model = Corporation
    template_name = 'corporation.html'
    context_object_name = 'corp'

    def get_context_data(self, **kwargs):
        context = super(CorporationView, self).get_context_data(**kwargs)
        # Create a structure like:
        #   {'2011-2012': {'avails': [avail1, avail2, ...], 'stats': {'fil': num}},
        #    '2012-2013': ...}
        school_years = SortedDict()
        for av in Availability.objects.filter(corporation=self.object
                ).select_related('training__student__klass', 'period__section'
                ).order_by('period__start_date'):
            if av.period.school_year not in school_years:
                school_years[av.period.school_year] = {'avails': [], 'stats': {}}
            school_years[av.period.school_year]['avails'].append(av)
            if av.period.section.name not in school_years[av.period.school_year]['stats']:
                school_years[av.period.school_year]['stats'][av.period.section.name] = 0
            try:
                av.training
                # Only add to stats if training exists
                school_years[av.period.school_year]['stats'][av.period.section.name] += av.period.weeks
            except Training.DoesNotExist:
                pass

        context['years'] = school_years
        return context

class AttributionView(TemplateView):
    """
    Base view for the attribution screen. Populate sections and referents.
    All other data are retrieved through AJAX requests:
      * training periods: section_period
      * corp. availabilities for current period: period_availabilities
      * already planned training for current period: TrainingsByPeriodView
      * student list targetted by current period: period_students
    When an availability is chosen:
      * corp. contact list: CorpContactJSONView
    When a student is chosen;
      * details of a student: StudentSummaryView
    """
    template_name = 'attribution.html'

    def get_context_data(self, **kwargs):
        context = super(AttributionView, self).get_context_data(**kwargs)
        # Need 2 queries, because referents with no training item would not appear in the second query
        referents = Referent.objects.all().order_by('last_name', 'first_name')

        # Populate each referent with the number of referencies done during the current school year
        ref_counts = dict([(ref.id, ref.num_refs)
                for ref in Referent.objects.filter(training__availability__period__end_date__gte=school_year_start
                ).annotate(num_refs=Count('training'))])
        for ref in referents:
            ref.num_refs = ref_counts.get(ref.id, 0)

        context.update({
            #'period_form': PeriodForm(),
            'sections': Section.objects.all(),
            'referents': referents,
        })
        return context

class StudentSummaryView(DetailView):
    model = Student
    template_name = 'student_summary.html'

    def get_context_data(self, **kwargs):
        context = super(StudentSummaryView, self).get_context_data(**kwargs)
        context['previous_stages'] = self.object.training_set.all(
            ).select_related('availability__corporation').order_by('availability__period__end_date')
        return context

class AvailabilitySummaryView(DetailView):
    model = Availability
    template_name = 'availability_summary.html'

class TrainingsByPeriodView(ListView):
    template_name = 'trainings_list.html'
    context_object_name = 'trainings'

    def get_queryset(self):
        return Training.objects.select_related('student__klass', 'availability__corporation', 'availability__domain'
            ).filter(availability__period__pk=self.kwargs['pk']
            ).order_by('student__last_name', 'student__first_name')

class CorpContactJSONView(ListView):
    """ Return all contacts from a given corporation """
    return_fields = ['id', 'first_name', 'last_name', 'role', 'is_main']

    def get_queryset(self):
        return CorpContact.objects.filter(corporation__pk=self.kwargs['pk'])

    def render_to_response(self, context):
        serialized = [dict([(field, getattr(obj, field)) for field in self.return_fields])
                      for obj in context['object_list']]
        return HttpResponse(json.dumps(serialized), content_type="application/json")

# AJAX views:

def section_periods(request, pk):
    """ Return all periods (until 2 years ago) from a section (JSON) """
    section = get_object_or_404(Section, pk=pk)
    two_years_ago = datetime.now() - timedelta(days=365 * 2)
    periods = [{'id': p.id, 'dates': p.dates, 'title': p.title}
               for p in section.period_set.filter(start_date__gt=two_years_ago).order_by('-start_date')]
    return HttpResponse(json.dumps(periods), content_type="application/json")

def section_classes(request, pk):
    section = get_object_or_404(Section, pk=pk)
    classes = [(k.id, k.name) for k in section.klass_set.all()]
    return HttpResponse(json.dumps(classes), content_type="application/json")

def period_students(request, pk):
    """
    Return all active students from period's section and level,
    with corresponding Training if existing (JSON)
    """
    period = get_object_or_404(Period, pk=pk)
    students = Student.objects.filter(
        archived=False, klass__section=period.section, klass__level=period.relative_level
        ).order_by('last_name')
    trainings = dict((t.student_id, t.id) for t in Training.objects.filter(availability__period=period))
    data = [{
        'name': unicode(s),
        'id': s.id,
        'training_id': trainings.get(s.id),
        'klass': s.klass.name} for s in students]
    return HttpResponse(json.dumps(data), content_type="application/json")

def period_availabilities(request, pk):
    """ Return all availabilities in the specified period """
    period = get_object_or_404(Period, pk=pk)
    corps = [{'id': av.id, 'id_corp': av.corporation.id, 'corp_name': av.corporation.name,
              'domain': av.domain.name, 'free': av.free}
             for av in period.availability_set.select_related('corporation').all()]
    return HttpResponse(json.dumps(corps), content_type="application/json")

def new_training(request):
    if request.method != 'POST':
        return HttpResponseNotAllowed()
    ref_key = request.POST.get('referent')
    cont_key = request.POST.get('contact')
    try:
        ref = Referent.objects.get(pk=ref_key) if ref_key else None
        contact = CorpContact.objects.get(pk=cont_key) if cont_key else None
        avail = Availability.objects.get(pk=request.POST.get('avail'))
        training = Training.objects.create(
            student=Student.objects.get(pk=request.POST.get('student')),
            availability=avail,
            referent=ref,
        )
        if avail.contact != contact:
            avail.contact = contact
            avail.save()
    except Exception as exc:
        return HttpResponse(str(exc))
    return HttpResponse('OK')

def del_training(request):
    """ Delete training and return the referent id """
    if request.method != 'POST':
        return HttpResponseNotAllowed()
    training = get_object_or_404(Training, pk=request.POST.get('pk'))
    ref_id = training.referent_id
    training.delete()
    return HttpResponse(json.dumps({'ref_id': ref_id}), content_type="application/json")

EXPORT_FIELDS = [
    ('Prénom', 'student__first_name'), ('Nom', 'student__last_name'),
    ('ID externe', 'student__ext_id'),
    ('Classe', 'student__klass__name'), ('Filière', 'student__klass__section__name'),
    ('Nom du stage', 'availability__period__title'),
    ('Début', 'availability__period__start_date'), ('Fin', 'availability__period__end_date'),
    ('Remarques stage', 'comment'),
    ('Prénom référent', 'referent__first_name'), ('Nom référent', 'referent__last_name'),
    ('Courriel référent', 'referent__email'),
    ('Institution', 'availability__corporation__name'),
    ('ID externe Inst', 'availability__corporation__ext_id'),
    ('Rue Inst', 'availability__corporation__street'),
    ('NPA Inst', 'availability__corporation__pcode'),
    ('Ville Inst', 'availability__corporation__city'),
    ('Tél Inst', 'availability__corporation__tel'),
    ('Domaine', 'availability__domain__name'),
    ('Remarques Inst', 'availability__comment'),
    ('Civilité contact', 'availability__contact__title'),
    ('Prénom contact', 'availability__contact__first_name'),
    ('Nom contact', 'availability__contact__last_name'),
    ('ID externe contact', 'availability__contact__ext_id'),
    ('Tél contact', 'availability__contact__tel'),
    ('Courriel contact', 'availability__contact__email'),
    ('Courriel contact - copie', None),
]

NON_ATTR_EXPORT_FIELDS = [
    ('Filière', 'period__section__name'),
    ('Nom du stage', 'period__title'),
    ('Début', 'period__start_date'), ('Fin', 'period__end_date'),
    ('Institution', 'corporation__name'),
    ('Rue Inst', 'corporation__street'),
    ('NPA Inst', 'corporation__pcode'),
    ('Ville Inst', 'corporation__city'),
    ('Tél Inst', 'corporation__tel'),
    ('Domaine', 'domain__name'),
    ('Remarques Inst', 'comment'),
    ('Civilité contact', 'contact__title'),
    ('Prénom contact', 'contact__first_name'),
    ('Nom contact', 'contact__last_name'),
    ('Tél contact', 'contact__tel'),
    ('Courriel contact', 'contact__email'),
    ('Courriel contact - copie', None),
]

def stages_export(request, scope=None):
    from datetime import date
    from openpyxl import Workbook
    from openpyxl.styles import Font, Style
    from openpyxl.writer.excel import save_virtual_workbook

    period_filter = request.GET.get('period')
    non_attributed = bool(int(request.GET.get('non_attr', 0)))

    export_fields = EXPORT_FIELDS
    contact_test_field = 'availability__contact__last_name'
    corp_name_field = 'availability__corporation__name'

    if period_filter:
        if non_attributed:
            # Export non attributed availabilities for a specific period
            query = Availability.objects.filter(period_id=period_filter, training__isnull=True)
            export_fields = NON_ATTR_EXPORT_FIELDS
            contact_test_field = 'contact__last_name'
            corp_name_field = 'corporation__name'
        else:
            # Export trainings for a specific period
            query = Training.objects.filter(availability__period_id=period_filter)
    else:
        if scope and scope == 'all':
            # Export all trainings in the database
            query = Training.objects.all()
        else:
            query = Training.objects.filter(availability__period__end_date__gt=school_year_start())

    # Prepare "default" contacts (when not defined on training)
    default_contacts = dict((c, '') for c in Corporation.objects.all().values_list('name', flat=True))
    always_ccs = dict((c, []) for c in Corporation.objects.all().values_list('name', flat=True))
    for contact in CorpContact.objects.all().select_related('corporation').order_by('corporation'):
        if not default_contacts[contact.corporation.name] or contact.is_main is True:
            default_contacts[contact.corporation.name] = contact
        if contact.always_cc:
            always_ccs[contact.corporation.name].append(contact)

    wb = Workbook()
    ws = wb.get_active_sheet()
    ws.title = 'Stages'
    bold = Style(font=Font(bold=True))
    # Headers
    for col_idx, header in enumerate([f[0] for f in export_fields], start=1):
        cell = ws.cell(row=1, column=col_idx)
        cell.value = header
        cell.style = bold
    # Data
    query_keys = [f[1] for f in export_fields if f[1] is not None]
    for row_idx, tr in enumerate(query.values(*query_keys), start=2):
        for col_idx, field in enumerate(query_keys, start=1):
            ws.cell(row=row_idx, column=col_idx).value = tr[field]
        if tr[contact_test_field] is None:
            # Use default contact
            contact = default_contacts.get(tr[corp_name_field])
            if contact:
                ws.cell(row=row_idx, column=col_idx-3).value = contact.title
                ws.cell(row=row_idx, column=col_idx-2).value = contact.first_name
                ws.cell(row=row_idx, column=col_idx-1).value = contact.last_name
                ws.cell(row=row_idx, column=col_idx).value = contact.email
        if always_ccs[tr[corp_name_field]]:
            ws.cell(row=row_idx, column=col_idx+1).value = "; ".join([c.email for c in always_ccs[tr[corp_name_field]]])

    response = HttpResponse(save_virtual_workbook(wb), content_type='application/ms-excel')
    response['Content-Disposition'] = 'attachment; filename=%s%s.xlsx' % (
          'stages_export_', date.strftime(date.today(), '%Y-%m-%d'))
    return response
<commit_msg>Replaced obsolete SortedDict and fix Referent queryset<commit_after># -*- encoding: utf-8 -*-
from __future__ import unicode_literals

import json
from collections import OrderedDict
from datetime import date, datetime, timedelta

from django.db.models import Count
from django.http import HttpResponse, HttpResponseNotAllowed
from django.shortcuts import get_object_or_404
from django.views.generic import DetailView, TemplateView, ListView

from .forms import PeriodForm
from .models import (Section, Student, Corporation, CorpContact, Period,
    Training, Referent, Availability)

def school_year_start():
    """ Return first official day of current school year """
    current_year = date.today().year
    if date(current_year, 8, 1) > date.today():
        return date(current_year-1, 8, 1)
    else:
        return date(current_year, 8, 1)

class CorporationListView(ListView):
    model = Corporation
    template_name = 'corporations.html'

class CorporationView(DetailView):
    model = Corporation
    template_name = 'corporation.html'
    context_object_name = 'corp'

    def get_context_data(self, **kwargs):
        context = super(CorporationView, self).get_context_data(**kwargs)
        # Create a structure like:
        #   {'2011-2012': {'avails': [avail1, avail2, ...], 'stats': {'fil': num}},
        #    '2012-2013': ...}
        school_years = OrderedDict()
        for av in Availability.objects.filter(corporation=self.object
                ).select_related('training__student__klass', 'period__section'
                ).order_by('period__start_date'):
            if av.period.school_year not in school_years:
                school_years[av.period.school_year] = {'avails': [], 'stats': {}}
            school_years[av.period.school_year]['avails'].append(av)
            if av.period.section.name not in school_years[av.period.school_year]['stats']:
                school_years[av.period.school_year]['stats'][av.period.section.name] = 0
            try:
                av.training
                # Only add to stats if training exists
                school_years[av.period.school_year]['stats'][av.period.section.name] += av.period.weeks
            except Training.DoesNotExist:
                pass

        context['years'] = school_years
        return context

class AttributionView(TemplateView):
    """
    Base view for the attribution screen. Populate sections and referents.
    All other data are retrieved through AJAX requests:
      * training periods: section_period
      * corp. availabilities for current period: period_availabilities
      * already planned training for current period: TrainingsByPeriodView
      * student list targetted by current period: period_students
    When an availability is chosen:
      * corp. contact list: CorpContactJSONView
    When a student is chosen;
      * details of a student: StudentSummaryView
    """
    template_name = 'attribution.html'

    def get_context_data(self, **kwargs):
        context = super(AttributionView, self).get_context_data(**kwargs)
        # Need 2 queries, because referents with no training item would not appear in the second query
        referents = Referent.objects.all().order_by('last_name', 'first_name')

        # Populate each referent with the number of referencies done during the current school year
        ref_counts = dict([(ref.id, ref.num_refs)
                for ref in Referent.objects.filter(training__availability__period__end_date__gte=school_year_start()
                ).annotate(num_refs=Count('training'))])
        for ref in referents:
            ref.num_refs = ref_counts.get(ref.id, 0)

        context.update({
            #'period_form': PeriodForm(),
            'sections': Section.objects.all(),
            'referents': referents,
        })
        return context

class StudentSummaryView(DetailView):
    model = Student
    template_name = 'student_summary.html'

    def get_context_data(self, **kwargs):
        context = super(StudentSummaryView, self).get_context_data(**kwargs)
        context['previous_stages'] = self.object.training_set.all(
            ).select_related('availability__corporation').order_by('availability__period__end_date')
        return context

class AvailabilitySummaryView(DetailView):
    model = Availability
    template_name = 'availability_summary.html'

class TrainingsByPeriodView(ListView):
    template_name = 'trainings_list.html'
    context_object_name = 'trainings'

    def get_queryset(self):
        return Training.objects.select_related('student__klass', 'availability__corporation', 'availability__domain'
            ).filter(availability__period__pk=self.kwargs['pk']
            ).order_by('student__last_name', 'student__first_name')

class CorpContactJSONView(ListView):
    """ Return all contacts from a given corporation """
    return_fields = ['id', 'first_name', 'last_name', 'role', 'is_main']

    def get_queryset(self):
        return CorpContact.objects.filter(corporation__pk=self.kwargs['pk'])

    def render_to_response(self, context):
        serialized = [dict([(field, getattr(obj, field)) for field in self.return_fields])
                      for obj in context['object_list']]
        return HttpResponse(json.dumps(serialized), content_type="application/json")

# AJAX views:

def section_periods(request, pk):
    """ Return all periods (until 2 years ago) from a section (JSON) """
    section = get_object_or_404(Section, pk=pk)
    two_years_ago = datetime.now() - timedelta(days=365 * 2)
    periods = [{'id': p.id, 'dates': p.dates, 'title': p.title}
               for p in section.period_set.filter(start_date__gt=two_years_ago).order_by('-start_date')]
    return HttpResponse(json.dumps(periods), content_type="application/json")

def section_classes(request, pk):
    section = get_object_or_404(Section, pk=pk)
    classes = [(k.id, k.name) for k in section.klass_set.all()]
    return HttpResponse(json.dumps(classes), content_type="application/json")

def period_students(request, pk):
    """
    Return all active students from period's section and level,
    with corresponding Training if existing (JSON)
    """
    period = get_object_or_404(Period, pk=pk)
    students = Student.objects.filter(
        archived=False, klass__section=period.section, klass__level=period.relative_level
        ).order_by('last_name')
    trainings = dict((t.student_id, t.id) for t in Training.objects.filter(availability__period=period))
    data = [{
        'name': unicode(s),
        'id': s.id,
        'training_id': trainings.get(s.id),
        'klass': s.klass.name} for s in students]
    return HttpResponse(json.dumps(data), content_type="application/json")

def period_availabilities(request, pk):
    """ Return all availabilities in the specified period """
    period = get_object_or_404(Period, pk=pk)
    corps = [{'id': av.id, 'id_corp': av.corporation.id, 'corp_name': av.corporation.name,
              'domain': av.domain.name, 'free': av.free}
             for av in period.availability_set.select_related('corporation').all()]
    return HttpResponse(json.dumps(corps), content_type="application/json")

def new_training(request):
    if request.method != 'POST':
        return HttpResponseNotAllowed()
    ref_key = request.POST.get('referent')
    cont_key = request.POST.get('contact')
    try:
        ref = Referent.objects.get(pk=ref_key) if ref_key else None
        contact = CorpContact.objects.get(pk=cont_key) if cont_key else None
        avail = Availability.objects.get(pk=request.POST.get('avail'))
        training = Training.objects.create(
            student=Student.objects.get(pk=request.POST.get('student')),
            availability=avail,
            referent=ref,
        )
        if avail.contact != contact:
            avail.contact = contact
            avail.save()
    except Exception as exc:
        return HttpResponse(str(exc))
    return HttpResponse('OK')

def del_training(request):
    """ Delete training and return the referent id """
    if request.method != 'POST':
        return HttpResponseNotAllowed()
    training = get_object_or_404(Training, pk=request.POST.get('pk'))
    ref_id = training.referent_id
    training.delete()
    return HttpResponse(json.dumps({'ref_id': ref_id}), content_type="application/json")

EXPORT_FIELDS = [
    ('Prénom', 'student__first_name'), ('Nom', 'student__last_name'),
    ('ID externe', 'student__ext_id'),
    ('Classe', 'student__klass__name'), ('Filière', 'student__klass__section__name'),
    ('Nom du stage', 'availability__period__title'),
    ('Début', 'availability__period__start_date'), ('Fin', 'availability__period__end_date'),
    ('Remarques stage', 'comment'),
    ('Prénom référent', 'referent__first_name'), ('Nom référent', 'referent__last_name'),
    ('Courriel référent', 'referent__email'),
    ('Institution', 'availability__corporation__name'),
    ('ID externe Inst', 'availability__corporation__ext_id'),
    ('Rue Inst', 'availability__corporation__street'),
    ('NPA Inst', 'availability__corporation__pcode'),
    ('Ville Inst', 'availability__corporation__city'),
    ('Tél Inst', 'availability__corporation__tel'),
    ('Domaine', 'availability__domain__name'),
    ('Remarques Inst', 'availability__comment'),
    ('Civilité contact', 'availability__contact__title'),
    ('Prénom contact', 'availability__contact__first_name'),
    ('Nom contact', 'availability__contact__last_name'),
    ('ID externe contact', 'availability__contact__ext_id'),
    ('Tél contact', 'availability__contact__tel'),
    ('Courriel contact', 'availability__contact__email'),
    ('Courriel contact - copie', None),
]

NON_ATTR_EXPORT_FIELDS = [
    ('Filière', 'period__section__name'),
    ('Nom du stage', 'period__title'),
    ('Début', 'period__start_date'), ('Fin', 'period__end_date'),
    ('Institution', 'corporation__name'),
    ('Rue Inst', 'corporation__street'),
    ('NPA Inst', 'corporation__pcode'),
    ('Ville Inst', 'corporation__city'),
    ('Tél Inst', 'corporation__tel'),
    ('Domaine', 'domain__name'),
    ('Remarques Inst', 'comment'),
    ('Civilité contact', 'contact__title'),
    ('Prénom contact', 'contact__first_name'),
    ('Nom contact', 'contact__last_name'),
    ('Tél contact', 'contact__tel'),
    ('Courriel contact', 'contact__email'),
    ('Courriel contact - copie', None),
]

def stages_export(request, scope=None):
    from datetime import date
    from openpyxl import Workbook
    from openpyxl.styles import Font, Style
    from openpyxl.writer.excel import save_virtual_workbook

    period_filter = request.GET.get('period')
    non_attributed = bool(int(request.GET.get('non_attr', 0)))

    export_fields = EXPORT_FIELDS
    contact_test_field = 'availability__contact__last_name'
    corp_name_field = 'availability__corporation__name'

    if period_filter:
        if non_attributed:
            # Export non attributed availabilities for a specific period
            query = Availability.objects.filter(period_id=period_filter, training__isnull=True)
            export_fields = NON_ATTR_EXPORT_FIELDS
            contact_test_field = 'contact__last_name'
            corp_name_field = 'corporation__name'
        else:
            # Export trainings for a specific period
            query = Training.objects.filter(availability__period_id=period_filter)
    else:
        if scope and scope == 'all':
            # Export all trainings in the database
            query = Training.objects.all()
        else:
            query = Training.objects.filter(availability__period__end_date__gt=school_year_start())

    # Prepare "default" contacts (when not defined on training)
    default_contacts = dict((c, '') for c in Corporation.objects.all().values_list('name', flat=True))
    always_ccs = dict((c, []) for c in Corporation.objects.all().values_list('name', flat=True))
    for contact in CorpContact.objects.all().select_related('corporation').order_by('corporation'):
        if not default_contacts[contact.corporation.name] or contact.is_main is True:
            default_contacts[contact.corporation.name] = contact
        if contact.always_cc:
            always_ccs[contact.corporation.name].append(contact)

    wb = Workbook()
    ws = wb.get_active_sheet()
    ws.title = 'Stages'
    bold = Style(font=Font(bold=True))
    # Headers
    for col_idx, header in enumerate([f[0] for f in export_fields], start=1):
        cell = ws.cell(row=1, column=col_idx)
        cell.value = header
        cell.style = bold
    # Data
    query_keys = [f[1] for f in export_fields if f[1] is not None]
    for row_idx, tr in enumerate(query.values(*query_keys), start=2):
        for col_idx, field in enumerate(query_keys, start=1):
            ws.cell(row=row_idx, column=col_idx).value = tr[field]
        if tr[contact_test_field] is None:
            # Use default contact
            contact = default_contacts.get(tr[corp_name_field])
            if contact:
                ws.cell(row=row_idx, column=col_idx-3).value = contact.title
                ws.cell(row=row_idx, column=col_idx-2).value = contact.first_name
                ws.cell(row=row_idx, column=col_idx-1).value = contact.last_name
                ws.cell(row=row_idx, column=col_idx).value = contact.email
        if always_ccs[tr[corp_name_field]]:
            ws.cell(row=row_idx, column=col_idx+1).value = "; ".join([c.email for c in always_ccs[tr[corp_name_field]]])

    response = HttpResponse(save_virtual_workbook(wb), content_type='application/ms-excel')
    response['Content-Disposition'] = 'attachment; filename=%s%s.xlsx' % (
          'stages_export_', date.strftime(date.today(), '%Y-%m-%d'))
    return response
<|endoftext|>