<commit_before># Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Helpful things for user interfaces."""

import datetime
import os
import tempfile
import better_webbrowser

import benchmark
import config
import geoip
import data_sources
import nameserver_list
import reporter
import site_connector
import util

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

def GenerateOutputFilename(template):
  # used for resolv.conf
  if '.' in template:
    filename = template
  else:
    output_base = 'namebench_%s' % datetime.datetime.strftime(datetime.datetime.now(),
                                                              '%Y-%m-%d %H%M')
    output_base = output_base.replace(':', '').replace(' ', '_')
    filename = '.'.join((output_base, template))

  output_dir = tempfile.gettempdir()
  return os.path.join(output_dir, filename)

class BaseUI(object):
  """Common methods for all UI implementations."""

  def __init__(self):
    self.SetupDataStructures()

  def SetupDataStructures(self):
    self.reporter = None
    self.nameservers = None
    self.bmark = None
    self.report_path = None
    self.csv_path = None
    self.geodata = None
    self.country = None
    self.sources = {}
    self.url = None
    self.share_state = None
    self.test_records = []

  def UpdateStatus(self, msg, **kwargs):
    """Update the little status message on the bottom of the window."""
    if hasattr(self, 'status_callback') and self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print msg

  def DebugMsg(self, message):
    self.UpdateStatus(message, debug=True)

  def LoadDataSources(self):
    self.data_src = data_sources.DataSources(status_callback=self.UpdateStatus)

  def PrepareTestRecords(self):
    if self.options.input_source:
      src_type = self.options.input_source
      src_name = self.data_src.GetNameForSource(src_type)
    else:
      (src_type, src_name) = self.data_src.GetBestSourceDetails()[:2]
      self.options.input_source = src_type

    self.test_records = self.data_src.GetTestsFromSource(src_type,
                                                         self.options.query_count,
                                                         select_mode=self.options.select_mode)

  def PrepareNameServers(self):
    """Setup self.nameservers to have a list of healthy fast servers."""

    self.nameservers = nameserver_list.NameServers(
        self.supplied_ns,
        global_servers=self.global_ns,
        regional_servers=self.regional_ns,
        include_internal=self.include_internal,
        num_servers=self.options.num_servers,
        timeout=self.options.timeout,
        ping_timeout=self.options.ping_timeout,
        health_timeout=self.options.health_timeout,
        ipv6_only=self.options.ipv6_only,
        status_callback=self.UpdateStatus
    )
    if self.options.invalidate_cache:
      self.nameservers.InvalidateSecondaryCache()

    self.nameservers.cache_dir = tempfile.gettempdir()

    # Don't waste time checking the health of the only nameserver in the list.
    if len(self.nameservers) > 1:
      self.nameservers.thread_count = int(self.options.health_thread_count)
      self.nameservers.cache_dir = tempfile.gettempdir()

    self.UpdateStatus('Checking latest sanity reference')
    (primary_checks, secondary_checks, censor_tests) = config.GetLatestSanityChecks()
    if not self.options.enable_censorship_checks:
      censor_tests = []
    else:
      self.UpdateStatus('Censorship checks enabled: %s found.' % len(censor_tests))

    self.UpdateStatus('Checking nameserver health')
    self.nameservers.CheckHealth(primary_checks, secondary_checks, censor_tests=censor_tests)

  def PrepareBenchmark(self):
    """Setup the benchmark object with the appropriate dataset."""
    if len(self.nameservers) == 1:
      thread_count = 1
    else:
      thread_count = self.options.benchmark_thread_count

    self.bmark = benchmark.Benchmark(self.nameservers,
                                     query_count=self.options.query_count,
                                     run_count=self.options.run_count,
                                     thread_count=thread_count,
                                     status_callback=self.UpdateStatus)

  def RunBenchmark(self):
    """Run the benchmark."""
    results = self.bmark.Run(self.test_records)
    index = []

    if self.options.upload_results:
      connector = site_connector.SiteConnector(self.options)
      index_hosts = connector.GetIndexHosts()
      if index_hosts:
        index = self.bmark.RunIndex(index_hosts)
      else:
        index = []
      self.DiscoverLocation()
      self.nameservers.RunPortBehaviorThreads()

    self.reporter = reporter.ReportGenerator(self.options, self.nameservers,
                                             results, index=index, geodata=self.geodata)
  def DiscoverLocation(self):
    if not getattr(self, 'geodata', None):
      self.geodata = geoip.GetGeoData()
      self.country = self.geodata.get('country_name', None)
    return self.geodata

  def RunAndOpenReports(self):
    """Run the benchmark and open up the report on completion."""
    self.RunBenchmark()
    best = self.reporter.BestOverallNameServer()
    self.CreateReports()
    if self.options.template == 'html':
      self.DisplayHtmlReport()
    if self.url:
      self.UpdateStatus('Complete! Your results: %s' % self.url)
    else:
      self.UpdateStatus('Complete! %s [%s] is the best.' % (best.name, best.ip))

  def CreateReports(self):
    """Create CSV & HTML reports for the latest run."""

    if self.options.output_file:
      self.report_path = self.options.output_file
    else:
      self.report_path = GenerateOutputFilename(self.options.template)

    if self.options.csv_file:
      self.csv_path = self.options_csv_file
    else:
      self.csv_path = GenerateOutputFilename('csv')

    if self.options.upload_results:
      self.UpdateStatus('Uploading results to %s' % self.options.site_url)
      json_data = self.reporter.CreateJsonData()
      connector = site_connector.SiteConnector(self.options)
      self.url, self.share_state = connector.UploadJsonResults(json_data, hide_results=self.options.hide_results)
      if self.url:
        self.UpdateStatus("Your sharing URL: %s (%s)" % (self.url, self.share_state))

    self.UpdateStatus('Saving report to %s' % self.report_path)
    f = open(self.report_path, 'w')
    self.reporter.CreateReport(format=self.options.template, output_fp=f,
                               csv_path=self.csv_path, sharing_url=self.url, sharing_state=self.share_state)
    f.close()

    self.UpdateStatus('Saving detailed results to %s' % self.csv_path)
    self.reporter.SaveResultsToCsv(self.csv_path)

  def DisplayHtmlReport(self):
    self.UpdateStatus('Opening %s' % self.report_path)
    better_webbrowser.output = self.DebugMsg
    better_webbrowser.open(self.report_path)

<commit_msg>Don't bother with port behavior checking for a single host<commit_after># Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Helpful things for user interfaces."""

import datetime
import os
import tempfile
import better_webbrowser

import benchmark
import config
import geoip
import data_sources
import nameserver_list
import reporter
import site_connector
import util

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

def GenerateOutputFilename(template):
  # used for resolv.conf
  if '.' in template:
    filename = template
  else:
    output_base = 'namebench_%s' % datetime.datetime.strftime(datetime.datetime.now(),
                                                              '%Y-%m-%d %H%M')
    output_base = output_base.replace(':', '').replace(' ', '_')
    filename = '.'.join((output_base, template))

  output_dir = tempfile.gettempdir()
  return os.path.join(output_dir, filename)

class BaseUI(object):
  """Common methods for all UI implementations."""

  def __init__(self):
    self.SetupDataStructures()

  def SetupDataStructures(self):
    self.reporter = None
    self.nameservers = None
    self.bmark = None
    self.report_path = None
    self.csv_path = None
    self.geodata = None
    self.country = None
    self.sources = {}
    self.url = None
    self.share_state = None
    self.test_records = []

  def UpdateStatus(self, msg, **kwargs):
    """Update the little status message on the bottom of the window."""
    if hasattr(self, 'status_callback') and self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print msg

  def DebugMsg(self, message):
    self.UpdateStatus(message, debug=True)

  def LoadDataSources(self):
    self.data_src = data_sources.DataSources(status_callback=self.UpdateStatus)

  def PrepareTestRecords(self):
    if self.options.input_source:
      src_type = self.options.input_source
      src_name = self.data_src.GetNameForSource(src_type)
    else:
      (src_type, src_name) = self.data_src.GetBestSourceDetails()[:2]
      self.options.input_source = src_type

    self.test_records = self.data_src.GetTestsFromSource(src_type,
                                                         self.options.query_count,
                                                         select_mode=self.options.select_mode)

  def PrepareNameServers(self):
    """Setup self.nameservers to have a list of healthy fast servers."""

    self.nameservers = nameserver_list.NameServers(
        self.supplied_ns,
        global_servers=self.global_ns,
        regional_servers=self.regional_ns,
        include_internal=self.include_internal,
        num_servers=self.options.num_servers,
        timeout=self.options.timeout,
        ping_timeout=self.options.ping_timeout,
        health_timeout=self.options.health_timeout,
        ipv6_only=self.options.ipv6_only,
        status_callback=self.UpdateStatus
    )
    if self.options.invalidate_cache:
      self.nameservers.InvalidateSecondaryCache()

    self.nameservers.cache_dir = tempfile.gettempdir()

    # Don't waste time checking the health of the only nameserver in the list.
    if len(self.nameservers) > 1:
      self.nameservers.thread_count = int(self.options.health_thread_count)
      self.nameservers.cache_dir = tempfile.gettempdir()

    self.UpdateStatus('Checking latest sanity reference')
    (primary_checks, secondary_checks, censor_tests) = config.GetLatestSanityChecks()
    if not self.options.enable_censorship_checks:
      censor_tests = []
    else:
      self.UpdateStatus('Censorship checks enabled: %s found.' % len(censor_tests))

    self.UpdateStatus('Checking nameserver health')
    self.nameservers.CheckHealth(primary_checks, secondary_checks, censor_tests=censor_tests)

  def PrepareBenchmark(self):
    """Setup the benchmark object with the appropriate dataset."""
    if len(self.nameservers) == 1:
      thread_count = 1
    else:
      thread_count = self.options.benchmark_thread_count

    self.bmark = benchmark.Benchmark(self.nameservers,
                                     query_count=self.options.query_count,
                                     run_count=self.options.run_count,
                                     thread_count=thread_count,
                                     status_callback=self.UpdateStatus)

  def RunBenchmark(self):
    """Run the benchmark."""
    results = self.bmark.Run(self.test_records)
    index = []

    if self.options.upload_results:
      connector = site_connector.SiteConnector(self.options)
      index_hosts = connector.GetIndexHosts()
      if index_hosts:
        index = self.bmark.RunIndex(index_hosts)
      else:
        index = []
      self.DiscoverLocation()
      if len(self.nameservers) > 1:
        self.nameservers.RunPortBehaviorThreads()

    self.reporter = reporter.ReportGenerator(self.options, self.nameservers,
                                             results, index=index, geodata=self.geodata)
  def DiscoverLocation(self):
    if not getattr(self, 'geodata', None):
      self.geodata = geoip.GetGeoData()
      self.country = self.geodata.get('country_name', None)
    return self.geodata

  def RunAndOpenReports(self):
    """Run the benchmark and open up the report on completion."""
    self.RunBenchmark()
    best = self.reporter.BestOverallNameServer()
    self.CreateReports()
    if self.options.template == 'html':
      self.DisplayHtmlReport()
    if self.url:
      self.UpdateStatus('Complete! Your results: %s' % self.url)
    else:
      self.UpdateStatus('Complete! %s [%s] is the best.' % (best.name, best.ip))

  def CreateReports(self):
    """Create CSV & HTML reports for the latest run."""

    if self.options.output_file:
      self.report_path = self.options.output_file
    else:
      self.report_path = GenerateOutputFilename(self.options.template)

    if self.options.csv_file:
      self.csv_path = self.options_csv_file
    else:
      self.csv_path = GenerateOutputFilename('csv')

    if self.options.upload_results:
      self.UpdateStatus('Uploading results to %s' % self.options.site_url)
      json_data = self.reporter.CreateJsonData()
      connector = site_connector.SiteConnector(self.options)
      self.url, self.share_state = connector.UploadJsonResults(json_data, hide_results=self.options.hide_results)
      if self.url:
        self.UpdateStatus("Your sharing URL: %s (%s)" % (self.url, self.share_state))

    self.UpdateStatus('Saving report to %s' % self.report_path)
    f = open(self.report_path, 'w')
    self.reporter.CreateReport(format=self.options.template, output_fp=f,
                               csv_path=self.csv_path, sharing_url=self.url, sharing_state=self.share_state)
    f.close()

    self.UpdateStatus('Saving detailed results to %s' % self.csv_path)
    self.reporter.SaveResultsToCsv(self.csv_path)

  def DisplayHtmlReport(self):
    self.UpdateStatus('Opening %s' % self.report_path)
    better_webbrowser.output = self.DebugMsg
    better_webbrowser.open(self.report_path)

<|endoftext|>


<commit_before>#!/usr/bin/env python

# Copyright (c) 2016-2017 Terry D. Ott
#
#   You may use, distribute, and modify this code under the terms
#   of the MIT License.
#
#   You should have received a copy of the MIT License with this
#   file. If not, please visit:
#
#   https://github.com/TerryOtt/ntpserverstats/blob/master/LICENSE
#
#   or
#
#   https://opensource.org/licenses/MIT

import logging

logger = logging.getLogger()
logger.setLevel(level=logging.DEBUG)

def s3_RRDWriteHandler(event, context):
  logger.info("Got event: {0}".format(format(event)))
  return "Hello, world"
<commit_msg>Need to merge in changes from other branch<commit_after>#!/usr/bin/env python

# Copyright (c) 2016-2017 Terry D. Ott
#
#   You may use, distribute, and modify this code under the terms
#   of the MIT License.
#
#   You should have received a copy of the MIT License with this
#   file. If not, please visit:
#
#   https://github.com/TerryOtt/ntpserverstats/blob/master/LICENSE
#
#   or
#
#   https://opensource.org/licenses/MIT

import boto3        # AWS SDK for Python
import logging      # AWS kindly puts anything logged into CloudWatch
import tempfile     # Used for reading RRD, writing of PNG

# Set up global logging
logger = logging.getLogger()
logger.setLevel(level=logging.DEBUG)

def s3_RRDWriteHandler(event, context):
  for currWriteEvent in event['Records']:
    _processS3Write(currWriteEvent['s3'])

def _processS3Write( s3WriteRecord ):
  s3_sourceBucket = s3WriteRecord['bucket']['name']
  s3_sourceRRD    = s3WriteRecord['object']['key']

  logging.info( "S3 write; bucket: {0}, rrd: {1}".format(
    s3_sourceBucket, s3_sourceRRD) )

  # Pull the RRD down

<|endoftext|>


<commit_before>import hashlib
import urllib, urllib2

def flatten_nested_hash(hash_table):
    """
    Flatten nested dictionary for GET / POST / DELETE API request
    """
    def flatten(hash_table, brackets=True):
        f = {}
        for key,value in hash_table.items():
            _key = '[' + str(key) + ']' if (brackets == True) else str(key)
            if type(value) is dict:
                for k,v in flatten(value).items():
                    f[_key + k] = v
            elif type(value) is list:
                temp_hash = {}
                for i,v in enumerate(value):
                    temp_hash[str(i)] = v
                for k,v in flatten(temp_hash).items():
                    f[ _key + k] = v
            else:
                f[_key] = value
        return f
    return flatten(hash_table, False)

def extract_params(params):
    """
    Extracts the values of a set of parameters, recursing into nested dictionaries.
    """
    values = []
    for key, value in params.items():
        if type(value) == type(dict()):
            values.extend(extract_params(value))
        elif type(value) == type(list()):
            values.extend(value)
        else:
            values.append(value)
    return values

def get_signature_string(params, secret):
    """
    Returns the unhashed signature string (secret + sorted list of param values) for an API call.
    @param params: dictionary values to generate signature string
    @param sercret: secret string
    """
    str_list = []
    for item in extract_params(params):
        str_list.append(str(item))
    print str_list
    str_list.sort()
    print str_list
    print "".join(str_list)
    return secret + "".join(str_list)

def get_signature_hash(params, secret):
    """
    Returns an MD5 hash of the signature string for an API call.
    @param params: dictionary values to generate signature hash
    @param sercret: secret string
    """
    return hashlib.md5(get_signature_string(params, secret)).hexdigest()

class SailthruClient:

    """
    This class makes HTTP Request to Sailthru API server
    Response from server depends on the format being queried
    If 'json' format is requested, client will recieve JSON object
    XML format is also available but XML response have not been tested thoroughly
    https://github.com/sailthru/sailthru-python-client

    Usage:
        from sailthru import SailthruClient
        api_key = "your-api-key"
        api_secret = "api-secret"
        client = SailthruCLient(api_key, api_secret)
    """

    def __init__(self, api_key, secret, api_url=None):
        self.api_key = api_key
        self.secret = secret
        self.api_url = api_url if (api_url is not None) else 'https://api.sailthru.com'
        self.user_agent = 'Sailthru API Python Client'

    def send(self, template, email, _vars = {}, options = {}, schedule_time = None):
        """
        Remotely send an email template to a single email address.
        http://docs.sailthru.com/api/send
        @param template: template string
        @param email: Email value
        @param _vars: a key/value hash of the replacement vars to use in the send. Each var may be referenced as {varname} within the template itself
        @param options: optional dictionary to include replyto and/or test keys
        @param schedule_time: do not send the email immediately, but at some point in the future. Any date recognized by PHP's strtotime function is valid, but be sure to specify timezone or use a UTC time to avoid confusion
        """
        data = {}
        data['template'] = template
        data['email'] = email
        data['vars'] = _vars
        data['options'] = options
        if schedule_time is not None:
            data['schedule_time'] = schedule_time
        return self._api_post('send', data)

    def multi_send(self, template, emails, _vars = {}, evars = {}, options = {}):
        """
        Remotely send an email template to multiple email addresses.
        http://docs.sailthru.com/api/send
        @param template: template string
        @param emails: List with email values or comma separated email string
        @param _vars: a key/value hash of the replacement vars to use in the send. Each var may be referenced as {varname} within the template itself
        @param options: optional dictionary to include replyto and/or test keys
        @param schedule_time: do not send the email immediately, but at some point in the future. Any date recognized by PHP's strtotime function is valid, but be sure to specify timezone or use a UTC time to avoid confusion
        """
        data = {}
        data['template'] = template
        data['email'] = ','.join(emails) if type(emails) is list else emails
        data['vars'] = _vars
        data['evars'] = evars
        data['options'] = options
        return self._api_post('send', data)

    def get_send(self, send_id):
        """
        Get the status of a send
        """
        return self._api_get('send', {'send_id': send_id})

    def cancel_send(self, send_id):
        """
        Cancels an email that you had previously scheduled for future sending with the schedule_time parameter. It is not possible to cancel an email that has not been scheduled.
        """
        return self._api_delete('send', {'send_id': send_id})

    def get_email(self, email):
        """
        get user email data
        http://docs.sailthru.com/api/email
        """
        data = {'email': email}
        return self._api_request('email', data, 'GET')

    def set_email(self, email, _vars={}, lists=[], templates=[], verified=0, optout=None, send=None, send_vars=[]):
        """
        Update information about one of your users, including adding and removing the user from lists.
        http://docs.sailthru.com/api/email
        """
        data = {}
        data['email'] = email
        data['vars'] = _vars
        data['lists'] = lists
        data['templates'] = templates
        data['verified'] = int(verified)
        if optout is not None:
            data['optout'] = optout
        if send is not None:
            data['send'] = send
        data['send_vars'] = send_vars
        return self._api_post('email', data)

    def schedule_blast(self, name, list, schedule_time, from_name, from_email, subject, content_html, content_text, options={}):
        """
        Schedule a mass mail blast
        http://docs.sailthru.com/api/blast
        @param name: name to give to this new blast
        @param list: mailing list name to send to
        @param schedule_time:  when the blast should send. Dates in the past will be scheduled for immediate delivery. Any English textual datetime format known to PHP's strtotime function is acceptable, such as 2009-03-18 23:57:22 UTC, now (immediate delivery), +3 hours (3 hours from now), or February 14, 9:30 EST. Be sure to specify a timezone if you use an exact time.
        @param from_name: name appearing in the "From" of the email
        @param from_email: email address to use as the "from" – choose from any of your verified emails
        @param subject: subject line of the email
        @param content_html: HTML format version of the email
        @param content_text: Text format version of the email
        @param options: optional parameters dictionary
            blast_id
            copy_blast
            copy_template
            replyto
            report_email
            is_link_tracking
            is_google_analytics
            is_public
            suppress_list
            test_vars
            email_hour_range
            abtest
            test_percent
            data_feed_url
        """
        data = options
        data['name'] = name
        data['list'] = list
        data['schedule_time'] = schedule_time
        data['from_name'] = from_name
        data['from_email'] = from_email
        data['subject'] = subject
        data['content_html'] = content_html
        data['content_text'] = content_text
        return self._api_post('blast', data)

    def update_blast(self, blast_id, name=None, list=None, schedule_time=None, from_name=None, from_email=None, subject=None, content_html=None, content_text=None, options={}):
        """
        updates existing blast
        http://docs.sailthru.com/api/blast
        @param blast_id: blast id
        @param name: name of the blast
        @param list: blast list
        @param schedule_time: new schedule time
        @param from_name: name appearing in the "From" of the email
        @param from_email: email address to use as the "from" – choose from any of your verified emails
        @param subject: subject line of the email
        @param content_html: HTML format version of the email
        @param content_text: Text format version of the email
        @param options: optional parameters dictionary
            blast_id
            copy_blast
            copy_template
            replyto
            report_email
            is_link_tracking
            is_google_analytics
            is_public
            suppress_list
            test_vars
            email_hour_range
            abtest
            test_percent
            data_feed_url
        """
        data = options
        data['blast_id'] = blast_id
        if name is not None:
            data['name'] = name
        if list is not None:
            data['list'] = list
        if schedule_time is not None:
            data['schedule_time'] = schedule_time
        if from_name is not None:
            data['from_name'] = from_name
        if from_email is not None:
            data['from_email'] = from_email
        if subject is not None:
            data['subject'] = subject
        if content_html is not None:
            data['content_html'] = content_html
        if content_text is not None:
            data['content_text'] = content_text
        return self._api_post('blast', data)

    def get_blast(self, blast_id):
        """
        Get Blast information
        http://docs.sailthru.com/api/blast
        """
        return self._api_get('blast', {'blast_id': blast_id})

    def delete_blast(self, blast_id):
        """
        delete existing blast
        """
        return self._api_delete('blast', {'blast_id': blast_id})

    def cancel_blast(self, blast_id):
        """
        Cancel a scheduled Blast
        """
        data = {}
        data['blast_id'] = blast_id
        data['schedule_time'] = ''
        return self._api_post('blast', data)

    def get_template(self, template_name):
        return self._api_get('template', {'template': template})

    def save_template(self, template, template_fields = {}):
        data = template_fields
        data['template'] = template
        return self._api_post('template', data)

    def get_list(self, list, format='txt'):
        """
        Download a list. Obviously, this can potentially be a very large download.
        'txt' is default format since, its more compact as compare to others
        http://docs.sailthru.com/api/list
        """
        data = {}
        data['list'] = list
        data['format'] = format
        return self._api_get('list', data)

    def get_lists(self):
        """
        Get metadata for all lists
        """
        data = {'list': ''} #blank list
        return self._api_get('list', data)

    def save_list(self, list, emails):
        """
        Upload a list. The list import job is queued and will happen shortly after the API request.
        http://docs.sailthru.com/api/list
        @param list: list name
        @param emails: List of email values or comma separated string
        """
        data = {}
        data['list'] = list
        data['emails'] = ','.join(emails) if emails is list else emails
        return self._api_post('list', data)

    def delete_list(self, list):
        """
        delete given list
        http://docs.sailthru.com/api/list
        """
        return self._api_delete('list', {'list': list})

    def _api_get(self, action, data):
        """
        Perform an HTTP GET request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'GET')

    def _api_post(self, action, data):
        """
        Perform an HTTP POST request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'POST')

    def _api_delete(self, action, data):
        """
        Perform an HTTP DELETE request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'DELETE')

    def _api_request(self, action, data, request_type):
        """
        Make Request to Sailthru API with given data and api key, format and signature hash
        """
        data['api_key'] = self.api_key
        data['format'] = data.get('format', 'json')
        data['sig'] = get_signature_hash(data, self.secret)
        return self._http_request(self.api_url+'/'+action, data, request_type)

    def _http_request(self, url, data, method='POST'):
        """
        Perform an HTTP GET / POST / DELETE request
        """
        data = flatten_nested_hash(data)
        data = urllib.urlencode(data, True)
        opener = urllib2.build_opener(urllib2.HTTPHandler)
        try:
            if method == 'POST':
                req = urllib2.Request(url, data)
            else:
                url += '?' + data
                req = urllib2.Request(url)
            req.add_header('User-Agent', self.user_agent)
            req.get_method = lambda: method
            response = opener.open(req)
            response_data = response.read()
            response.close()
            return response_data
        except urllib2.URLError, e:
            return str(e)
        except urllib2.HTTPError, e:
            return e.read()<commit_msg>get_templates<commit_after># -*- coding: utf-8 -*-

import hashlib
import urllib, urllib2

def flatten_nested_hash(hash_table):
    """
    Flatten nested dictionary for GET / POST / DELETE API request
    """
    def flatten(hash_table, brackets=True):
        f = {}
        for key,value in hash_table.items():
            _key = '[' + str(key) + ']' if (brackets == True) else str(key)
            if type(value) is dict:
                for k,v in flatten(value).items():
                    f[_key + k] = v
            elif type(value) is list:
                temp_hash = {}
                for i,v in enumerate(value):
                    temp_hash[str(i)] = v
                for k,v in flatten(temp_hash).items():
                    f[ _key + k] = v
            else:
                f[_key] = value
        return f
    return flatten(hash_table, False)

def extract_params(params):
    """
    Extracts the values of a set of parameters, recursing into nested dictionaries.
    """
    values = []
    for key, value in params.items():
        if type(value) == type(dict()):
            values.extend(extract_params(value))
        elif type(value) == type(list()):
            values.extend(value)
        else:
            values.append(value)
    return values

def get_signature_string(params, secret):
    """
    Returns the unhashed signature string (secret + sorted list of param values) for an API call.
    @param params: dictionary values to generate signature string
    @param sercret: secret string
    """
    str_list = []
    for item in extract_params(params):
        str_list.append(str(item))
    str_list.sort()
    return secret + "".join(str_list)

def get_signature_hash(params, secret):
    """
    Returns an MD5 hash of the signature string for an API call.
    @param params: dictionary values to generate signature hash
    @param sercret: secret string
    """
    return hashlib.md5(get_signature_string(params, secret)).hexdigest()

class SailthruClient:

    """
    This class makes HTTP Request to Sailthru API server
    Response from server depends on the format being queried
    If 'json' format is requested, client will recieve JSON object
    XML format is also available but XML response have not been tested thoroughly
    https://github.com/sailthru/sailthru-python-client

    Usage:
        from sailthru import SailthruClient
        api_key = "your-api-key"
        api_secret = "api-secret"
        client = SailthruCLient(api_key, api_secret)
    """

    def __init__(self, api_key, secret, api_url=None):
        self.api_key = api_key
        self.secret = secret
        self.api_url = api_url if (api_url is not None) else 'https://api.sailthru.com'
        self.user_agent = 'Sailthru API Python Client'

    def send(self, template, email, _vars = {}, options = {}, schedule_time = None):
        """
        Remotely send an email template to a single email address.
        http://docs.sailthru.com/api/send
        @param template: template string
        @param email: Email value
        @param _vars: a key/value hash of the replacement vars to use in the send. Each var may be referenced as {varname} within the template itself
        @param options: optional dictionary to include replyto and/or test keys
        @param schedule_time: do not send the email immediately, but at some point in the future. Any date recognized by PHP's strtotime function is valid, but be sure to specify timezone or use a UTC time to avoid confusion
        """
        data = {}
        data['template'] = template
        data['email'] = email
        data['vars'] = _vars
        data['options'] = options
        if schedule_time is not None:
            data['schedule_time'] = schedule_time
        return self._api_post('send', data)

    def multi_send(self, template, emails, _vars = {}, evars = {}, options = {}):
        """
        Remotely send an email template to multiple email addresses.
        http://docs.sailthru.com/api/send
        @param template: template string
        @param emails: List with email values or comma separated email string
        @param _vars: a key/value hash of the replacement vars to use in the send. Each var may be referenced as {varname} within the template itself
        @param options: optional dictionary to include replyto and/or test keys
        @param schedule_time: do not send the email immediately, but at some point in the future. Any date recognized by PHP's strtotime function is valid, but be sure to specify timezone or use a UTC time to avoid confusion
        """
        data = {}
        data['template'] = template
        data['email'] = ','.join(emails) if type(emails) is list else emails
        data['vars'] = _vars
        data['evars'] = evars
        data['options'] = options
        return self._api_post('send', data)

    def get_send(self, send_id):
        """
        Get the status of a send
        """
        return self._api_get('send', {'send_id': send_id})

    def cancel_send(self, send_id):
        """
        Cancels an email that you had previously scheduled for future sending with the schedule_time parameter. It is not possible to cancel an email that has not been scheduled.
        """
        return self._api_delete('send', {'send_id': send_id})

    def get_email(self, email):
        """
        get user email data
        http://docs.sailthru.com/api/email
        """
        data = {'email': email}
        return self._api_request('email', data, 'GET')

    def set_email(self, email, _vars={}, lists=[], templates=[], verified=0, optout=None, send=None, send_vars=[]):
        """
        Update information about one of your users, including adding and removing the user from lists.
        http://docs.sailthru.com/api/email
        """
        data = {}
        data['email'] = email
        data['vars'] = _vars
        data['lists'] = lists
        data['templates'] = templates
        data['verified'] = int(verified)
        if optout is not None:
            data['optout'] = optout
        if send is not None:
            data['send'] = send
        data['send_vars'] = send_vars
        return self._api_post('email', data)

    def schedule_blast(self, name, list, schedule_time, from_name, from_email, subject, content_html, content_text, options={}):
        """
        Schedule a mass mail blast
        http://docs.sailthru.com/api/blast
        @param name: name to give to this new blast
        @param list: mailing list name to send to
        @param schedule_time:  when the blast should send. Dates in the past will be scheduled for immediate delivery. Any English textual datetime format known to PHP's strtotime function is acceptable, such as 2009-03-18 23:57:22 UTC, now (immediate delivery), +3 hours (3 hours from now), or February 14, 9:30 EST. Be sure to specify a timezone if you use an exact time.
        @param from_name: name appearing in the "From" of the email
        @param from_email: email address to use as the "from" – choose from any of your verified emails
        @param subject: subject line of the email
        @param content_html: HTML format version of the email
        @param content_text: Text format version of the email
        @param options: optional parameters dictionary
            blast_id
            copy_blast
            copy_template
            replyto
            report_email
            is_link_tracking
            is_google_analytics
            is_public
            suppress_list
            test_vars
            email_hour_range
            abtest
            test_percent
            data_feed_url
        """
        data = options
        data['name'] = name
        data['list'] = list
        data['schedule_time'] = schedule_time
        data['from_name'] = from_name
        data['from_email'] = from_email
        data['subject'] = subject
        data['content_html'] = content_html
        data['content_text'] = content_text
        return self._api_post('blast', data)

    def update_blast(self, blast_id, name=None, list=None, schedule_time=None, from_name=None, from_email=None, subject=None, content_html=None, content_text=None, options={}):
        """
        updates existing blast
        http://docs.sailthru.com/api/blast
        @param blast_id: blast id
        @param name: name of the blast
        @param list: blast list
        @param schedule_time: new schedule time
        @param from_name: name appearing in the "From" of the email
        @param from_email: email address to use as the "from" – choose from any of your verified emails
        @param subject: subject line of the email
        @param content_html: HTML format version of the email
        @param content_text: Text format version of the email
        @param options: optional parameters dictionary
            blast_id
            copy_blast
            copy_template
            replyto
            report_email
            is_link_tracking
            is_google_analytics
            is_public
            suppress_list
            test_vars
            email_hour_range
            abtest
            test_percent
            data_feed_url
        """
        data = options
        data['blast_id'] = blast_id
        if name is not None:
            data['name'] = name
        if list is not None:
            data['list'] = list
        if schedule_time is not None:
            data['schedule_time'] = schedule_time
        if from_name is not None:
            data['from_name'] = from_name
        if from_email is not None:
            data['from_email'] = from_email
        if subject is not None:
            data['subject'] = subject
        if content_html is not None:
            data['content_html'] = content_html
        if content_text is not None:
            data['content_text'] = content_text
        return self._api_post('blast', data)

    def get_blast(self, blast_id):
        """
        Get Blast information
        http://docs.sailthru.com/api/blast
        """
        return self._api_get('blast', {'blast_id': blast_id})

    def delete_blast(self, blast_id):
        """
        delete existing blast
        """
        return self._api_delete('blast', {'blast_id': blast_id})

    def cancel_blast(self, blast_id):
        """
        Cancel a scheduled Blast
        """
        data = {}
        data['blast_id'] = blast_id
        data['schedule_time'] = ''
        return self._api_post('blast', data)

    def get_template(self, template_name):
        """
        get information of a given template
        """
        return self._api_get('template', {'template': template_name})

    def get_templates(self):
        """
        get metadata for all user templates
        """
        data = {'template': ''}
        return self._api_get('template', data)

    def save_template(self, template, template_fields = {}):
        data = template_fields
        data['template'] = template
        return self._api_post('template', data)

    def get_list(self, list, format='txt'):
        """
        Download a list. Obviously, this can potentially be a very large download.
        'txt' is default format since, its more compact as compare to others
        http://docs.sailthru.com/api/list
        """
        data = {}
        data['list'] = list
        data['format'] = format
        return self._api_get('list', data)

    def get_lists(self):
        """
        Get metadata for all lists
        """
        data = {'list': ''} #blank list
        return self._api_get('list', data)

    def save_list(self, list, emails):
        """
        Upload a list. The list import job is queued and will happen shortly after the API request.
        http://docs.sailthru.com/api/list
        @param list: list name
        @param emails: List of email values or comma separated string
        """
        data = {}
        data['list'] = list
        data['emails'] = ','.join(emails) if emails is list else emails
        return self._api_post('list', data)

    def delete_list(self, list):
        """
        delete given list
        http://docs.sailthru.com/api/list
        """
        return self._api_delete('list', {'list': list})

    def _api_get(self, action, data):
        """
        Perform an HTTP GET request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'GET')

    def _api_post(self, action, data):
        """
        Perform an HTTP POST request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'POST')

    def _api_delete(self, action, data):
        """
        Perform an HTTP DELETE request, using the shared-secret auth hash.
        @param action: API action call
        @param data: dictionary values
        """
        return self._api_request(action, data, 'DELETE')

    def _api_request(self, action, data, request_type):
        """
        Make Request to Sailthru API with given data and api key, format and signature hash
        """
        data['api_key'] = self.api_key
        data['format'] = data.get('format', 'json')
        data['sig'] = get_signature_hash(data, self.secret)
        return self._http_request(self.api_url+'/'+action, data, request_type)

    def _http_request(self, url, data, method='POST'):
        """
        Perform an HTTP GET / POST / DELETE request
        """
        data = flatten_nested_hash(data)
        data = urllib.urlencode(data, True)
        opener = urllib2.build_opener(urllib2.HTTPHandler)
        try:
            if method == 'POST':
                req = urllib2.Request(url, data)
            else:
                url += '?' + data
                req = urllib2.Request(url)
            req.add_header('User-Agent', self.user_agent)
            req.get_method = lambda: method
            response = opener.open(req)
            response_data = response.read()
            response.close()
            return response_data
        except urllib2.URLError, e:
            return str(e)
        except urllib2.HTTPError, e:
            return e.read()<|endoftext|>


<commit_before>from __future__ import absolute_import

# need to get system mendeley library
from mendeley.exception import MendeleyException
import mendeley as mendeley_lib
import os
import datetime
from sqlalchemy.dialects.postgresql import JSONB
from collections import defaultdict
import requests

from app import db
from app import logger
from doi import Doi
from event import UnpaywallEvent
from event import CedEvent
from util import remove_punctuation

discipline_lookup = {
    u'Arts and Humanities': u'arts and humanities',
    u'Design': u'arts and humanities',
    u'Philosophy': u'arts and humanities',
    u'Agricultural and Biological Sciences': u'life science',
    u'Biochemistry, Genetics and Molecular Biology': u'life science',
    u'Immunology and Microbiology': u'life science',
    u'Veterinary Science and Veterinary Medicine': u'life science',
    u'Business, Management and Accounting': u'business',
    u'Decision Sciences': u'business',
    u'Economics, Econometrics and Finance': u'business',
    u'Chemistry': u'chemistry',
    u'Computer Science': u'computing',
    u'Chemical Engineering': u'engineering',
    u'Engineering': u'engineering',
    u'Materials Science': u'engineering',
    u'Earth and Planetary Sciences': u'environment',
    u'Energy': u'environment',
    u'Environmental Science': u'environment',
    u'Medicine and Dentistry': u'health',
    u'Nursing and Health Professions': u'health',
    u'Pharmacology, Toxicology and Pharmaceutical Science': u'health',
    u'Sports and Recreations': u'health',
    u'Mathematics': u'mathematics',
    u'Physics and Astronomy': u'physics and astronomy',
    u'Neuroscience': u'brain and mind',
    u'Psychology': u'brain and mind',
    u'Linguistics': u'social science',
    u'Social Sciences': u'social science',
    u'Unspecified': u'unspecified'
}

class WeeklyStats(db.Model):
    id = db.Column(db.Text, primary_key=True)
    updated = db.Column(db.DateTime)
    mendeley_api_raw = db.Column(JSONB)
    oadoi_api_raw = db.Column(JSONB)
    sources = db.Column(JSONB)
    main_discipline = db.Column(db.Text)
    num_main_discipline = db.Column(db.Numeric)
    num_unpaywall_events = db.Column(db.Numeric)
    num_twitter_events = db.Column(db.Numeric)
    num_ced_events = db.Column(db.Numeric)
    num_academic_unpaywall_events = db.Column(db.Numeric)
    num_nonacademic_unpaywall_events = db.Column(db.Numeric)
    radio_academic_unpaywall_events = db.Column(db.Numeric)
    is_open_access = db.Column(db.Boolean)
    week = db.Column(db.Numeric)

    def __init__(self, **kwargs):
        self.updated = datetime.datetime.utcnow()
        super(WeeklyStats, self).__init__(**kwargs)

    def set_sources(self):
        event_count_dict = defaultdict(int)

        ced_events = CedEvent.query.filter(CedEvent.doi==self.id).all()
        ced_events_this_week = [e for e in ced_events if e.week==self.week]
        for e in ced_events_this_week:
            event_count_dict[e.source_id] += 1

        unpaywall_events = UnpaywallEvent.query.filter(UnpaywallEvent.doi==self.id).all()
        unpaywall_events_this_week = [e for e in unpaywall_events if e.week==self.week]
        event_count_dict["unpaywall_views"] = len(unpaywall_events_this_week)
        event_count_dict["unpaywall_views_academic"] = len([e for e in unpaywall_events_this_week if e.is_academic_location])
        event_count_dict["unpaywall_views_nonacademic"] = len([e for e in unpaywall_events_this_week if not e.is_academic_location])

        sources = []
        for (source_id, num) in event_count_dict.iteritems():
            sources.append({
                "source_id": source_id,
                "event_count": num
            })
        self.sources = sources

    def get_event_count(self, source_id):
        try:
            return [s["event_count"] for s in self.sources if s["source_id"]==source_id][0]
        except IndexError:
            return 0

    def run(self):
        self.updated = datetime.datetime.utcnow()
        url = "http://api.oadoi.org/v2/{}?email=paperbuzz@impactstory.org".format(self.id)
        r = requests.get(url)
        self.oadoi_api_raw = r.json()
        if self.oadoi_api_raw and "is_oa" in self.oadoi_api_raw:
            self.is_open_access = self.oadoi_api_raw["is_oa"]

        self.set_sources()
        self.num_twitter_events = self.get_event_count("twitter")

        self.num_unpaywall_events = self.get_event_count("unpaywall_views")
        self.num_academic_unpaywall_events = self.get_event_count("unpaywall_views_academic")
        self.num_nonacademic_unpaywall_events = self.get_event_count("unpaywall_views_nonacademic")
        if self.num_unpaywall_events:
            self.radio_academic_unpaywall_events = float(self.num_academic_unpaywall_events)/self.num_unpaywall_events

    def run_mendeley(self):
        self.updated = datetime.datetime.utcnow()
        self.mendeley_api_raw = set_mendeley_data(self.id)

        if not self.mendeley_api_raw:
            return
        if not self.mendeley_api_raw.get("reader_count_by_subdiscipline", None):
            return

        self.num_main_discipline = 0
        self.main_discipline = None
        discipline_dict = defaultdict(int)
        for discipline in self.mendeley_api_raw["reader_count_by_subdiscipline"]:
            discipline_dict[discipline_lookup[discipline]] += self.mendeley_api_raw["reader_count_by_subdiscipline"][discipline][discipline]
        for (discipline, num) in discipline_dict.iteritems():
            if num > self.num_main_discipline:
                self.num_main_discipline = num
                self.main_discipline = discipline
        logger.info(u"discipline for {} is {}={}".format(self.id, self.main_discipline, self.num_main_discipline))

    def sources_dicts_no_events(self):
        my_doi_obj = Doi(self.id)
        my_doi_obj.altmetrics.get()
        sources_dicts_with_events = my_doi_obj.altmetrics_dict_including_unpaywall_views()
        sources_dicts_without_events = []
        for sources_dict in sources_dicts_with_events["sources"]:
            sources_dict_new = sources_dict
            del sources_dict_new["events"]
            sources_dicts_without_events.append(sources_dict_new)
        return {"sources": sources_dicts_without_events}

    def to_dict_hotness(self):
        metadata_keys = ["year", "journal_name", "journal_authors", "title"]
        open_access_keys = ["best_oa_location", "is_oa", "journal_is_oa"]
        metadata = dict((k, v) for k, v in self.oadoi_api_raw.iteritems() if k in metadata_keys)
        if self.mendeley_api_raw:
            metadata["abstract"] = self.mendeley_api_raw.get("abstract", None)
        else:
            metadata["abstract"] = None

        ret = {
            "doi": self.id,
            "metadata": metadata,
            "open_access": dict((k, v) for k, v in self.oadoi_api_raw.iteritems() if k in open_access_keys),
            "sources": self.sources
        }
        return ret

    def __repr__(self):
        return u"<WeeklyStats ({}, {}={})>".format(self.id, self.main_discipline, self.num_main_discipline)

def get_mendeley_session():
    mendeley_client = mendeley_lib.Mendeley(
        client_id=os.getenv("MENDELEY_OAUTH2_CLIENT_ID"),
        client_secret=os.getenv("MENDELEY_OAUTH2_SECRET"))
    auth = mendeley_client.start_client_credentials_flow()
    session = auth.authenticate()
    return session

def set_mendeley_data(doi):

    resp = None
    doc = None

    try:
        mendeley_session = get_mendeley_session()
        method = "doi"
        try:
            doc = mendeley_session.catalog.by_identifier(
                    doi=doi,
                    view='stats')
        except (UnicodeEncodeError, IndexError):
            return None

        if not doc:
            return None

        # print u"\nMatch! got the mendeley paper! for title {}".format(biblio_title)
        # print "got mendeley for {} using {}".format(product.id, method)
        resp = {}
        resp["reader_count"] = doc.reader_count
        resp["reader_count_by_academic_status"] = doc.reader_count_by_academic_status
        resp["reader_count_by_subdiscipline"] = doc.reader_count_by_subdiscipline
        resp["reader_count_by_country"] = doc.reader_count_by_country
        resp["mendeley_url"] = doc.link
        resp["abstract"] = doc.abstract
        resp["method"] = method

    except (KeyError, MendeleyException):
        pass

    return resp

<commit_msg>rename ratio<commit_after>from __future__ import absolute_import

# need to get system mendeley library
from mendeley.exception import MendeleyException
import mendeley as mendeley_lib
import os
import datetime
from sqlalchemy.dialects.postgresql import JSONB
from collections import defaultdict
import requests

from app import db
from app import logger
from doi import Doi
from event import UnpaywallEvent
from event import CedEvent
from util import remove_punctuation

discipline_lookup = {
    u'Arts and Humanities': u'arts and humanities',
    u'Design': u'arts and humanities',
    u'Philosophy': u'arts and humanities',
    u'Agricultural and Biological Sciences': u'life science',
    u'Biochemistry, Genetics and Molecular Biology': u'life science',
    u'Immunology and Microbiology': u'life science',
    u'Veterinary Science and Veterinary Medicine': u'life science',
    u'Business, Management and Accounting': u'business',
    u'Decision Sciences': u'business',
    u'Economics, Econometrics and Finance': u'business',
    u'Chemistry': u'chemistry',
    u'Computer Science': u'computing',
    u'Chemical Engineering': u'engineering',
    u'Engineering': u'engineering',
    u'Materials Science': u'engineering',
    u'Earth and Planetary Sciences': u'environment',
    u'Energy': u'environment',
    u'Environmental Science': u'environment',
    u'Medicine and Dentistry': u'health',
    u'Nursing and Health Professions': u'health',
    u'Pharmacology, Toxicology and Pharmaceutical Science': u'health',
    u'Sports and Recreations': u'health',
    u'Mathematics': u'mathematics',
    u'Physics and Astronomy': u'physics and astronomy',
    u'Neuroscience': u'brain and mind',
    u'Psychology': u'brain and mind',
    u'Linguistics': u'social science',
    u'Social Sciences': u'social science',
    u'Unspecified': u'unspecified'
}

class WeeklyStats(db.Model):
    id = db.Column(db.Text, primary_key=True)
    updated = db.Column(db.DateTime)
    mendeley_api_raw = db.Column(JSONB)
    oadoi_api_raw = db.Column(JSONB)
    sources = db.Column(JSONB)
    main_discipline = db.Column(db.Text)
    num_main_discipline = db.Column(db.Numeric)
    num_unpaywall_events = db.Column(db.Numeric)
    num_twitter_events = db.Column(db.Numeric)
    num_ced_events = db.Column(db.Numeric)
    num_academic_unpaywall_events = db.Column(db.Numeric)
    num_nonacademic_unpaywall_events = db.Column(db.Numeric)
    ratio_academic_unpaywall_events = db.Column(db.Numeric)
    is_open_access = db.Column(db.Boolean)
    week = db.Column(db.Numeric)

    def __init__(self, **kwargs):
        self.updated = datetime.datetime.utcnow()
        super(WeeklyStats, self).__init__(**kwargs)

    def set_sources(self):
        event_count_dict = defaultdict(int)

        ced_events = CedEvent.query.filter(CedEvent.doi==self.id).all()
        ced_events_this_week = [e for e in ced_events if e.week==self.week]
        for e in ced_events_this_week:
            event_count_dict[e.source_id] += 1

        unpaywall_events = UnpaywallEvent.query.filter(UnpaywallEvent.doi==self.id).all()
        unpaywall_events_this_week = [e for e in unpaywall_events if e.week==self.week]
        event_count_dict["unpaywall_views"] = len(unpaywall_events_this_week)
        event_count_dict["unpaywall_views_academic"] = len([e for e in unpaywall_events_this_week if e.is_academic_location])
        event_count_dict["unpaywall_views_nonacademic"] = len([e for e in unpaywall_events_this_week if not e.is_academic_location])

        sources = []
        for (source_id, num) in event_count_dict.iteritems():
            sources.append({
                "source_id": source_id,
                "event_count": num
            })
        self.sources = sources

    def get_event_count(self, source_id):
        try:
            return [s["event_count"] for s in self.sources if s["source_id"]==source_id][0]
        except IndexError:
            return 0

    def run(self):
        self.updated = datetime.datetime.utcnow()
        url = "http://api.oadoi.org/v2/{}?email=paperbuzz@impactstory.org".format(self.id)
        r = requests.get(url)
        self.oadoi_api_raw = r.json()
        if self.oadoi_api_raw and "is_oa" in self.oadoi_api_raw:
            self.is_open_access = self.oadoi_api_raw["is_oa"]

        self.set_sources()
        self.num_twitter_events = self.get_event_count("twitter")

        self.num_unpaywall_events = self.get_event_count("unpaywall_views")
        self.num_academic_unpaywall_events = self.get_event_count("unpaywall_views_academic")
        self.num_nonacademic_unpaywall_events = self.get_event_count("unpaywall_views_nonacademic")
        if self.num_unpaywall_events:
            self.ratio_academic_unpaywall_events = float(self.num_academic_unpaywall_events)/self.num_unpaywall_events

    def run_mendeley(self):
        self.updated = datetime.datetime.utcnow()
        self.mendeley_api_raw = set_mendeley_data(self.id)

        if not self.mendeley_api_raw:
            return
        if not self.mendeley_api_raw.get("reader_count_by_subdiscipline", None):
            return

        self.num_main_discipline = 0
        self.main_discipline = None
        discipline_dict = defaultdict(int)
        for discipline in self.mendeley_api_raw["reader_count_by_subdiscipline"]:
            discipline_dict[discipline_lookup[discipline]] += self.mendeley_api_raw["reader_count_by_subdiscipline"][discipline][discipline]
        for (discipline, num) in discipline_dict.iteritems():
            if num > self.num_main_discipline:
                self.num_main_discipline = num
                self.main_discipline = discipline
        logger.info(u"discipline for {} is {}={}".format(self.id, self.main_discipline, self.num_main_discipline))

    def sources_dicts_no_events(self):
        my_doi_obj = Doi(self.id)
        my_doi_obj.altmetrics.get()
        sources_dicts_with_events = my_doi_obj.altmetrics_dict_including_unpaywall_views()
        sources_dicts_without_events = []
        for sources_dict in sources_dicts_with_events["sources"]:
            sources_dict_new = sources_dict
            del sources_dict_new["events"]
            sources_dicts_without_events.append(sources_dict_new)
        return {"sources": sources_dicts_without_events}

    def to_dict_hotness(self):
        metadata_keys = ["year", "journal_name", "journal_authors", "title"]
        open_access_keys = ["best_oa_location", "is_oa", "journal_is_oa"]
        metadata = dict((k, v) for k, v in self.oadoi_api_raw.iteritems() if k in metadata_keys)
        if self.mendeley_api_raw:
            metadata["abstract"] = self.mendeley_api_raw.get("abstract", None)
        else:
            metadata["abstract"] = None

        ret = {
            "doi": self.id,
            "metadata": metadata,
            "open_access": dict((k, v) for k, v in self.oadoi_api_raw.iteritems() if k in open_access_keys),
            "sources": self.sources
        }
        return ret

    def __repr__(self):
        return u"<WeeklyStats ({}, {}={})>".format(self.id, self.main_discipline, self.num_main_discipline)

def get_mendeley_session():
    mendeley_client = mendeley_lib.Mendeley(
        client_id=os.getenv("MENDELEY_OAUTH2_CLIENT_ID"),
        client_secret=os.getenv("MENDELEY_OAUTH2_SECRET"))
    auth = mendeley_client.start_client_credentials_flow()
    session = auth.authenticate()
    return session

def set_mendeley_data(doi):

    resp = None
    doc = None

    try:
        mendeley_session = get_mendeley_session()
        method = "doi"
        try:
            doc = mendeley_session.catalog.by_identifier(
                    doi=doi,
                    view='stats')
        except (UnicodeEncodeError, IndexError):
            return None

        if not doc:
            return None

        # print u"\nMatch! got the mendeley paper! for title {}".format(biblio_title)
        # print "got mendeley for {} using {}".format(product.id, method)
        resp = {}
        resp["reader_count"] = doc.reader_count
        resp["reader_count_by_academic_status"] = doc.reader_count_by_academic_status
        resp["reader_count_by_subdiscipline"] = doc.reader_count_by_subdiscipline
        resp["reader_count_by_country"] = doc.reader_count_by_country
        resp["mendeley_url"] = doc.link
        resp["abstract"] = doc.abstract
        resp["method"] = method

    except (KeyError, MendeleyException):
        pass

    return resp

<|endoftext|>


<commit_before>import itertools
import numpy as np

import chainer
from chainer.backends import cuda
import chainer.functions as F
from chainer.links import Convolution2D

from chainercv.links import Conv2DBNActiv
from chainercv import transforms
from chainercv import utils

from chainercv.links.model.ssd.ssd_vgg16 import _check_pretrained_model

def _leaky_relu(x):
    return F.leaky_relu(x, slope=0.1)

def _upsample(x):
    return F.unpooling_2d(x, 2, cover_all=False)

class ResidualBlock(chainer.ChainList):

    def __init__(self, *links):
        super().__init__(*links)

    def __call__(self, x):
        h = x
        for link in self:
            h = link(h)
        h += x
        return h

class Darknet53Extractor(chainer.ChainList):
    """A Darknet53 based feature extractor for YOLOv3.

    This is a feature extractor for :class:`~chainercv.links.model.yolo.YOLOv3`
    """

    insize = 416
    grids = (13, 26, 52)

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

        # Darknet53
        self.append(Conv2DBNActiv(32, 3, pad=1, activ=_leaky_relu))
        for k, n_block in enumerate((1, 2, 8, 8, 4)):
            self.append(Conv2DBNActiv(
                32 << (k + 1), 3, stride=2, pad=1, activ=_leaky_relu))
            for _ in range(n_block):
                self.append(ResidualBlock(
                    Conv2DBNActiv(32 << k, 1, activ=_leaky_relu),
                    Conv2DBNActiv(32 << (k + 1), 3, pad=1, activ=_leaky_relu)))

        # additional links
        for i, n in enumerate((512, 256, 128)):
            if i > 0:
                self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))

    def __call__(self, x):
        """Compute feature maps from a batch of images.

        This method extracts feature maps from 3 layers.

        Args:
            x (ndarray): An array holding a batch of images.
                The images should be resized to :math:`416\\times 416`.

        Returns:
            list of Variable:
            Each variable contains a feature map.
        """

        ys = []
        h = x
        hs = []
        for i, link in enumerate(self):
            h = link(h)
            if i in {33, 39, 45}:
                ys.append(h)
            elif i in {14, 23}:
                hs.append(h)
            elif i in {34, 40}:
                h = F.concat((_upsample(h), hs.pop()))
        return ys

class YOLOv3(chainer.Chain):
    """YOLOv3.

    This is a model of YOLOv3 [#]_.
    This model uses :class:`~chainercv.links.model.yolo.Darknet53Extractor` as
    its feature extractor.

    .. [#] Joseph Redmon, Ali Farhadi.
       YOLOv3: An Incremental Improvement. arXiv 2018.

    Args:
       n_fg_class (int): The number of classes excluding the background.
       pretrained_model (str): The weight file to be loaded.
           This can take :obj:`'voc0712'`, `filepath` or :obj:`None`.
           The default value is :obj:`None`.

            * :obj:`'voc0712'`: Load weights trained on trainval split of \
                PASCAL VOC 2007 and 2012. \
                The weight file is downloaded and cached automatically. \
                :obj:`n_fg_class` must be :obj:`20` or :obj:`None`. \
                These weights were converted from the darknet model. \
                The conversion code is \
                `chainercv/examples/yolo/darknet2npz.py`.
            * `filepath`: A path of npz file. In this case, :obj:`n_fg_class` \
                must be specified properly.
            * :obj:`None`: Do not load weights.

    """

    _models = {
        'voc0712': {
            'n_fg_class': 20,
            'url': 'https://github.com/yuyu2172/share-weights/releases/'
            'download/0.0.6/yolov3_voc0712_2018_05_01.npz'
        },
    }

    anchors = (
        ((90, 116), (198, 156), (326, 373)),
        ((61, 30), (45, 62), (119, 59)),
        ((13, 10), (30, 16), (23, 33)))

    def __init__(self, n_fg_class=None, pretrained_model=None):
        super().__init__()

        n_fg_class, path = _check_pretrained_model(
            n_fg_class, pretrained_model, self._models)

        self.n_fg_class = n_fg_class
        self.use_preset('visualize')

        with self.init_scope():
            self.extractor = Darknet53Extractor()
            self.subnet = chainer.ChainList()

        for n in (512, 256, 128):
            self.subnet.append(chainer.Sequential(
                Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu),
                Convolution2D(3 * (4 + 1 + self.n_fg_class), 1)))

        default_bbox = []
        step = []
        for k, grid in enumerate(self.extractor.grids):
            for v, u in itertools.product(range(grid), repeat=2):
                for h, w in self.anchors[k]:
                    default_bbox.append((v, u, h, w))
                    step.append(self.insize / grid)
        self._default_bbox = np.array(default_bbox, dtype=np.float32)
        self._step = np.array(step, dtype=np.float32)

        if path:
            chainer.serializers.load_npz(path, self, strict=False)

    @property
    def insize(self):
        return self.extractor.insize

    def to_cpu(self):
        super().to_cpu()
        self._default_bbox = cuda.to_cpu(self._default_bbox)
        self._step = cuda.to_cpu(self._step)

    def to_gpu(self, device=None):
        super().to_gpu(device)
        self._default_bbox = cuda.to_gpu(self._default_bbox, device)
        self._step = cuda.to_gpu(self._step, device)

    def __call__(self, x):
        """Compute localization and classification from a batch of images.

        This method computes a variable.
        :func:`self._decode` converts this variable to bounding box
        coordinates and confidence scores.
        This variable is also used in training YOLOv3.

        Args:
            x (chainer.Variable): A variable holding a batch of images.
                The images are preprocessed by :meth:`_prepare`.

        Returns:
            chainer.Variable:
            A variable of float arrays of shape
            :math:`(B, K, 4 + 1 + n\_fg\_class)`,
            where :math:`B` is the number of samples in the batch and
            :math:`K` is the number of default bounding boxes.
        """

        ys = []
        for i, h in enumerate(self.extractor(x)):
            h = self.subnet[i](h)
            h = F.transpose(h, (0, 2, 3, 1))
            h = F.reshape(h, (h.shape[0], -1, 4 + 1 + self.n_fg_class))
            ys.append(h)
        return F.concat(ys)

    def use_preset(self, preset):
        """Use the given preset during prediction.

        This method changes values of :obj:`nms_thresh` and
        :obj:`score_thresh`. These values are a threshold value
        used for non maximum suppression and a threshold value
        to discard low confidence proposals in :meth:`predict`,
        respectively.

        If the attributes need to be changed to something
        other than the values provided in the presets, please modify
        them by directly accessing the public attributes.

        Args:
            preset ({'visualize', 'evaluate'}): A string to determine the
                preset to use.
        """

        if preset == 'visualize':
            self.nms_thresh = 0.45
            self.score_thresh = 0.5
        elif preset == 'evaluate':
            self.nms_thresh = 0.45
            self.score_thresh = 0.005
        else:
            raise ValueError('preset must be visualize or evaluate')

    def _decode(self, loc, conf):
        raw_bbox = self._default_bbox.copy()
        raw_bbox[:, :2] += 1 / (1 + self.xp.exp(-loc[:, :2]))
        raw_bbox[:, :2] *= self._step[:, None]
        raw_bbox[:, 2:] *= self.xp.exp(loc[:, 2:])
        raw_bbox[:, :2] -= raw_bbox[:, 2:] / 2
        raw_bbox[:, 2:] += raw_bbox[:, :2]

        conf = 1 / (1 + self.xp.exp(-conf))
        raw_score = conf[:, 0, None] * conf[:, 1:]

        bbox = []
        label = []
        score = []
        for l in range(self.n_fg_class):
            bbox_l = raw_bbox
            score_l = raw_score[:, l]

            mask = score_l >= self.score_thresh
            bbox_l = bbox_l[mask]
            score_l = score_l[mask]

            indices = utils.non_maximum_suppression(
                bbox_l, self.nms_thresh, score_l)
            bbox_l = bbox_l[indices]
            score_l = score_l[indices]

            bbox.append(cuda.to_cpu(bbox_l))
            label.append(np.array((l,) * len(bbox_l)))
            score.append(cuda.to_cpu(score_l))

        bbox = np.vstack(bbox).astype(np.float32)
        label = np.hstack(label).astype(np.int32)
        score = np.hstack(score).astype(np.float32)

        return bbox, label, score

    def predict(self, imgs):
        """Detect objects from images.

        This method predicts objects for each image.

        Args:
            imgs (iterable of numpy.ndarray): Arrays holding images.
                All images are in CHW and RGB format
                and the range of their value is :math:`[0, 255]`.

        Returns:
           tuple of lists:
           This method returns a tuple of three lists,
           :obj:`(bboxes, labels, scores)`.

           * **bboxes**: A list of float arrays of shape :math:`(R, 4)`, \
               where :math:`R` is the number of bounding boxes in a image. \
               Each bouding box is organized by \
               :math:`(y_{min}, x_{min}, y_{max}, x_{max})` \
               in the second axis.
           * **labels** : A list of integer arrays of shape :math:`(R,)`. \
               Each value indicates the class of the bounding box. \
               Values are in range :math:`[0, L - 1]`, where :math:`L` is the \
               number of the foreground classes.
           * **scores** : A list of float arrays of shape :math:`(R,)`. \
               Each value indicates how confident the prediction is.

        """

        x = []
        params = []
        for img in imgs:
            _, H, W = img.shape
            img, param = transforms.resize_contain(
                img / 255, (self.insize, self.insize), fill=0.5,
                return_param=True)
            x.append(self.xp.array(img))
            param['size'] = (H, W)
            params.append(param)

        with chainer.using_config('train', False), \
                chainer.function.no_backprop_mode():
            y = self(self.xp.stack(x)).array
        locs = y[:, :, :4]
        confs = y[:, :, 4:]

        bboxes = []
        labels = []
        scores = []
        for loc, conf, param in zip(locs, confs, params):
            bbox, label, score = self._decode(loc, conf)

            bbox = transforms.translate_bbox(
                bbox, -self.insize / 2, -self.insize / 2)
            bbox = transforms.resize_bbox(
                bbox, param['scaled_size'], param['size'])
            bbox = transforms.translate_bbox(
                bbox, param['size'][0] / 2, param['size'][1] / 2)

            bboxes.append(bbox)
            labels.append(label)
            scores.append(score)

        return bboxes, labels, scores
<commit_msg>add docs<commit_after>import itertools
import numpy as np

import chainer
from chainer.backends import cuda
import chainer.functions as F
from chainer.links import Convolution2D

from chainercv.links import Conv2DBNActiv
from chainercv import transforms
from chainercv import utils

from chainercv.links.model.ssd.ssd_vgg16 import _check_pretrained_model

def _leaky_relu(x):
    return F.leaky_relu(x, slope=0.1)

def _upsample(x):
    return F.unpooling_2d(x, 2, cover_all=False)

class ResidualBlock(chainer.ChainList):
    """ChainList with a residual connection."""

    def __init__(self, *links):
        super().__init__(*links)

    def __call__(self, x):
        h = x
        for link in self:
            h = link(h)
        h += x
        return h

class Darknet53Extractor(chainer.ChainList):
    """A Darknet53 based feature extractor for YOLOv3.

    This is a feature extractor for :class:`~chainercv.links.model.yolo.YOLOv3`
    """

    insize = 416
    grids = (13, 26, 52)

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

        # Darknet53
        self.append(Conv2DBNActiv(32, 3, pad=1, activ=_leaky_relu))
        for k, n_block in enumerate((1, 2, 8, 8, 4)):
            self.append(Conv2DBNActiv(
                32 << (k + 1), 3, stride=2, pad=1, activ=_leaky_relu))
            for _ in range(n_block):
                self.append(ResidualBlock(
                    Conv2DBNActiv(32 << k, 1, activ=_leaky_relu),
                    Conv2DBNActiv(32 << (k + 1), 3, pad=1, activ=_leaky_relu)))

        # additional links
        for i, n in enumerate((512, 256, 128)):
            if i > 0:
                self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu))
            self.append(Conv2DBNActiv(n, 1, activ=_leaky_relu))

    def __call__(self, x):
        """Compute feature maps from a batch of images.

        This method extracts feature maps from 3 layers.

        Args:
            x (ndarray): An array holding a batch of images.
                The images should be resized to :math:`416\\times 416`.

        Returns:
            list of Variable:
            Each variable contains a feature map.
        """

        ys = []
        h = x
        hs = []
        for i, link in enumerate(self):
            h = link(h)
            if i in {33, 39, 45}:
                ys.append(h)
            elif i in {14, 23}:
                hs.append(h)
            elif i in {34, 40}:
                h = F.concat((_upsample(h), hs.pop()))
        return ys

class YOLOv3(chainer.Chain):
    """YOLOv3.

    This is a model of YOLOv3 [#]_.
    This model uses :class:`~chainercv.links.model.yolo.Darknet53Extractor` as
    its feature extractor.

    .. [#] Joseph Redmon, Ali Farhadi.
       YOLOv3: An Incremental Improvement. arXiv 2018.

    Args:
       n_fg_class (int): The number of classes excluding the background.
       pretrained_model (str): The weight file to be loaded.
           This can take :obj:`'voc0712'`, `filepath` or :obj:`None`.
           The default value is :obj:`None`.

            * :obj:`'voc0712'`: Load weights trained on trainval split of \
                PASCAL VOC 2007 and 2012. \
                The weight file is downloaded and cached automatically. \
                :obj:`n_fg_class` must be :obj:`20` or :obj:`None`. \
                These weights were converted from the darknet model. \
                The conversion code is \
                `chainercv/examples/yolo/darknet2npz.py`.
            * `filepath`: A path of npz file. In this case, :obj:`n_fg_class` \
                must be specified properly.
            * :obj:`None`: Do not load weights.

    """

    _models = {
        'voc0712': {
            'n_fg_class': 20,
            'url': 'https://github.com/yuyu2172/share-weights/releases/'
            'download/0.0.6/yolov3_voc0712_2018_05_01.npz'
        },
    }

    anchors = (
        ((90, 116), (198, 156), (326, 373)),
        ((61, 30), (45, 62), (119, 59)),
        ((13, 10), (30, 16), (23, 33)))

    def __init__(self, n_fg_class=None, pretrained_model=None):
        super().__init__()

        n_fg_class, path = _check_pretrained_model(
            n_fg_class, pretrained_model, self._models)

        self.n_fg_class = n_fg_class
        self.use_preset('visualize')

        with self.init_scope():
            self.extractor = Darknet53Extractor()
            self.subnet = chainer.ChainList()

        for n in (512, 256, 128):
            self.subnet.append(chainer.Sequential(
                Conv2DBNActiv(n * 2, 3, pad=1, activ=_leaky_relu),
                Convolution2D(3 * (4 + 1 + self.n_fg_class), 1)))

        default_bbox = []
        step = []
        for k, grid in enumerate(self.extractor.grids):
            for v, u in itertools.product(range(grid), repeat=2):
                for h, w in self.anchors[k]:
                    default_bbox.append((v, u, h, w))
                    step.append(self.insize / grid)
        self._default_bbox = np.array(default_bbox, dtype=np.float32)
        self._step = np.array(step, dtype=np.float32)

        if path:
            chainer.serializers.load_npz(path, self, strict=False)

    @property
    def insize(self):
        return self.extractor.insize

    def to_cpu(self):
        super().to_cpu()
        self._default_bbox = cuda.to_cpu(self._default_bbox)
        self._step = cuda.to_cpu(self._step)

    def to_gpu(self, device=None):
        super().to_gpu(device)
        self._default_bbox = cuda.to_gpu(self._default_bbox, device)
        self._step = cuda.to_gpu(self._step, device)

    def __call__(self, x):
        """Compute localization and classification from a batch of images.

        This method computes a variable.
        :func:`self._decode` converts this variable to bounding box
        coordinates and confidence scores.
        This variable is also used in training YOLOv3.

        Args:
            x (chainer.Variable): A variable holding a batch of images.
                The images are preprocessed by :meth:`_prepare`.

        Returns:
            chainer.Variable:
            A variable of float arrays of shape
            :math:`(B, K, 4 + 1 + n\_fg\_class)`,
            where :math:`B` is the number of samples in the batch and
            :math:`K` is the number of default bounding boxes.
        """

        ys = []
        for i, h in enumerate(self.extractor(x)):
            h = self.subnet[i](h)
            h = F.transpose(h, (0, 2, 3, 1))
            h = F.reshape(h, (h.shape[0], -1, 4 + 1 + self.n_fg_class))
            ys.append(h)
        return F.concat(ys)

    def use_preset(self, preset):
        """Use the given preset during prediction.

        This method changes values of :obj:`nms_thresh` and
        :obj:`score_thresh`. These values are a threshold value
        used for non maximum suppression and a threshold value
        to discard low confidence proposals in :meth:`predict`,
        respectively.

        If the attributes need to be changed to something
        other than the values provided in the presets, please modify
        them by directly accessing the public attributes.

        Args:
            preset ({'visualize', 'evaluate'}): A string to determine the
                preset to use.
        """

        if preset == 'visualize':
            self.nms_thresh = 0.45
            self.score_thresh = 0.5
        elif preset == 'evaluate':
            self.nms_thresh = 0.45
            self.score_thresh = 0.005
        else:
            raise ValueError('preset must be visualize or evaluate')

    def _decode(self, loc, conf):
        raw_bbox = self._default_bbox.copy()
        raw_bbox[:, :2] += 1 / (1 + self.xp.exp(-loc[:, :2]))
        raw_bbox[:, :2] *= self._step[:, None]
        raw_bbox[:, 2:] *= self.xp.exp(loc[:, 2:])
        raw_bbox[:, :2] -= raw_bbox[:, 2:] / 2
        raw_bbox[:, 2:] += raw_bbox[:, :2]

        conf = 1 / (1 + self.xp.exp(-conf))
        raw_score = conf[:, 0, None] * conf[:, 1:]

        bbox = []
        label = []
        score = []
        for l in range(self.n_fg_class):
            bbox_l = raw_bbox
            score_l = raw_score[:, l]

            mask = score_l >= self.score_thresh
            bbox_l = bbox_l[mask]
            score_l = score_l[mask]

            indices = utils.non_maximum_suppression(
                bbox_l, self.nms_thresh, score_l)
            bbox_l = bbox_l[indices]
            score_l = score_l[indices]

            bbox.append(cuda.to_cpu(bbox_l))
            label.append(np.array((l,) * len(bbox_l)))
            score.append(cuda.to_cpu(score_l))

        bbox = np.vstack(bbox).astype(np.float32)
        label = np.hstack(label).astype(np.int32)
        score = np.hstack(score).astype(np.float32)

        return bbox, label, score

    def predict(self, imgs):
        """Detect objects from images.

        This method predicts objects for each image.

        Args:
            imgs (iterable of numpy.ndarray): Arrays holding images.
                All images are in CHW and RGB format
                and the range of their value is :math:`[0, 255]`.

        Returns:
           tuple of lists:
           This method returns a tuple of three lists,
           :obj:`(bboxes, labels, scores)`.

           * **bboxes**: A list of float arrays of shape :math:`(R, 4)`, \
               where :math:`R` is the number of bounding boxes in a image. \
               Each bouding box is organized by \
               :math:`(y_{min}, x_{min}, y_{max}, x_{max})` \
               in the second axis.
           * **labels** : A list of integer arrays of shape :math:`(R,)`. \
               Each value indicates the class of the bounding box. \
               Values are in range :math:`[0, L - 1]`, where :math:`L` is the \
               number of the foreground classes.
           * **scores** : A list of float arrays of shape :math:`(R,)`. \
               Each value indicates how confident the prediction is.

        """

        x = []
        params = []
        for img in imgs:
            _, H, W = img.shape
            img, param = transforms.resize_contain(
                img / 255, (self.insize, self.insize), fill=0.5,
                return_param=True)
            x.append(self.xp.array(img))
            param['size'] = (H, W)
            params.append(param)

        with chainer.using_config('train', False), \
                chainer.function.no_backprop_mode():
            y = self(self.xp.stack(x)).array
        locs = y[:, :, :4]
        confs = y[:, :, 4:]

        bboxes = []
        labels = []
        scores = []
        for loc, conf, param in zip(locs, confs, params):
            bbox, label, score = self._decode(loc, conf)

            bbox = transforms.translate_bbox(
                bbox, -self.insize / 2, -self.insize / 2)
            bbox = transforms.resize_bbox(
                bbox, param['scaled_size'], param['size'])
            bbox = transforms.translate_bbox(
                bbox, param['size'][0] / 2, param['size'][1] / 2)

            bboxes.append(bbox)
            labels.append(label)
            scores.append(score)

        return bboxes, labels, scores
<|endoftext|>


<commit_before># Tai Sakuma <tai.sakuma@gmail.com>
import logging
import pytest
from collections import deque

try:
    import unittest.mock as mock
except ImportError:
    import mock

from alphatwirl.concurrently import TaskPackageDropbox

##__________________________________________________________________||
@pytest.fixture()
def packages():
    return [
        mock.Mock(name='package0'),
        mock.Mock(name='package1'),
        mock.Mock(name='package2'),
        mock.Mock(name='package3'),
        mock.Mock(name='package4'),
    ]

@pytest.fixture()
def pkgidx_result_pairs(results):
    # e.g., [(0, package0), (1, package1), ...]
    return list(enumerate(results))

@pytest.fixture()
def results():
    return [
        mock.Mock(name='result0'),
        mock.Mock(name='result1'),
        mock.Mock(name='result2'),
        mock.Mock(name='result3'),
        mock.Mock(name='result4'),
    ]

@pytest.fixture()
def collect_results(results):
    ret = {i: deque([r]) for i, r in enumerate(results)}
    # e.g., {0: deque([result0]), 1:deque([result1]), ...}

    # 3rd job fails twice before succeed
    ret[2].extendleft([None, None])
    # e.g., deque([None, None, result2])

    return ret

@pytest.fixture()
def workingarea(collect_results):
    ret = mock.MagicMock()
    ret.put_package.side_effect = [0, 1, 2, 3, 4] # package indices
    ret.collect_result.side_effect = lambda i: collect_results[i].popleft()
    return ret

@pytest.fixture()
def dispatcher():
    ret = mock.MagicMock()
    ret.run_multiple.return_value = [1000, 1001, 1002, 1003, 1004]

    # jobs finish in steps
    ret.poll.side_effect = [
        [1001, 1003],
        [ ],
        [1002],
        [1000, 1005],
        [1006, 1004]
    ]

    # for resubmission
    ret.run.side_effect = [1005, 1006]

    return ret

@pytest.fixture()
def obj(workingarea, dispatcher, packages):
    ret = TaskPackageDropbox(workingArea=workingarea, dispatcher=dispatcher, sleep=0.01)
    ret.open()
    ret.put_multiple(packages)
    yield ret
    ret.close()

##__________________________________________________________________||
def test_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        assert pkgidx_result_pairs == obj.receive()
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_dispatcher_received_failed_runids(obj, dispatcher):
    with mock.patch('time.sleep') as sleep:
        obj.receive()
        assert [mock.call([1002]), mock.call([1005])] == dispatcher.failed_runids.call_args_list
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_logging_resubmission(obj, caplog):
    with mock.patch('time.sleep') as sleep:
        with caplog.at_level(logging.WARNING):
            obj.receive()
        assert len(caplog.records) == 2
        assert caplog.records[0].levelname == 'WARNING'
        assert caplog.records[1].levelname == 'WARNING'
        assert 'TaskPackageDropbox' in caplog.records[0].name
        assert 'TaskPackageDropbox' in caplog.records[1].name
        assert 'resubmitting' in caplog.records[0].msg
        assert 'resubmitting' in caplog.records[1].msg
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_in_one_step(obj, pkgidx_result_pairs, dispatcher, collect_results):
    with mock.patch('time.sleep') as sleep:
        # make all jobs finish by the first poll
        dispatcher.poll.side_effect = [[1000, 1001, 1002, 1003, 1004]]

        # make the 3rd job successful by removing two None's
        collect_results[2].popleft() # deque([None, result2])
        collect_results[2].popleft() # deque([result2])

        assert pkgidx_result_pairs == obj.receive()
        assert [] == sleep.call_args_list

##__________________________________________________________________||
def test_poll(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            actual.extend(obj.poll())
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [] == sleep.call_args_list

def test_poll_then_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        actual.extend(obj.poll())
        actual.extend(obj.receive())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*3 == sleep.call_args_list

##__________________________________________________________________||
def test_receive_one(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            pair = obj.receive_one()
            if pair is None:
                break
            actual.append(pair)
        assert obj.receive_one() is None
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*4 == sleep.call_args_list

@pytest.mark.parametrize('dispatcher_poll', [
    pytest.param(
        [[1001, 1003], [ ], [1002], [1000, 1005], [1006, 1004]],
        id='as_in_fixture'
    ),
    pytest.param(
        [[1001, 1003], [ ], [1002], [1000, 1005], [1006], [1004]],
        id='last_one'
    ),
    pytest.param(
        [[ ], [1001, 1003], [ ], [1002], [1000, 1005], [1006, 1004]],
        id='empty_first'
    ),
])
def test_receive_one_param(obj, pkgidx_result_pairs, dispatcher, dispatcher_poll):
    with mock.patch('time.sleep') as sleep:
        dispatcher.poll.side_effect = dispatcher_poll

        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            pair = obj.receive_one()
            if pair is None:
                break
            actual.append(pair)
        assert obj.receive_one() is None
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*(len(dispatcher_poll)-1) == sleep.call_args_list

def test_receive_one_then_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]

        actual.append(obj.receive_one())

        actual.extend(obj.receive())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_one_then_poll(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]

        actual.append(obj.receive_one())
        actual.append(obj.receive_one())

        actual.extend(obj.poll())
        actual.extend(obj.poll())
        actual.extend(obj.poll())
        actual.extend(obj.poll())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)] == sleep.call_args_list

##__________________________________________________________________||
<commit_msg>Test single call to receive_one<commit_after># Tai Sakuma <tai.sakuma@gmail.com>
import logging
import pytest
from collections import deque

try:
    import unittest.mock as mock
except ImportError:
    import mock

from alphatwirl.concurrently import TaskPackageDropbox

##__________________________________________________________________||
@pytest.fixture()
def packages():
    return [
        mock.Mock(name='package0'),
        mock.Mock(name='package1'),
        mock.Mock(name='package2'),
        mock.Mock(name='package3'),
        mock.Mock(name='package4'),
    ]

@pytest.fixture()
def pkgidx_result_pairs(results):
    # e.g., [(0, package0), (1, package1), ...]
    return list(enumerate(results))

@pytest.fixture()
def results():
    return [
        mock.Mock(name='result0'),
        mock.Mock(name='result1'),
        mock.Mock(name='result2'),
        mock.Mock(name='result3'),
        mock.Mock(name='result4'),
    ]

@pytest.fixture()
def collect_results(results):
    ret = {i: deque([r]) for i, r in enumerate(results)}
    # e.g., {0: deque([result0]), 1:deque([result1]), ...}

    # 3rd job fails twice before succeed
    ret[2].extendleft([None, None])
    # e.g., deque([None, None, result2])

    return ret

@pytest.fixture()
def workingarea(collect_results):
    ret = mock.MagicMock()
    ret.put_package.side_effect = [0, 1, 2, 3, 4] # package indices
    ret.collect_result.side_effect = lambda i: collect_results[i].popleft()
    return ret

@pytest.fixture()
def dispatcher():
    ret = mock.MagicMock()
    ret.run_multiple.return_value = [1000, 1001, 1002, 1003, 1004]

    # jobs finish in steps
    ret.poll.side_effect = [
        [1001, 1003],
        [ ],
        [1002],
        [1000, 1005],
        [1006, 1004]
    ]

    # for resubmission
    ret.run.side_effect = [1005, 1006]

    return ret

@pytest.fixture()
def obj(workingarea, dispatcher, packages):
    ret = TaskPackageDropbox(workingArea=workingarea, dispatcher=dispatcher, sleep=0.01)
    ret.open()
    ret.put_multiple(packages)
    yield ret
    ret.close()

##__________________________________________________________________||
def test_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        assert pkgidx_result_pairs == obj.receive()
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_dispatcher_received_failed_runids(obj, dispatcher):
    with mock.patch('time.sleep') as sleep:
        obj.receive()
        assert [mock.call([1002]), mock.call([1005])] == dispatcher.failed_runids.call_args_list
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_logging_resubmission(obj, caplog):
    with mock.patch('time.sleep') as sleep:
        with caplog.at_level(logging.WARNING):
            obj.receive()
        assert len(caplog.records) == 2
        assert caplog.records[0].levelname == 'WARNING'
        assert caplog.records[1].levelname == 'WARNING'
        assert 'TaskPackageDropbox' in caplog.records[0].name
        assert 'TaskPackageDropbox' in caplog.records[1].name
        assert 'resubmitting' in caplog.records[0].msg
        assert 'resubmitting' in caplog.records[1].msg
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_in_one_step(obj, pkgidx_result_pairs, dispatcher, collect_results):
    with mock.patch('time.sleep') as sleep:
        # make all jobs finish by the first poll
        dispatcher.poll.side_effect = [[1000, 1001, 1002, 1003, 1004]]

        # make the 3rd job successful by removing two None's
        collect_results[2].popleft() # deque([None, result2])
        collect_results[2].popleft() # deque([result2])

        assert pkgidx_result_pairs == obj.receive()
        assert [] == sleep.call_args_list

##__________________________________________________________________||
def test_poll(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            actual.extend(obj.poll())
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [] == sleep.call_args_list

def test_poll_then_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        actual.extend(obj.poll())
        actual.extend(obj.receive())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*3 == sleep.call_args_list

##__________________________________________________________________||
def test_receive_one(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            pair = obj.receive_one()
            if pair is None:
                break
            actual.append(pair)
        assert obj.receive_one() is None
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_one_single_call(obj, pkgidx_result_pairs, workingarea):
    with mock.patch('time.sleep') as sleep:
        pair = obj.receive_one()
        assert pair in pkgidx_result_pairs
        assert [mock.call(pair[0])] == workingarea.collect_result.call_args_list
        assert [] == sleep.call_args_list

@pytest.mark.parametrize('dispatcher_poll', [
    pytest.param(
        [[1001, 1003], [ ], [1002], [1000, 1005], [1006, 1004]],
        id='as_in_fixture'
    ),
    pytest.param(
        [[1001, 1003], [ ], [1002], [1000, 1005], [1006], [1004]],
        id='last_one'
    ),
    pytest.param(
        [[ ], [1001, 1003], [ ], [1002], [1000, 1005], [1006, 1004]],
        id='empty_first'
    ),
])
def test_receive_one_param(obj, pkgidx_result_pairs, dispatcher, dispatcher_poll):
    with mock.patch('time.sleep') as sleep:
        dispatcher.poll.side_effect = dispatcher_poll

        actual = [ ]
        while len(actual) < len(pkgidx_result_pairs):
            pair = obj.receive_one()
            if pair is None:
                break
            actual.append(pair)
        assert obj.receive_one() is None
        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*(len(dispatcher_poll)-1) == sleep.call_args_list

def test_receive_one_then_receive(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]

        actual.append(obj.receive_one())

        actual.extend(obj.receive())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)]*4 == sleep.call_args_list

def test_receive_one_then_poll(obj, pkgidx_result_pairs):
    with mock.patch('time.sleep') as sleep:
        actual = [ ]

        actual.append(obj.receive_one())
        actual.append(obj.receive_one())

        actual.extend(obj.poll())
        actual.extend(obj.poll())
        actual.extend(obj.poll())
        actual.extend(obj.poll())

        assert sorted(pkgidx_result_pairs) == sorted(actual)
        assert [mock.call(0.01)] == sleep.call_args_list

##__________________________________________________________________||
<|endoftext|>