<commit_before>import collections
import json
import os
import unittest
import mock

from .utils import placebo_session

from zappa.wsgi import create_wsgi_request, common_log
from zappa.zappa import Zappa, ASSUME_POLICY, ATTACH_POLICY

class TestZappa(unittest.TestCase):
    ##
    # Sanity Tests
    ##

    def test_test(self):
        self.assertTrue(True)
    ##
    # Basic Tests
    ##

    def test_zappa(self):
        self.assertTrue(True)
        Zappa()

    def test_create_lambda_package(self):
        self.assertTrue(True)
        z = Zappa()
        path = z.create_lambda_zip()
        self.assertTrue(os.path.isfile(path))
        os.remove(path)

    def test_load_credentials(self):
        z = Zappa()
        z.aws_region = 'us-east-1'
        z.load_credentials()
        self.assertEqual(z.boto_session.region_name, 'us-east-1')
        self.assertEqual(z.aws_region, 'us-east-1')

        z.aws_region = 'eu-west-1'
        z.load_credentials()
        self.assertEqual(z.boto_session.region_name, 'eu-west-1')
        self.assertEqual(z.aws_region, 'eu-west-1')

        creds = {
            'AWS_ACCESS_KEY_ID': 'AK123',
            'AWS_SECRET_ACCESS_KEY': 'JKL456',
            'AWS_DEFAULT_REGION': 'us-west-1'
        }
        with mock.patch.dict('os.environ', creds):
            z.aws_region = None
            z.load_credentials()
            loaded_creds = z.boto_session._session.get_credentials()

        self.assertEqual(loaded_creds.access_key, 'AK123')
        self.assertEqual(loaded_creds.secret_key, 'JKL456')
        self.assertEqual(z.boto_session.region_name, 'us-west-1')

    @placebo_session
    def test_upload_remove_s3(self, session):
        bucket_name = 'test_zappa_upload_s3'
        z = Zappa(session)
        zip_path = z.create_lambda_zip()
        res = z.upload_to_s3(zip_path, bucket_name)
        os.remove(zip_path)
        self.assertTrue(res)
        s3 = session.resource('s3')

        # will throw ClientError with 404 if bucket doesn't exist
        s3.meta.client.head_bucket(Bucket=bucket_name)

        # will throw ClientError with 404 if object doesn't exist
        s3.meta.client.head_object(
            Bucket=bucket_name,
            Key=zip_path,
        )
        res = z.remove_from_s3(zip_path, bucket_name)
        self.assertTrue(res)

    @placebo_session
    def test_create_iam_roles(self, session):
        z = Zappa(session)
        arn = z.create_iam_roles()
        self.assertEqual(arn, "arn:aws:iam::123:role/{}".format(z.role_name))

    @placebo_session
    def test_create_api_gateway_routes(self, session):
        z = Zappa(session)
        z.parameter_depth = 1
        z.integration_response_codes = [200]
        z.method_response_codes = [200]
        z.http_methods = ['GET']
        z.credentials_arn = 'arn:aws:iam::12345:role/ZappaLambdaExecution'
        lambda_arn = 'arn:aws:lambda:us-east-1:12345:function:helloworld'
        with mock.patch('time.time', return_value=123.456):
            api_id = z.create_api_gateway_routes(lambda_arn)
        self.assertEqual(api_id, 'j27idab94h')

    def test_policy_json(self):
        # ensure the policy docs are valid JSON
        json.loads(ASSUME_POLICY)
        json.loads(ATTACH_POLICY)

    ##
    # Logging
    ##

    def test_logging(self):
        """
        TODO
        """
        Zappa()

    ##
    # WSGI
    ##

    def test_wsgi_event(self):

        event = {
            "body": {},
            "headers": {
                "Via": "1.1 e604e934e9195aaf3e36195adbcb3e18.cloudfront.net (CloudFront)",
                "Accept-Language": "en-US,en;q=0.5",
                "Accept-Encoding": "gzip",
                "CloudFront-Is-SmartTV-Viewer": "false",
                "CloudFront-Forwarded-Proto": "https",
                "X-Forwarded-For": "109.81.209.118, 216.137.58.43",
                "CloudFront-Viewer-Country": "CZ",
                "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                "X-Forwarded-Proto": "https",
                "X-Amz-Cf-Id": "LZeP_TZxBgkDt56slNUr_H9CHu1Us5cqhmRSswOh1_3dEGpks5uW-g==",
                "CloudFront-Is-Tablet-Viewer": "false",
                "X-Forwarded-Port": "443",
                "CloudFront-Is-Mobile-Viewer": "false",
                "CloudFront-Is-Desktop-Viewer": "true"
            },
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {
                "dead": "beef"
            }
        }
        request = create_wsgi_request(event)

    def test_wsgi_path_info(self):
        # Test no parameters (site.com/)
        event = {
            "body": {},
            "headers": {},
            "params": {},
            "method": "GET",
            "query": {}
        }

        request = create_wsgi_request(event, trailing_slash=True)
        self.assertEqual("/", request['PATH_INFO'])

        request = create_wsgi_request(event, trailing_slash=False)
        self.assertEqual("/", request['PATH_INFO'])

        # Test parameters (site.com/asdf1/asdf2 or site.com/asdf1/asdf2/)
        event = {
            "body": {},
            "headers": {},
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {}
        }

        request = create_wsgi_request(event, trailing_slash=True)
        self.assertEqual("/asdf1/asdf2/", request['PATH_INFO'])

        request = create_wsgi_request(event, trailing_slash=False)
        self.assertEqual("/asdf1/asdf2", request['PATH_INFO'])

    def test_wsgi_logging(self):
        event = {
            "body": {},
            "headers": {},
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {}
        }
        environ = create_wsgi_request(event, trailing_slash=False)
        response_tuple = collections.namedtuple('Response', ['status_code', 'content'])
        response = response_tuple(200, 'hello')
        le = common_log(environ, response, response_time=True)
        le = common_log(environ, response, response_time=False)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Added WSGI tests.<commit_after># -*- coding: utf8 -*-
import collections
import json
import os
import unittest
import mock

from werkzeug.wrappers import Response

from .utils import placebo_session

from zappa.wsgi import create_wsgi_request, common_log
from zappa.zappa import Zappa, ASSUME_POLICY, ATTACH_POLICY
from zappa.middleware import ZappaWSGIMiddleware

class TestZappa(unittest.TestCase):
    ##
    # Sanity Tests
    ##

    def test_test(self):
        self.assertTrue(True)
    ##
    # Basic Tests
    ##

    def test_zappa(self):
        self.assertTrue(True)
        Zappa()

    def test_create_lambda_package(self):
        self.assertTrue(True)
        z = Zappa()
        path = z.create_lambda_zip()
        self.assertTrue(os.path.isfile(path))
        os.remove(path)

    def test_load_credentials(self):
        z = Zappa()
        z.aws_region = 'us-east-1'
        z.load_credentials()
        self.assertEqual(z.boto_session.region_name, 'us-east-1')
        self.assertEqual(z.aws_region, 'us-east-1')

        z.aws_region = 'eu-west-1'
        z.load_credentials()
        self.assertEqual(z.boto_session.region_name, 'eu-west-1')
        self.assertEqual(z.aws_region, 'eu-west-1')

        creds = {
            'AWS_ACCESS_KEY_ID': 'AK123',
            'AWS_SECRET_ACCESS_KEY': 'JKL456',
            'AWS_DEFAULT_REGION': 'us-west-1'
        }
        with mock.patch.dict('os.environ', creds):
            z.aws_region = None
            z.load_credentials()
            loaded_creds = z.boto_session._session.get_credentials()

        self.assertEqual(loaded_creds.access_key, 'AK123')
        self.assertEqual(loaded_creds.secret_key, 'JKL456')
        self.assertEqual(z.boto_session.region_name, 'us-west-1')

    @placebo_session
    def test_upload_remove_s3(self, session):
        bucket_name = 'test_zappa_upload_s3'
        z = Zappa(session)
        zip_path = z.create_lambda_zip()
        res = z.upload_to_s3(zip_path, bucket_name)
        os.remove(zip_path)
        self.assertTrue(res)
        s3 = session.resource('s3')

        # will throw ClientError with 404 if bucket doesn't exist
        s3.meta.client.head_bucket(Bucket=bucket_name)

        # will throw ClientError with 404 if object doesn't exist
        s3.meta.client.head_object(
            Bucket=bucket_name,
            Key=zip_path,
        )
        res = z.remove_from_s3(zip_path, bucket_name)
        self.assertTrue(res)

    @placebo_session
    def test_create_iam_roles(self, session):
        z = Zappa(session)
        arn = z.create_iam_roles()
        self.assertEqual(arn, "arn:aws:iam::123:role/{}".format(z.role_name))

    @placebo_session
    def test_create_api_gateway_routes(self, session):
        z = Zappa(session)
        z.parameter_depth = 1
        z.integration_response_codes = [200]
        z.method_response_codes = [200]
        z.http_methods = ['GET']
        z.credentials_arn = 'arn:aws:iam::12345:role/ZappaLambdaExecution'
        lambda_arn = 'arn:aws:lambda:us-east-1:12345:function:helloworld'
        with mock.patch('time.time', return_value=123.456):
            api_id = z.create_api_gateway_routes(lambda_arn)
        self.assertEqual(api_id, 'j27idab94h')

    def test_policy_json(self):
        # ensure the policy docs are valid JSON
        json.loads(ASSUME_POLICY)
        json.loads(ATTACH_POLICY)

    ##
    # Logging
    ##

    def test_logging(self):
        """
        TODO
        """
        Zappa()

    ##
    # WSGI
    ##

    def test_wsgi_event(self):

        event = {
            "body": {},
            "headers": {
                "Via": "1.1 e604e934e9195aaf3e36195adbcb3e18.cloudfront.net (CloudFront)",
                "Accept-Language": "en-US,en;q=0.5",
                "Accept-Encoding": "gzip",
                "CloudFront-Is-SmartTV-Viewer": "false",
                "CloudFront-Forwarded-Proto": "https",
                "X-Forwarded-For": "109.81.209.118, 216.137.58.43",
                "CloudFront-Viewer-Country": "CZ",
                "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                "X-Forwarded-Proto": "https",
                "X-Amz-Cf-Id": "LZeP_TZxBgkDt56slNUr_H9CHu1Us5cqhmRSswOh1_3dEGpks5uW-g==",
                "CloudFront-Is-Tablet-Viewer": "false",
                "X-Forwarded-Port": "443",
                "CloudFront-Is-Mobile-Viewer": "false",
                "CloudFront-Is-Desktop-Viewer": "true"
            },
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {
                "dead": "beef"
            }
        }
        request = create_wsgi_request(event)

    def test_wsgi_path_info(self):
        # Test no parameters (site.com/)
        event = {
            "body": {},
            "headers": {},
            "params": {},
            "method": "GET",
            "query": {}
        }

        request = create_wsgi_request(event, trailing_slash=True)
        self.assertEqual("/", request['PATH_INFO'])

        request = create_wsgi_request(event, trailing_slash=False)
        self.assertEqual("/", request['PATH_INFO'])

        # Test parameters (site.com/asdf1/asdf2 or site.com/asdf1/asdf2/)
        event = {
            "body": {},
            "headers": {},
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {}
        }

        request = create_wsgi_request(event, trailing_slash=True)
        self.assertEqual("/asdf1/asdf2/", request['PATH_INFO'])

        request = create_wsgi_request(event, trailing_slash=False)
        self.assertEqual("/asdf1/asdf2", request['PATH_INFO'])

    def test_wsgi_logging(self):
        event = {
            "body": {},
            "headers": {},
            "params": {
                "parameter_1": "asdf1",
                "parameter_2": "asdf2",
            },
            "method": "GET",
            "query": {}
        }
        environ = create_wsgi_request(event, trailing_slash=False)
        response_tuple = collections.namedtuple('Response', ['status_code', 'content'])
        response = response_tuple(200, 'hello')
        le = common_log(environ, response, response_time=True)
        le = common_log(environ, response, response_time=False)

class TestWSGIMiddleWare(unittest.TestCase):

    def setUp(self):
        print "setting up"
        self.headers = list()
        self.status = list()

    def _start_response(self, status, headers, exc_info=None):
        self.status[:] = [status]
        self.headers[:] = headers

    def test_wsgi_middleware_multiplecookies(self):
        def simple_app(environ, start_response):
            status = '200 OK'
            response_headers = [('Set-Cookie', 'foo=123'),
                                ('Set-Cookie', 'bar=456')]
            start_response(status, response_headers)
            return ['Hello zappa!']

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call with empty WSGI Environment
        resp = app(dict(), self._start_response)

        self.assertEqual(self.status[0], '200 OK')

        # Assert there is only one zappa cookie
        self.assertEqual(len(self.headers), 1)
        self.assertEqual(self.headers[0][0], 'Set-Cookie')
        self.assertTrue(self.headers[0][1].startswith('zappa='))

        self.assertEqual(''.join(resp), 'Hello zappa!')

    def test_wsgi_middleware_unpackcookies(self):
        # Setting the cookies
        def simple_app(environ, start_response):
            status = '200 OK'
            response_headers = [('Set-Cookie', 'foo=123'),
                                ('Set-Cookie', 'bar=456'),
                                ('Set-Cookie', 'baz=789'),
                                ('Set-Cookie', 'xyz=zappalicious')]
            start_response(status, response_headers)
            return ['Set cookies!']

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call with empty WSGI Environment
        resp = app(dict(), self._start_response)

        # Ensure the encoded zappa cookie is set
        self.assertEqual(self.headers[0][0], 'Set-Cookie')
        zappa_cookie = self.headers[0][1]
        self.assertTrue(zappa_cookie.startswith('zappa='))

        # Reads the hopefully decoded cookies
        def simple_app(environ, start_response):
            status = '200 OK'
            response_headers = []
            start_response(status, response_headers)
            return [environ['HTTP_COOKIE']]

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call the app with the encoded cookie in the environment
        resp = app({'HTTP_COOKIE': zappa_cookie}, self._start_response)

        # Assert that the simple_app, received the decoded cookies
        excpected = {'foo': '123', 'bar': '456', 'baz': '789',
                     'xyz': 'zappalicious'}
        received = ''.join(resp)
        # Split the response on ';', then on '=', then convert to dict for
        # comparison
        received = dict([cookie.split('=') for cookie in received.split(';')])
        self.assertDictEqual(received, excpected)

    def test_wsgi_middleware_cookieoverwrite(self):
        """ This method do:
            * Sets a bunch of cookies.
            * Fetches the zappa cookie from the response
            * Overwrites only some of the cookies.
            * Fetches the zappa cookie from the response
            * Reads the cookie
        """
        # Setting the cookies
        def set_cookies(environ, start_response):
            status = '200 OK'
            response_headers = [('Set-Cookie', 'foo=123'),
                                ('Set-Cookie', 'bar=456'),
                                ('Set-Cookie', 'baz=789'),
                                ('Set-Cookie', 'xyz=zappalicious')]
            start_response(status, response_headers)
            return ['Set cookies!']

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(set_cookies)

        # Call with empty WSGI Environment
        resp = app(dict(), self._start_response)

        # Retrieve the cookie
        zappa_cookie = self.headers[0][1]

        def change_cookies(environ, start_response):
            status = '200 OK'
            response_headers = [('Set-Cookie', 'foo=sdf'),
                                ('Set-Cookie', 'xyz=jkl')]
            start_response(status, response_headers)
            return ['Set cookies!']

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(change_cookies)

        # Call the app with the encoded cookie in the environment
        resp = app({'HTTP_COOKIE': zappa_cookie}, self._start_response)

        # Retrieve the cookie
        zappa_cookie = self.headers[0][1]

        # Reads the hopefully decoded cookies
        def read_cookies(environ, start_response):
            status = '200 OK'
            response_headers = []
            start_response(status, response_headers)
            return [environ['HTTP_COOKIE']]

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(read_cookies)

        # Call the app with the encoded cookie in the environment
        resp = app({'HTTP_COOKIE': zappa_cookie}, self._start_response)

        # Assert that read_cookies received the corrected decoded cookies
        excpected = {'foo': 'sdf', 'bar': '456', 'baz': '789',
                     'xyz': 'jkl'}
        received = ''.join(resp)
        # Split the response on ';', then on '=', then convert to dict for
        # comparison
        received = dict([cookie.split('=') for cookie in received.split(';')])
        self.assertDictEqual(received, excpected)

        # Call the app with the encoded cookie in the environment
        resp = app({'HTTP_COOKIE': zappa_cookie}, self._start_response)

        # Assert that read_cookies received the decoded cookies
        received = ''.join(resp)
        # Split the response on ';', then on '=', then convert to dict for
        # comparison
        received = dict([cookie.split('=') for cookie in received.split(';')])
        self.assertDictEqual(received, excpected)

    def test_wsgi_middleware_realcall(self):
        print("1: Setting the cookies.")
        event = {
            u'method': u'POST',
            u'params': {u'parameter_1': u'set_cookie'},
            u'body': u'foo=xxx&bar=yyy',
            u'headers': {
                # u'Cookie': u'zappa=8fFdF9whnxFRZaZbRgqvyb2wupbHuAjFLY21u5ukqvL4qCte6XJgduJU4edADwNwpZ2cwmizh7SQ4VhTuXy76KDaHVmAt4J4Pch1ZTjUHBmKmU52yoCpWYD', u'Accept': u'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
                # u'Content-Type': u'application/x-www-form-urlencoded'
            },
            u'query': {}}

        def set_cookies(environ, start_response):
            status = '200 OK'
            print environ
            response_headers = [('Set-Cookie', 'foo=123'),
                                ('Set-Cookie', 'bar=456'),
                                ('Set-Cookie', 'baz=789'),
                                ('Set-Cookie', 'xyz=zappalicious')]
            start_response(status, response_headers)
            return ['Set cookies!']

        app = ZappaWSGIMiddleware(set_cookies)

        environ = create_wsgi_request(event, script_name='http://zappa.com/',
                                      trailing_slash=False)

        response = Response.from_app(app, environ)

        # for (header, value) in response.headers:
        #     print header, value

        # Filter the headers for Set-Cookie header
        zappa_cookie = [x[1] for x in response.headers if x[0] == 'Set-Cookie']
        self.assertEqual(len(zappa_cookie), 1)
        zappa_cookie0 = zappa_cookie[0]
        self.assertTrue(zappa_cookie0.startswith('zappa='))

        print("2: Changing 1 cookie")
        event = {
            u'method': u'POST',
            u'params': {u'parameter_1': u'set_cookie'},
            u'body': u'foo=qwe',
            u'headers': {
                u'Cookie': zappa_cookie0
                # u'Content-Type': u'application/x-www-form-urlencoded'
            },
            u'query': {}
        }

        environ = create_wsgi_request(event, script_name='http://zappa.com/',
                                      trailing_slash=False)

        def change_cookie(environ, start_response):
            status = '200 OK'
            print 'environ', environ
            response_headers = [('Set-Cookie', 'foo=new_value')]
            start_response(status, response_headers)
            return ['Set cookies!']

        app = ZappaWSGIMiddleware(change_cookie)

        response = Response.from_app(app, environ)
        # raise Exception("stop")
        # Filter the headers for Set-Cookie header
        zappa_cookie = [x[1] for x in response.headers if x[0] == 'Set-Cookie']
        self.assertEqual(len(zappa_cookie), 1)
        zappa_cookie1 = zappa_cookie[0]
        self.assertTrue(zappa_cookie1.startswith('zappa='))

        # We have changed foo, so they should be different
        print zappa_cookie0
        print zappa_cookie1
        self.assertNotEqual(zappa_cookie0, zappa_cookie1)

        print("3: Reading the cookies")
        event['headers']['Cookie'] = zappa_cookie1

        def read_cookies(environ, start_response):
            status = '200 OK'
            print 'environ', environ
            response_headers = []
            start_response(status, response_headers)
            return [environ['HTTP_COOKIE']]

        app = ZappaWSGIMiddleware(read_cookies)

        environ = create_wsgi_request(event, script_name='http://zappa.com/',
                                      trailing_slash=False)

        response = Response.from_app(app, environ)
        cookies = dict([x.split('=') for x in response.data.split(';')])
        self.assertEqual(cookies['foo'], 'new_value')
        print cookies

        print("4: zappa cookie")

        raise Exception("the end")

    def test_wsgi_middleware_redirect(self):
        url = 'http://bogus.com/target'
        body = 'Moved. Click <a href="' + url + '">here</a>!'

        # 301
        def simple_app(environ, start_response):
            status = '301 Moved Permanently'
            response_headers = [('Location', url),
                                ('Set-Cookie', 'foo=456')]
            start_response(status, response_headers)
            return [body]

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call with empty WSGI Environment
        resp = app(dict(), self._start_response)

        self.assertEqual(self.status[0], '301 Moved Permanently')

        # Assert there is only one zappa cookie
        self.assertEqual(len(self.headers), 2)

        self.assertEqual(self.headers[0][0], 'Location')
        self.assertEqual(self.headers[0][1], url)

        self.assertEqual(self.headers[1][0], 'Set-Cookie')
        self.assertTrue(self.headers[1][1].startswith('zappa='))

        self.assertEqual(''.join(resp), body)

        # Same as above but with 302
        def simple_app(environ, start_response):
            status = '302 Found'
            response_headers = [('Location', url),
                                ('Set-Cookie', 'foo=456')]
            start_response(status, response_headers)
            return [body]

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call with empty WSGI Environment
        resp = app(dict(), self._start_response)

        self.assertEqual(self.status[0], '302 Found')

        # Assert there is only one zappa cookie
        self.assertEqual(len(self.headers), 2)

        self.assertEqual(self.headers[0][0], 'Location')
        self.assertEqual(self.headers[0][1], url)

        self.assertEqual(self.headers[1][0], 'Set-Cookie')
        self.assertTrue(self.headers[1][1].startswith('zappa='))

        self.assertEqual(''.join(resp), body)

    # TODO: This test breaks in the middleware with weird unicode chars
    def test_wsgi_middleware_unicode(self):
        # Pass some unicode through the middleware
        def simple_app(environ, start_response):
            # String of weird characters
            status = '200 OK'
            response_headers = []
            start_response(status, response_headers)
            return [unicode('asåøæd', encoding='utf8')]

        # Wrap the app with the middleware
        app = ZappaWSGIMiddleware(simple_app)

        # Call with empty WSGI Environment
        _ = app(dict(), self._start_response)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>"""Generation of human-friendly strings"""

import logging
import secrets

logger = logging.getLogger(__name__)

_adjectives = [
    "aged",
    "ancient",
    "autumn",
    "billowing",
    "bitter",
    "black",
    "blue",
    "bold",
    "broad",
    "broken",
    "calm",
    "cold",
    "cool",
    "crimson",
    "curly",
    "damp",
    "dark",
    "dawn",
    "delicate",
    "divine",
    "dry",
    "empty",
    "falling",
    "fancy",
    "flat",
    "floral",
    "fragrant",
    "frosty",
    "gentle",
    "green",
    "hidden",
    "holy",
    "icy",
    "jolly",
    "late",
    "lingering",
    "little",
    "lively",
    "long",
    "lucky",
    "misty",
    "morning",
    "muddy",
    "mute",
    "nameless",
    "noisy",
    "odd",
    "old",
    "orange",
    "patient",
    "plain",
    "polished",
    "proud",
    "purple",
    "quiet",
    "rapid",
    "raspy",
    "red",
    "restless",
    "rough",
    "round",
    "royal",
    "shiny",
    "shrill",
    "shy",
    "silent",
    "small",
    "snowy",
    "soft",
    "solitary",
    "sparkling",
    "spring",
    "square",
    "steep",
    "still",
    "summer",
    "super",
    "sweet",
    "throbbing",
    "tight",
    "tiny",
    "twilight",
    "wandering",
    "weathered",
    "white",
    "wild",
    "winter",
    "wispy",
    "withered",
    "yellow",
    "young",
]
_nouns = [
    "art",
    "band",
    "bar",
    "base",
    "bird",
    "block",
    "boat",
    "bonus",
    "bread",
    "breeze",
    "brook",
    "bush",
    "butterfly",
    "cake",
    "cell",
    "cherry",
    "cloud",
    "credit",
    "darkness",
    "dawn",
    "dew",
    "disk",
    "dream",
    "dust",
    "feather",
    "field",
    "fire",
    "firefly",
    "flower",
    "fog",
    "forest",
    "frog",
    "frost",
    "glade",
    "glitter",
    "grass",
    "hall",
    "hat",
    "haze",
    "heart",
    "hill",
    "king",
    "lab",
    "lake",
    "leaf",
    "limit",
    "math",
    "meadow",
    "mode",
    "moon",
    "morning",
    "mountain",
    "mouse",
    "mud",
    "night",
    "paper",
    "pine",
    "poetry",
    "pond",
    "queen",
    "rain",
    "recipe",
    "resonance",
    "rice",
    "river",
    "salad",
    "scene",
    "sea",
    "shadow",
    "shape",
    "silence",
    "sky",
    "smoke",
    "snow",
    "snowflake",
    "sound",
    "star",
    "sun",
    "sun",
    "sunset",
    "surf",
    "term",
    "thunder",
    "tooth",
    "tree",
    "truth",
    "union",
    "unit",
    "violet",
    "voice",
    "water",
    "waterfall",
    "wave",
    "wildflower",
    "wind",
    "wood",
]

def human_time(seconds: float):
    if seconds > 1:
        return "{} seconds".format(round(seconds, 2))
    else:
        return "{} milliseconds".format(round(seconds * 1000, 2))

def generate_human_friendly_name():
    return "{}-{}-{}".format(
        secrets.choice(_adjectives), secrets.choice(_nouns), secrets.randbelow(999) + 1
    )
<commit_msg>Human name is now generated lazily<commit_after>"""Generation of human-friendly strings"""

import logging

logger = logging.getLogger(__name__)

_adjectives = [
    "aged",
    "ancient",
    "autumn",
    "billowing",
    "bitter",
    "black",
    "blue",
    "bold",
    "broad",
    "broken",
    "calm",
    "cold",
    "cool",
    "crimson",
    "curly",
    "damp",
    "dark",
    "dawn",
    "delicate",
    "divine",
    "dry",
    "empty",
    "falling",
    "fancy",
    "flat",
    "floral",
    "fragrant",
    "frosty",
    "gentle",
    "green",
    "hidden",
    "holy",
    "icy",
    "jolly",
    "late",
    "lingering",
    "little",
    "lively",
    "long",
    "lucky",
    "misty",
    "morning",
    "muddy",
    "mute",
    "nameless",
    "noisy",
    "odd",
    "old",
    "orange",
    "patient",
    "plain",
    "polished",
    "proud",
    "purple",
    "quiet",
    "rapid",
    "raspy",
    "red",
    "restless",
    "rough",
    "round",
    "royal",
    "shiny",
    "shrill",
    "shy",
    "silent",
    "small",
    "snowy",
    "soft",
    "solitary",
    "sparkling",
    "spring",
    "square",
    "steep",
    "still",
    "summer",
    "super",
    "sweet",
    "throbbing",
    "tight",
    "tiny",
    "twilight",
    "wandering",
    "weathered",
    "white",
    "wild",
    "winter",
    "wispy",
    "withered",
    "yellow",
    "young",
]
_nouns = [
    "art",
    "band",
    "bar",
    "base",
    "bird",
    "block",
    "boat",
    "bonus",
    "bread",
    "breeze",
    "brook",
    "bush",
    "butterfly",
    "cake",
    "cell",
    "cherry",
    "cloud",
    "credit",
    "darkness",
    "dawn",
    "dew",
    "disk",
    "dream",
    "dust",
    "feather",
    "field",
    "fire",
    "firefly",
    "flower",
    "fog",
    "forest",
    "frog",
    "frost",
    "glade",
    "glitter",
    "grass",
    "hall",
    "hat",
    "haze",
    "heart",
    "hill",
    "king",
    "lab",
    "lake",
    "leaf",
    "limit",
    "math",
    "meadow",
    "mode",
    "moon",
    "morning",
    "mountain",
    "mouse",
    "mud",
    "night",
    "paper",
    "pine",
    "poetry",
    "pond",
    "queen",
    "rain",
    "recipe",
    "resonance",
    "rice",
    "river",
    "salad",
    "scene",
    "sea",
    "shadow",
    "shape",
    "silence",
    "sky",
    "smoke",
    "snow",
    "snowflake",
    "sound",
    "star",
    "sun",
    "sun",
    "sunset",
    "surf",
    "term",
    "thunder",
    "tooth",
    "tree",
    "truth",
    "union",
    "unit",
    "violet",
    "voice",
    "water",
    "waterfall",
    "wave",
    "wildflower",
    "wind",
    "wood",
]

def human_time(seconds: float):
    if seconds > 1:
        return "{} seconds".format(round(seconds, 2))
    else:
        return "{} milliseconds".format(round(seconds * 1000, 2))

class LazyHumanName:
    def __init__(self):
        self._s = None

    def __str__(self):
        if self._s is None:
            import secrets

            self._s = "{}-{}-{}".format(
                secrets.choice(_adjectives), secrets.choice(_nouns), secrets.randbelow(999) + 1
            )
        return self._s

def generate_human_friendly_name():
    return LazyHumanName()
<|endoftext|>


<commit_before># Symbolic constants for Tk

# Booleans
NO=FALSE=OFF=0
YES=TRUE=ON=1

# -anchor and -sticky
N='n'
S='s'
W='w'
E='e'
NW='nw'
SW='sw'
NE='ne'
SE='se'
NS='ns'
EW='ew'
NSEW='nsew'
CENTER='center'

# -fill
NONE='none'
X='x'
Y='y'
BOTH='both'

# -side
LEFT='left'
TOP='top'
RIGHT='right'
BOTTOM='bottom'

# -relief
RAISED='raised'
SUNKEN='sunken'
FLAT='flat'
RIDGE='ridge'
GROOVE='groove'
SOLID = 'solid'

# -orient
HORIZONTAL='horizontal'
VERTICAL='vertical'

# -tabs
NUMERIC='numeric'

# -wrap
CHAR='char'
WORD='word'

# -align
BASELINE='baseline'

# -bordermode
INSIDE='inside'
OUTSIDE='outside'

# Special tags, marks and insert positions
SEL='sel'
SEL_FIRST='sel.first'
SEL_LAST='sel.last'
END='end'
INSERT='insert'
CURRENT='current'
ANCHOR='anchor'
ALL='all' # e.g. Canvas.delete(ALL)

# Text widget and button states
NORMAL='normal'
DISABLED='disabled'
ACTIVE='active'

# Menu item types
CASCADE='cascade'
CHECKBUTTON='checkbutton'
COMMAND='command'
RADIOBUTTON='radiobutton'
SEPARATOR='separator'

# Selection modes for list boxes
SINGLE='single'
BROWSE='browse'
MULTIPLE='multiple'
EXTENDED='extended'

# Various canvas styles
PIESLICE='pieslice'
CHORD='chord'
ARC='arc'
FIRST='first'
LAST='last'
BUTT='butt'
PROJECTING='projecting'
ROUND='round'
BEVEL='bevel'
MITER='miter'

# Arguments to xview/yview
MOVETO='moveto'
SCROLL='scroll'
UNITS='units'
PAGES='pages'
<commit_msg>Add HIDDEN. Fixes #777664.<commit_after># Symbolic constants for Tk

# Booleans
NO=FALSE=OFF=0
YES=TRUE=ON=1

# -anchor and -sticky
N='n'
S='s'
W='w'
E='e'
NW='nw'
SW='sw'
NE='ne'
SE='se'
NS='ns'
EW='ew'
NSEW='nsew'
CENTER='center'

# -fill
NONE='none'
X='x'
Y='y'
BOTH='both'

# -side
LEFT='left'
TOP='top'
RIGHT='right'
BOTTOM='bottom'

# -relief
RAISED='raised'
SUNKEN='sunken'
FLAT='flat'
RIDGE='ridge'
GROOVE='groove'
SOLID = 'solid'

# -orient
HORIZONTAL='horizontal'
VERTICAL='vertical'

# -tabs
NUMERIC='numeric'

# -wrap
CHAR='char'
WORD='word'

# -align
BASELINE='baseline'

# -bordermode
INSIDE='inside'
OUTSIDE='outside'

# Special tags, marks and insert positions
SEL='sel'
SEL_FIRST='sel.first'
SEL_LAST='sel.last'
END='end'
INSERT='insert'
CURRENT='current'
ANCHOR='anchor'
ALL='all' # e.g. Canvas.delete(ALL)

# Text widget and button states
NORMAL='normal'
DISABLED='disabled'
ACTIVE='active'
# Canvas state
HIDDEN='hidden'

# Menu item types
CASCADE='cascade'
CHECKBUTTON='checkbutton'
COMMAND='command'
RADIOBUTTON='radiobutton'
SEPARATOR='separator'

# Selection modes for list boxes
SINGLE='single'
BROWSE='browse'
MULTIPLE='multiple'
EXTENDED='extended'

# Various canvas styles
PIESLICE='pieslice'
CHORD='chord'
ARC='arc'
FIRST='first'
LAST='last'
BUTT='butt'
PROJECTING='projecting'
ROUND='round'
BEVEL='bevel'
MITER='miter'

# Arguments to xview/yview
MOVETO='moveto'
SCROLL='scroll'
UNITS='units'
PAGES='pages'
<|endoftext|>


<commit_before>#! /usr/bin/python3
# -*- coding: utf-8 -*-
"""
Print a markdown-formatted bullet point with a link to the URL in your
clipboard.
"""
from glob import glob
from html.parser import HTMLParser
from http.cookiejar import CookieJar
from os import path
from subprocess import Popen, PIPE
from urllib.error import URLError
from urllib.parse import unquote, urlparse
from urllib.request import build_opener, HTTPCookieProcessor
import json
import re

from bs4 import BeautifulSoup

def main():
    """Returns the .md formatted bullet point for the URL in the clipboard."""
    url = get_url_from_clipboard()
    site = urlparse(url).netloc.replace('www.', '')
    title = get_title(url)
    title = transform_title(title, site)
    return '  * [{title} [{site}]]({url})'.format(
        title=title, site=site, url=url)

def get_url_from_clipboard():
    """Returns the contents of your clipboard, hopefully a url."""
    process = Popen(['xclip', '-out'], stdout=PIPE)
    clipboard_contents = process.communicate()[0]
    url = clipboard_contents.strip()
    url = str(url, encoding='utf-8')
    url = unquote(url)
    return url

def get_title(url):
    """Get the title of the page matching the `url`."""
    title = get_title_from_firefox_session(url)
    if not title:
        title = get_title_from_internet(url)
    return title

def transform_title(title, site):
    """Miscellaneous title transformations.

    Handle some unicode, unescape HTML, simplify hierarchical titles, ...

    """
    title = HTMLParser().unescape(title)
    title = title.strip()
    title = parse_fancy_titles(title, site)
    return title

def get_title_from_firefox_session(url):
    """Get the title for the `url` from Firefox, if it's recently opened."""
    def get_firefox_data():
        """Get a dictionary with the Firefox state from disk."""
        try:
            session = get_session_file()
            with open(session, 'r') as json_session_file:
                firefox_data = json.loads(json_session_file.read())
        except IOError:
            return ''
        except IndexError:
            return ''
        return firefox_data

    def walk(structure, criteria=lambda x: True):
        """Yield each substructure in the nested `structure`.

        Only yields things that match the `criteria` passed in.

        """
        if criteria(structure):
            yield structure
        if isinstance(structure, dict):
            for value in structure.values():
                for structure in walk(value, criteria=criteria):
                    yield structure
        elif isinstance(structure, list) or isinstance(structure, tuple):
            for item in structure:
                for structure in walk(item, criteria=criteria):
                    yield structure

    def search(structure, criteria):
        """Search the `structure` for a substruct that matches `criteria`."""
        for match in walk(structure, criteria):
            return match
        return {}

    def has_matching_url(element):
        """`element` has the url we are looking for."""
        isdict = isinstance(element, dict)
        return isdict and unquote(element.get('url', '')) == url

    firefox_data = get_firefox_data()
    matching_tab = search(firefox_data, has_matching_url)
    return matching_tab.get('title')

def get_session_file():
    """Get the path of the place where firefox stores its session info."""
    home = path.expanduser('~')
    glob_path = path.join(home, '.mozilla/firefox/*.default/sessionstore.js')
    session_file = glob(glob_path)
    return session_file[0]

def get_title_from_internet(url):
    """Get the title of the `url` from the internet."""
    try:
        opener = build_opener(HTTPCookieProcessor(CookieJar()))
        html = opener.open(url).read()
        soup = BeautifulSoup(html)
        title = soup.find('title').text
    except ValueError:
        # Usually happens when the clipboard doesn't contain a URL.
        title = ''
    except AttributeError:
        # Usually happens when soup = None; e.g. when URL is an image or PDF.
        title = ''
    except URLError:
        title = ''

    return title

def make_title_site_similarity_function(site):
    """Curry the title_site_similarity function to only require a title."""
    def remove_non_alphanumerics(word):
        """Returns the `word` with nonalphanumerics (and underscores)."""
        return re.sub(r'[^\w]', '', word)

    def title_site_similarity(title_part):
        """What portion of the words in the title part are in the site?

        Don't count very common words like "the" towards the score.
        http://www.world-english.org/english500.htm

        """
        common_words = ['the', 'of', 'to', 'and', 'a', 'in', 'is', 'it']
        result = 0.0
        words = re.split(r'[. ]', title_part)
        words = [word for word in words if word not in common_words]
        for word in words:
            word = word.lower()
            word = remove_non_alphanumerics(word)
            if word and word in site:
                result += 1.0 / len(words)

        return result

    return title_site_similarity

def parse_fancy_titles(title, site):
    """Pull out the title of the page when the title has the site in it.

    >>> parse_fancy_titles(
    ... 'Flubber - Wikipedia, the free encyclopedia', 'en.wikipedia.org')
    'Flubber'

    """
    separator = get_separator_in_title(title)
    if not separator:
        return title

    title_parts = title.split(separator)
    scores = score_parts_on_similarity(title_parts, site)
    if not any(score for score, title_part in scores):
        return title

    zero_score_title_parts = [part for score, part in scores if score == 0]
    if not zero_score_title_parts:
        return min(scores)[1]

    return separator.join(zero_score_title_parts)

def get_separator_in_title(title):
    """Return a separator in the title if there is one, otherwise None."""
    unicode_separators = [
        ' \u2010 ', ' \u002D ', ' \u00AD ', ' \u2011 ', ' \u2012 ', ' \u2013 ',
        ' \u2014 ', ' \u2015 ', ' \u2027 ', ' \u00B7 ', ' \u2043']
    ascii_separators = [' -- ', ' - ', ' | ', ' >> ', ' : ']
    separators = unicode_separators + ascii_separators
    separator = get_element_present_in_list(separators, title)
    return separator

def score_parts_on_similarity(title_parts, site):
    """Score each part of the title based on similarity to the site.

    Higher scores mean more similarity.
    Returns (score, title_part) pairs.

    """
    similarity_function = make_title_site_similarity_function(site)
    return [(similarity_function(part), part) for part in title_parts]

def get_element_present_in_list(elements, list_):
    """Return the first element that is present in `list`, otherwise None."""
    for element in elements:
        if element in list_:
            return element
    return None

if __name__ == '__main__':
    print(main())
<commit_msg>Set useragent=Mozilla/5.0 so kaspersky.com stops blocking<commit_after>#! /usr/bin/python3
# -*- coding: utf-8 -*-
"""
Print a markdown-formatted bullet point with a link to the URL in your
clipboard.
"""
from glob import glob
from html.parser import HTMLParser
from http.cookiejar import CookieJar
from os import path
from subprocess import Popen, PIPE
from urllib.error import URLError
from urllib.parse import unquote, urlparse
from urllib.request import build_opener, HTTPCookieProcessor
import json
import re

from bs4 import BeautifulSoup

def main():
    """Returns the .md formatted bullet point for the URL in the clipboard."""
    url = get_url_from_clipboard()
    site = urlparse(url).netloc.replace('www.', '')
    title = get_title(url)
    title = transform_title(title, site)
    return '  * [{title} [{site}]]({url})'.format(
        title=title, site=site, url=url)

def get_url_from_clipboard():
    """Returns the contents of your clipboard, hopefully a url."""
    process = Popen(['xclip', '-out'], stdout=PIPE)
    clipboard_contents = process.communicate()[0]
    url = clipboard_contents.strip()
    url = str(url, encoding='utf-8')
    url = unquote(url)
    return url

def get_title(url):
    """Get the title of the page matching the `url`."""
    title = get_title_from_firefox_session(url)
    if not title:
        title = get_title_from_internet(url)
    return title

def transform_title(title, site):
    """Miscellaneous title transformations.

    Handle some unicode, unescape HTML, simplify hierarchical titles, ...

    """
    title = HTMLParser().unescape(title)
    title = title.strip()
    title = parse_fancy_titles(title, site)
    return title

def get_title_from_firefox_session(url):
    """Get the title for the `url` from Firefox, if it's recently opened."""
    def get_firefox_data():
        """Get a dictionary with the Firefox state from disk."""
        try:
            session = get_session_file()
            with open(session, 'r') as json_session_file:
                firefox_data = json.loads(json_session_file.read())
        except IOError:
            return ''
        except IndexError:
            return ''
        return firefox_data

    def walk(structure, criteria=lambda x: True):
        """Yield each substructure in the nested `structure`.

        Only yields things that match the `criteria` passed in.

        """
        if criteria(structure):
            yield structure
        if isinstance(structure, dict):
            for value in structure.values():
                for structure in walk(value, criteria=criteria):
                    yield structure
        elif isinstance(structure, list) or isinstance(structure, tuple):
            for item in structure:
                for structure in walk(item, criteria=criteria):
                    yield structure

    def search(structure, criteria):
        """Search the `structure` for a substruct that matches `criteria`."""
        for match in walk(structure, criteria):
            return match
        return {}

    def has_matching_url(element):
        """`element` has the url we are looking for."""
        isdict = isinstance(element, dict)
        return isdict and unquote(element.get('url', '')) == url

    firefox_data = get_firefox_data()
    matching_tab = search(firefox_data, has_matching_url)
    return matching_tab.get('title')

def get_session_file():
    """Get the path of the place where firefox stores its session info."""
    home = path.expanduser('~')
    glob_path = path.join(home, '.mozilla/firefox/*.default/sessionstore.js')
    session_file = glob(glob_path)
    return session_file[0]

def get_title_from_internet(url):
    """Get the title of the `url` from the internet."""
    try:
        opener = build_opener(HTTPCookieProcessor(CookieJar()))
        opener.addheaders = [('User-agent', 'Mozilla/5.0')]
        html = opener.open(url).read()
        soup = BeautifulSoup(html)
        title = soup.find('title').text
    except ValueError:
        # Usually happens when the clipboard doesn't contain a URL.
        title = ''
    except AttributeError:
        # Usually happens when soup = None; e.g. when URL is an image or PDF.
        title = ''
    except URLError:
        title = ''

    return title

def make_title_site_similarity_function(site):
    """Curry the title_site_similarity function to only require a title."""
    def remove_non_alphanumerics(word):
        """Returns the `word` with nonalphanumerics (and underscores)."""
        return re.sub(r'[^\w]', '', word)

    def title_site_similarity(title_part):
        """What portion of the words in the title part are in the site?

        Don't count very common words like "the" towards the score.
        http://www.world-english.org/english500.htm

        """
        common_words = ['the', 'of', 'to', 'and', 'a', 'in', 'is', 'it']
        result = 0.0
        words = re.split(r'[. ]', title_part)
        words = [word for word in words if word not in common_words]
        for word in words:
            word = word.lower()
            word = remove_non_alphanumerics(word)
            if word and word in site:
                result += 1.0 / len(words)

        return result

    return title_site_similarity

def parse_fancy_titles(title, site):
    """Pull out the title of the page when the title has the site in it.

    >>> parse_fancy_titles(
    ... 'Flubber - Wikipedia, the free encyclopedia', 'en.wikipedia.org')
    'Flubber'

    """
    separator = get_separator_in_title(title)
    if not separator:
        return title

    title_parts = title.split(separator)
    scores = score_parts_on_similarity(title_parts, site)
    if not any(score for score, title_part in scores):
        return title

    zero_score_title_parts = [part for score, part in scores if score == 0]
    if not zero_score_title_parts:
        return min(scores)[1]

    return separator.join(zero_score_title_parts)

def get_separator_in_title(title):
    """Return a separator in the title if there is one, otherwise None."""
    unicode_separators = [
        ' \u2010 ', ' \u002D ', ' \u00AD ', ' \u2011 ', ' \u2012 ', ' \u2013 ',
        ' \u2014 ', ' \u2015 ', ' \u2027 ', ' \u00B7 ', ' \u2043']
    ascii_separators = [' -- ', ' - ', ' | ', ' >> ', ' : ']
    separators = unicode_separators + ascii_separators
    separator = get_element_present_in_list(separators, title)
    return separator

def score_parts_on_similarity(title_parts, site):
    """Score each part of the title based on similarity to the site.

    Higher scores mean more similarity.
    Returns (score, title_part) pairs.

    """
    similarity_function = make_title_site_similarity_function(site)
    return [(similarity_function(part), part) for part in title_parts]

def get_element_present_in_list(elements, list_):
    """Return the first element that is present in `list`, otherwise None."""
    for element in elements:
        if element in list_:
            return element
    return None

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


<commit_before>from pandaharvester.harvestercore.work_spec import WorkSpec
from pandaharvester.harvestercore.plugin_base import PluginBase
from pandaharvester.harvestercore import core_utils
from pandaharvester.harvestercloud.googlecloud import compute, ZONE, PROJECT

baseLogger = core_utils.setup_logger('google_monitor')

class GoogleMonitor(PluginBase):
    def __init__(self, **kwarg):
        PluginBase.__init__(self, **kwarg)

        # States taken from: https://cloud.google.com/compute/docs/instances/checking-instance-status
        self.vm_to_worker_status = {
                                     'RUNNING': WorkSpec.ST_running,
                                     'TERMINATED': WorkSpec.ST_finished,
                                     'STOPPING': WorkSpec.ST_finished,
                                     'PROVISIONING': WorkSpec.ST_submitted,
                                     'STAGING': WorkSpec.ST_submitted
                                     }

    def list_vms(self):
        """
        List the status of the running VMs
        :return:
        """

        try:
            result = compute.instances().list(project=PROJECT, zone=ZONE).execute()
            vm_instances = result['items']

            # make a list with the VM names
            vm_names = map(lambda vm_instance: vm_instance['name'], vm_instances)

            # make a dictionary so we can retrieve a VM by its name
            vm_name_to_status = {}
            for vm_instance in vm_instances:
                vm_name_to_status[vm_instance['name']] = vm_instance['status']

            return vm_names, vm_name_to_status

        except:
            return None, None

    def check_workers(self, workers):
        """
        This method takes a list of WorkSpecs as input argument and returns a list of worker's statuses.
        Nth element in the return list corresponds to the status of Nth WorkSpec in the given list.

        :param worker_list: a list of work specs instances
        :return: A tuple containing the return code (True for success, False otherwise) and a list of worker's statuses
        :rtype: (bool, [string,])
        """

        # running instances
        vm_names, vm_name_to_status = self.list_vms()
        if vm_names is None and vm_name_to_status is None:
            error_string = 'Could not list the VMs'
            baseLogger.error(error_string)
            return False, error_string

        # extract the list of batch IDs
        batch_IDs = map(lambda x: str(x.batchID), workers)
        baseLogger.debug('Batch IDs: {0}'.format(batch_IDs))

        ret_list = []
        for batch_ID in batch_IDs:
            tmp_log = core_utils.make_logger(baseLogger, 'batch ID={0}'.format(batch_ID), method_name='check_workers')

            if batch_ID not in vm_names:
                new_status = WorkSpec.ST_finished
                message = 'VM not found'
            else:
                try:
                    new_status = self.vm_to_worker_status[vm_name_to_status[batch_ID]]
                    message = 'VM status returned by GCE API'
                except KeyError:
                    new_status = WorkSpec.ST_missed
                    message = 'Unknown status to Harvester: {0}'.format(vm_name_to_status[batch_ID])

            tmp_log.debug('new_status={0}'.format(new_status))
            ret_list.append((new_status, message))

        baseLogger.debug('ret_list: {0}'.format(ret_list))
        return True, ret_list
<commit_msg>GCE: distinguish no running VMs and an exception listing them...<commit_after>from pandaharvester.harvestercore.work_spec import WorkSpec
from pandaharvester.harvestercore.plugin_base import PluginBase
from pandaharvester.harvestercore import core_utils
from pandaharvester.harvestercloud.googlecloud import compute, ZONE, PROJECT

baseLogger = core_utils.setup_logger('google_monitor')

class GoogleMonitor(PluginBase):
    def __init__(self, **kwarg):
        PluginBase.__init__(self, **kwarg)

        # States taken from: https://cloud.google.com/compute/docs/instances/checking-instance-status
        self.vm_to_worker_status = {
                                     'RUNNING': WorkSpec.ST_running,
                                     'TERMINATED': WorkSpec.ST_finished,
                                     'STOPPING': WorkSpec.ST_finished,
                                     'PROVISIONING': WorkSpec.ST_submitted,
                                     'STAGING': WorkSpec.ST_submitted
                                     }

    def list_vms(self):
        """
        List the status of the running VMs
        :return:
        """

        try:
            result = compute.instances().list(project=PROJECT, zone=ZONE).execute()

            try:
                vm_instances = result['items']
            except KeyError:
                # there are no VMs running
                return [], {}

            # make a list with the VM names
            vm_names = map(lambda vm_instance: vm_instance['name'], vm_instances)

            # make a dictionary so we can retrieve a VM by its name
            vm_name_to_status = {}
            for vm_instance in vm_instances:
                vm_name_to_status[vm_instance['name']] = vm_instance['status']

            return vm_names, vm_name_to_status

        except:
            return None, None

    def check_workers(self, workers):
        """
        This method takes a list of WorkSpecs as input argument and returns a list of worker's statuses.
        Nth element in the return list corresponds to the status of Nth WorkSpec in the given list.

        :param worker_list: a list of work specs instances
        :return: A tuple containing the return code (True for success, False otherwise) and a list of worker's statuses
        :rtype: (bool, [string,])
        """

        # running instances
        vm_names, vm_name_to_status = self.list_vms()
        if vm_names is None and vm_name_to_status is None:
            error_string = 'Could not list the VMs'
            baseLogger.error(error_string)
            return False, error_string

        # extract the list of batch IDs
        batch_IDs = map(lambda x: str(x.batchID), workers)
        baseLogger.debug('Batch IDs: {0}'.format(batch_IDs))

        ret_list = []
        for batch_ID in batch_IDs:
            tmp_log = core_utils.make_logger(baseLogger, 'batch ID={0}'.format(batch_ID), method_name='check_workers')

            if batch_ID not in vm_names:
                new_status = WorkSpec.ST_finished
                message = 'VM not found'
            else:
                try:
                    new_status = self.vm_to_worker_status[vm_name_to_status[batch_ID]]
                    message = 'VM status returned by GCE API'
                except KeyError:
                    new_status = WorkSpec.ST_missed
                    message = 'Unknown status to Harvester: {0}'.format(vm_name_to_status[batch_ID])

            tmp_log.debug('new_status={0}'.format(new_status))
            ret_list.append((new_status, message))

        baseLogger.debug('ret_list: {0}'.format(ret_list))
        return True, ret_list
<|endoftext|>


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

"""
Mock code to support testing of chrome_throttle.
"""

import unittest
import sys
import os
from signal import SIGSTOP, SIGCONT
import io
import re

from ..chrome_throttle import get_chromium_renderers

"""
 2 /usr/lib/chromium/chromium --password-store=detect
 2 /usr/lib/chromium/chromium --type=zygote
 1 /usr/lib/chromium/chromium-sandbox /usr/lib/chromium/chromium --type=zygote
15 /usr/lib/chromium/chromium --type=renderer --lang=en-US
"""

CMD_CR_PASSWD   = "/usr/lib/chromium/chromium --password-store=detect"
CMD_CR_ZYGOTE   = "/usr/lib/chromium/chromium --type=zygote"
CMD_CR_SANDBOX  = "/usr/lib/chromium/chromium-sandbox /usr/lib/chromium/chromium --type=zygote"
CMD_CR_RENDERER = "/usr/lib/chromium/chromium --type=renderer --lang=en-US"
CMD_BASH        = "/bin/bash"
CMD_STRACE      = "/usr/bin/strace /usr/lib/chromium/chromium"

class MockProcess(object):
    def __init__(self, pid, cmdline, status='S'):
        self.pid = pid
        self.cmdline = cmdline
        self.cmd = cmdline.split()[0].split('/')[-1]
        self.status = status

    def pause(self):
        self.status = 'T'

    def unpause(self):
        self.sleep()

    def run(self):
        self.status = 'R'

    def sleep(self):
        self.status = 'S'

    def status_file(self):
        status_str = {'S': 'sleeping', 'T': 'stopped', 'R': 'running'}[self.status]
        return "Name:\t{0}\nState:\t{1} ({2})\n".format(
                self.cmd, self.status, status_str)

class StringFile(io.BytesIO):
    pass

class MockProcListException(Exception):
    "Used to ensure only a single instance is active."

class MockProcList(object):

    @staticmethod
    def do_system_mocks():
        try:
            __builtins__['real_open']
        except KeyError:
            __builtins__['real_open'] = __builtins__['open']
            __builtins__['open'] = MockProcList.open

        try:
            os.real_listdir
        except AttributeError:
            os.real_listdir = os.listdir
            os.listdir = MockProcList.listdir

        try:
            os.real_kill
        except AttributeError:
            os.real_kill = os.kill
            os.kill = MockProcList.kill

    @staticmethod
    def listdir(dir_name):
        procs = MockProcList.procs
        if dir_name == "/proc":
            return sorted([str(i) for i in procs.keys()] + ["acpi", "self", "sys"])
        return os.real_listdir(dir_name)

    @staticmethod
    def kill(pid, sig):
        procs = MockProcList.procs
        try:
            if sig == SIGSTOP:
                procs[pid].pause()
            elif sig == SIGCONT:
                procs[pid].unpause()
        except KeyError:
            raise OSError("[Errno 1] Operation not permitted")

    @staticmethod
    def open(fname, mode='r'):   # FIXME: there is a buffered argument
        procs = MockProcList.procs
        try:
            try:
                pid = int(re.compile("/proc/(\d+)/status").match(fname).groups(0)[0])
                return StringFile(procs[pid].status_file())
            except AttributeError:
                pass
            try:
                pid = int(re.compile("/proc/(\d+)/cmdline").match(fname).groups(0)[0])
                return StringFile(procs[pid].cmdline)
            except AttributeError:
                pass

        except KeyError:
            raise IOError("IOError: [Errno 2] No such file or directory: '{}'".format(fname))

        return real_open(fname, mode)

    procs = None

    def __init__(self, *args):
        if MockProcList.procs is not None:
            raise MockProcListException("only one instance allowed at a time")
        MockProcList.do_system_mocks()
        self.procs = {t[0]: apply(MockProcess, t) for t in args}
        MockProcList.procs = self.procs

    def reset(self):
        for k in list(self.procs):
            del self.procs[k]
        self.procs = None
        MockProcList.procs = self.procs

class ChromeThrottleTest(unittest.TestCase):

    def setUp(self):
        self.mpl = MockProcList(
           (1000, CMD_BASH, 'S'),
           (2000, CMD_STRACE, 'R'),
           (3000, CMD_CR_PASSWD, 'S'),
           (3001, CMD_CR_PASSWD, 'S'),
           (3002, CMD_CR_ZYGOTE, 'S'),
           (3003, CMD_CR_ZYGOTE, 'S'),
           (3004, CMD_CR_SANDBOX, 'S'),
           (4000, CMD_CR_RENDERER, 'S'),
           (4001, CMD_CR_RENDERER, 'S'),
           (4002, CMD_CR_RENDERER, 'S'),
           (4003, CMD_CR_RENDERER, 'S'),
           (4004, CMD_CR_RENDERER, 'R'),
           (4005, CMD_CR_RENDERER, 'S'))

    def tearDown(self):
        self.mpl.reset()

    def test_placeholder(self):
        self.assertEqual(1, 1)

    def test_procs_baseline(self):
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'S'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_non_renderer(self):
        os.kill(3000, SIGSTOP)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'S'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_two_renderers(self):
        os.kill(4001, SIGSTOP)
        os.kill(4002, SIGSTOP)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'T'), (4002, 'T'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_two_renderers_and_continue_one(self):
        os.kill(4001, SIGSTOP)
        os.kill(4002, SIGSTOP)
        os.kill(4001, SIGCONT)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'T'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_non_existent_process(self):
        self.assertRaises(OSError, os.kill, 5000, SIGSTOP)

if __name__ == '__main__':
    unittest.main()

<commit_msg>The use of apply() is deprecated with Python 3.<commit_after>#!/usr/bin/env python

"""
Mock code to support testing of chrome_throttle.
"""

import unittest
import sys
import os
from signal import SIGSTOP, SIGCONT
import io
import re

from ..chrome_throttle import get_chromium_renderers

"""
 2 /usr/lib/chromium/chromium --password-store=detect
 2 /usr/lib/chromium/chromium --type=zygote
 1 /usr/lib/chromium/chromium-sandbox /usr/lib/chromium/chromium --type=zygote
15 /usr/lib/chromium/chromium --type=renderer --lang=en-US
"""

CMD_CR_PASSWD   = "/usr/lib/chromium/chromium --password-store=detect"
CMD_CR_ZYGOTE   = "/usr/lib/chromium/chromium --type=zygote"
CMD_CR_SANDBOX  = "/usr/lib/chromium/chromium-sandbox /usr/lib/chromium/chromium --type=zygote"
CMD_CR_RENDERER = "/usr/lib/chromium/chromium --type=renderer --lang=en-US"
CMD_BASH        = "/bin/bash"
CMD_STRACE      = "/usr/bin/strace /usr/lib/chromium/chromium"

class MockProcess(object):
    def __init__(self, pid, cmdline, status='S'):
        self.pid = pid
        self.cmdline = cmdline
        self.cmd = cmdline.split()[0].split('/')[-1]
        self.status = status

    def pause(self):
        self.status = 'T'

    def unpause(self):
        self.sleep()

    def run(self):
        self.status = 'R'

    def sleep(self):
        self.status = 'S'

    def status_file(self):
        status_str = {'S': 'sleeping', 'T': 'stopped', 'R': 'running'}[self.status]
        return "Name:\t{0}\nState:\t{1} ({2})\n".format(
                self.cmd, self.status, status_str)

class StringFile(io.BytesIO):
    pass

class MockProcListException(Exception):
    "Used to ensure only a single instance is active."

class MockProcList(object):

    @staticmethod
    def do_system_mocks():
        try:
            __builtins__['real_open']
        except KeyError:
            __builtins__['real_open'] = __builtins__['open']
            __builtins__['open'] = MockProcList.open

        try:
            os.real_listdir
        except AttributeError:
            os.real_listdir = os.listdir
            os.listdir = MockProcList.listdir

        try:
            os.real_kill
        except AttributeError:
            os.real_kill = os.kill
            os.kill = MockProcList.kill

    @staticmethod
    def listdir(dir_name):
        procs = MockProcList.procs
        if dir_name == "/proc":
            return sorted([str(i) for i in procs.keys()] + ["acpi", "self", "sys"])
        return os.real_listdir(dir_name)

    @staticmethod
    def kill(pid, sig):
        procs = MockProcList.procs
        try:
            if sig == SIGSTOP:
                procs[pid].pause()
            elif sig == SIGCONT:
                procs[pid].unpause()
        except KeyError:
            raise OSError("[Errno 1] Operation not permitted")

    @staticmethod
    def open(fname, mode='r'):   # FIXME: there is a buffered argument
        procs = MockProcList.procs
        try:
            try:
                pid = int(re.compile("/proc/(\d+)/status").match(fname).groups(0)[0])
                return StringFile(procs[pid].status_file())
            except AttributeError:
                pass
            try:
                pid = int(re.compile("/proc/(\d+)/cmdline").match(fname).groups(0)[0])
                return StringFile(procs[pid].cmdline)
            except AttributeError:
                pass

        except KeyError:
            raise IOError("IOError: [Errno 2] No such file or directory: '{}'".format(fname))

        return real_open(fname, mode)

    procs = None

    def __init__(self, *args):
        if MockProcList.procs is not None:
            raise MockProcListException("only one instance allowed at a time")
        MockProcList.do_system_mocks()
        self.procs = {t[0]: MockProcess(*t) for t in args}
        MockProcList.procs = self.procs

    def reset(self):
        for k in list(self.procs):
            del self.procs[k]
        self.procs = None
        MockProcList.procs = self.procs

class ChromeThrottleTest(unittest.TestCase):

    def setUp(self):
        self.mpl = MockProcList(
           (1000, CMD_BASH, 'S'),
           (2000, CMD_STRACE, 'R'),
           (3000, CMD_CR_PASSWD, 'S'),
           (3001, CMD_CR_PASSWD, 'S'),
           (3002, CMD_CR_ZYGOTE, 'S'),
           (3003, CMD_CR_ZYGOTE, 'S'),
           (3004, CMD_CR_SANDBOX, 'S'),
           (4000, CMD_CR_RENDERER, 'S'),
           (4001, CMD_CR_RENDERER, 'S'),
           (4002, CMD_CR_RENDERER, 'S'),
           (4003, CMD_CR_RENDERER, 'S'),
           (4004, CMD_CR_RENDERER, 'R'),
           (4005, CMD_CR_RENDERER, 'S'))

    def tearDown(self):
        self.mpl.reset()

    def test_placeholder(self):
        self.assertEqual(1, 1)

    def test_procs_baseline(self):
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'S'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_non_renderer(self):
        os.kill(3000, SIGSTOP)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'S'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_two_renderers(self):
        os.kill(4001, SIGSTOP)
        os.kill(4002, SIGSTOP)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'T'), (4002, 'T'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_two_renderers_and_continue_one(self):
        os.kill(4001, SIGSTOP)
        os.kill(4002, SIGSTOP)
        os.kill(4001, SIGCONT)
        self.assertEqual(get_chromium_renderers(),
             [(4000, 'S'), (4001, 'S'), (4002, 'T'),
              (4003, 'S'), (4004, 'R'), (4005, 'S')])

    def test_procs_stop_non_existent_process(self):
        self.assertRaises(OSError, os.kill, 5000, SIGSTOP)

if __name__ == '__main__':
    unittest.main()

<|endoftext|>


<commit_before>"""This script is intended to be run on an Amazon ECS container, so information
for the job either needs to be provided in environment variables (e.g., the
REACH version and path) or loaded from S3 (e.g., the list of PMIDs).
"""
from __future__ import absolute_import, print_function, unicode_literals
from builtins import dict, str
import boto3
import botocore
import logging
import sys
import os
import random
import numpy as np
from matplotlib import pyplot as plt
from datetime import datetime
from indra.tools.reading.db_reading.read_db import produce_readings, \
    produce_statements, get_id_dict
from indra.tools.reading.readers import get_readers

def plot_hist(agged, agg_over, data, s3, s3_base, bucket_name):
    fig = plt.figure()
    plt.hist(data)
    plt.xlabel('Number of %s for %s' % (agged, agg_over))
    plt.ylabel('Number of %s with a given number of %s' % (agg_over, agged))
    fname = '%s_per_%s.png' % (agged, agg_over)
    fig.savefig(fname)
    with open(fname, 'rb') as f:
        s3_key = s3_base + fname
        s3.put_object(Key=s3_key, Body=f.read(), Bucket=bucket_name)
    return

def report_statistics(reading_outputs, stmt_outputs, starts, ends, basename, s3,
                      bucket_name):
    s3_base = 'reading_results/%s/logs/run_db_reading/statisics/' % basename
    starts['stats'] = datetime.now()
    data_dict = {}
    hist_dict = {}
    data_dict['Total readings'] = len(reading_outputs)
    reading_stmts = [(rd.reading_id, rd.tcid, rd.reader, rd.get_statements())
                     for rd in reading_outputs]
    data_dict['Content processed'] = len({t[1] for t in reading_stmts})
    data_dict['Statements produced'] = len(stmt_outputs)

    # Do a bunch of aggregation
    tc_rd_dict = {}
    tc_stmt_dict = {}
    rd_stmt_dict = {}
    reader_stmts = {}
    reader_tcids = {}
    reader_rids = {}
    for rid, tcid, reader, stmts in reading_stmts:
        # Handle things keyed by tcid
        if tcid not in tc_rd_dict.keys():
            tc_rd_dict[tcid] = {rid}
            tc_stmt_dict[tcid] = set(stmts)
        else:
            tc_rd_dict[tcid].add(rid)
            tc_stmt_dict[tcid] |= set(stmts)

        # Handle things keyed by rid
        if rid not in rd_stmt_dict.keys():
            rd_stmt_dict[rid] = set(stmts)
        else:
            rd_stmt_dict[rid] |= set(stmts)  # this shouldn't really happen.

        # Handle things keyed by reader
        if reader not in reader_stmts.keys():
            reader_stmts[reader] = set(stmts)
            reader_tcids[reader] = {tcid}
            reader_rids[reader] = {rid}
        else:
            reader_stmts[reader] |= set(stmts)
            reader_tcids[reader].add(tcid)
            reader_rids[reader].add(rid)

    # Produce some numpy count arrays.
    hist_dict[('readings', 'text content')] = \
        np.array([len(rid_set) for rid_set in tc_rd_dict.values()])
    hist_dict[('statements', 'text_content')] = \
        np.array([len(stmts) for stmts in tc_stmt_dict.values()])
    hist_dict[('statements', 'readings')] = \
        np.array([len(stmts) for stmts in rd_stmt_dict.values()])
    hist_dict[('statements', 'readers')] = \
        np.array([len(stmts) for stmts in reader_stmts.values()])
    hist_dict[('text content', 'reader')] = \
        np.array([len(tcid_set) for tcid_set in reader_tcids.values()])
    hist_dict[('readings', 'reader')] = \
        np.array([len(rid_set) for rid_set in reader_rids.values()])

    # Produce the histograms
    for (agged, agg_over), data in hist_dict.items():
        plot_hist(agged, agg_over, data, s3, s3_base, bucket_name)
        label = '%s per %s' % (agged, agg_over)
        data_dict[label.capitalize()] = {'mean': data.mean(), 'std': data.std(),
                                         'median': np.median(data)}

    # Produce the summary report
    text_report_str = ''
    top_labels = ['Total readings', 'Content processed', 'Statements produced']
    for label in top_labels:
        text_report_str += '%s: %d\n' % (label, data_dict[label])

    for label, data in data_dict.items():
        if label in top_labels:
            continue
        if isinstance(data, dict):
            text_report_str += '%s:\n' % label
            text_report_str += '\n'.join(['\t%s: %d' % (k, v)
                                          for k, v in data.items()])
            text_report_str += '\n'
        else:
            text_report_str += '%s: %d\n' % (label, data)

    s3.put_object(Key=s3_base + 'summary.txt', Body=text_report_str,
                  Bucket=bucket_name)
    ends['stats'] = datetime.now()

    # Report on the timing
    timing_str = ''
    for step in ['reading', 'statement production', 'stats']:
        time_taken = ends[step] - starts[step]
        timing_str += ('%22s: start: %s, end: %s, duration: %s\n'
                       % (step, starts[step], ends[step], time_taken))

    s3.put_object(Key=s3_base + 'time.txt', Body=timing_str, Bucket=bucket_name)
    return

if __name__ == '__main__':
    from argparse import ArgumentParser
    parser = ArgumentParser(
        description=__doc__
        )
    parser.add_argument(
        dest='basename',
        help='The name of this run.'
        )
    parser.add_argument(
        dest='out_dir',
        help='The name of the temporary output directory'
        )
    parser.add_argument(
        dest='read_mode',
        choices=['all', 'unread', 'none'],
        help=("Set the reading mode. If 'all', read everything, if "
              "'unread', only read content that does not have pre-existing "
              "readings of the same reader and version, if 'none', only "
              "use pre-existing readings. Default is 'unread'.")
        )
    parser.add_argument(
        dest='stmt_mode',
        choices=['all', 'unread', 'none'],
        help=("Choose which readings should produce statements. If 'all', all "
              "readings that are produced or retrieved will be used to produce "
              "statements. If 'unread', only produce statements from "
              "previously unread content. If 'none', do not produce any "
              "statements (only readings will be produced).")
        )
    parser.add_argument(
        dest='num_cores',
        help='Select the number of cores on which to run.',
        type=int
        )
    parser.add_argument(
        dest='start_index',
        help='Select the index of the first pmid in the list to read.',
        type=int
        )
    parser.add_argument(
        dest='end_index',
        help='Select the index of the last pmid in the list to read.',
        type=int
        )
    parser.add_argument(
        '-r', '--readers',
        dest='readers',
        choices=['reach', 'sparser'],
        nargs='+',
        help='Choose which reader(s) to use.'
        )
    parser.add_argument(
        '--force_fulltext',
        action='store_true',
        help='Require that content be fulltext, skip anything that isn\'t.'
        )
    parser.add_argument(
        '--use_best_fulltext',
        action='store_true',
        help='Read only the "best" fulltext available for a given id.'
        )
    parser.add_argument(
        '--test',
        action='store_true',
        help="Use the test database."
        )
    args = parser.parse_args()

    logger = logging.getLogger('read_db_aws')
    logger.setLevel(logging.DEBUG)

    client = boto3.client('s3')
    bucket_name = 'bigmech'
    id_list_key = 'reading_inputs/%s/id_list' % args.basename
    readers = [reader_class(args.basename, args.num_cores)
               for reader_class in get_readers()
               if reader_class.name.lower() in args.readers]

    try:
        id_list_obj = client.get_object(
            Bucket=bucket_name,
            Key=id_list_key
            )
    # Handle a missing object gracefully
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == 'NoSuchKey':
            logger.info('Could not find PMID list file at %s, exiting' %
                        id_list_key)
            sys.exit(1)
        # If there was some other kind of problem, re-raise the exception
        else:
            raise e
    # Get the content from the object
    id_list_str = id_list_obj['Body'].read().decode('utf8').strip()
    id_str_list = id_list_str.splitlines()[args.start_index:args.end_index]
    random.shuffle(id_str_list)
    id_dict = get_id_dict([line.strip() for line in id_str_list])

    # Some combinations of options don't make sense:
    forbidden_combos = [('all', 'unread'), ('none', 'unread'), ('none', 'none')]
    assert (args.read_mode, args.stmt_mode) not in forbidden_combos, \
        ("The combination of reading mode %s and statement mode %s is not "
         "allowed." % (args.reading_mode, args.stmt_mode))

    # Init some timing dicts
    starts = {}
    ends = {}

    # Get a handle for the database
    if args.test:
        from indra.db import util as dbu
        db = dbu.get_test_db()
    else:
        db = None

    # Read everything ========================================
    starts['reading'] = datetime.now()
    outputs = produce_readings(id_dict, readers, verbose=True,
                               read_mode=args.read_mode,
                               get_preexisting=(args.stmt_mode == 'all'),
                               force_fulltext=args.force_fulltext,
                               prioritize=args.use_best_fulltext, db=db)
    ends['reading'] = datetime.now()

    # Preserve the sparser logs
    contents = os.listdir('.')
    sparser_logs = [fname for fname in contents
                    if fname.startswith('sparser') and fname.endswith('log')]
    sparser_log_dir = ('reading_results/%s/logs/run_db_reading_queue/'
                       'sparser_logs_%s/') % (
                           args.basename,
                           datetime.now().strftime('%Y%m%d_%H%M%S')
                           )
    for fname in sparser_logs:
        s3_key = sparser_log_dir + fname
        logger.info("Saving sparser logs to %s on s3 in %s."
                    % (s3_key, bucket_name))
        with open(fname, 'r') as f:
            client.put_object(Key=s3_key, Body=f.read(),
                              Bucket=bucket_name)

    # Convert the outputs to statements ==================================
    if args.stmt_mode != 'none':
        starts['statement production'] = datetime.now()
        stmt_data = produce_statements(outputs, n_proc=args.num_cores, db=db)
        ends['statement production'] = datetime.now()
    else:
        stmt_data = []

    report_statistics(outputs, stmt_data, starts, ends, args.basename, client,
                      bucket_name)

<commit_msg>Use display-less matplotlib backend.<commit_after>"""This script is intended to be run on an Amazon ECS container, so information
for the job either needs to be provided in environment variables (e.g., the
REACH version and path) or loaded from S3 (e.g., the list of PMIDs).
"""
from __future__ import absolute_import, print_function, unicode_literals
from builtins import dict, str
import boto3
import botocore
import logging
import sys
import os
import random
import numpy as np
import matplotlib as mpl
mpl.use('Agg')
from matplotlib import pyplot as plt
from datetime import datetime
from indra.tools.reading.db_reading.read_db import produce_readings, \
    produce_statements, get_id_dict
from indra.tools.reading.readers import get_readers

def plot_hist(agged, agg_over, data, s3, s3_base, bucket_name):
    fig = plt.figure()
    plt.hist(data)
    plt.xlabel('Number of %s for %s' % (agged, agg_over))
    plt.ylabel('Number of %s with a given number of %s' % (agg_over, agged))
    fname = '%s_per_%s.png' % (agged, agg_over)
    fig.savefig(fname)
    with open(fname, 'rb') as f:
        s3_key = s3_base + fname
        s3.put_object(Key=s3_key, Body=f.read(), Bucket=bucket_name)
    return

def report_statistics(reading_outputs, stmt_outputs, starts, ends, basename, s3,
                      bucket_name):
    s3_base = 'reading_results/%s/logs/run_db_reading/statisics/' % basename
    starts['stats'] = datetime.now()
    data_dict = {}
    hist_dict = {}
    data_dict['Total readings'] = len(reading_outputs)
    reading_stmts = [(rd.reading_id, rd.tcid, rd.reader, rd.get_statements())
                     for rd in reading_outputs]
    data_dict['Content processed'] = len({t[1] for t in reading_stmts})
    data_dict['Statements produced'] = len(stmt_outputs)

    # Do a bunch of aggregation
    tc_rd_dict = {}
    tc_stmt_dict = {}
    rd_stmt_dict = {}
    reader_stmts = {}
    reader_tcids = {}
    reader_rids = {}
    for rid, tcid, reader, stmts in reading_stmts:
        # Handle things keyed by tcid
        if tcid not in tc_rd_dict.keys():
            tc_rd_dict[tcid] = {rid}
            tc_stmt_dict[tcid] = set(stmts)
        else:
            tc_rd_dict[tcid].add(rid)
            tc_stmt_dict[tcid] |= set(stmts)

        # Handle things keyed by rid
        if rid not in rd_stmt_dict.keys():
            rd_stmt_dict[rid] = set(stmts)
        else:
            rd_stmt_dict[rid] |= set(stmts)  # this shouldn't really happen.

        # Handle things keyed by reader
        if reader not in reader_stmts.keys():
            reader_stmts[reader] = set(stmts)
            reader_tcids[reader] = {tcid}
            reader_rids[reader] = {rid}
        else:
            reader_stmts[reader] |= set(stmts)
            reader_tcids[reader].add(tcid)
            reader_rids[reader].add(rid)

    # Produce some numpy count arrays.
    hist_dict[('readings', 'text content')] = \
        np.array([len(rid_set) for rid_set in tc_rd_dict.values()])
    hist_dict[('statements', 'text_content')] = \
        np.array([len(stmts) for stmts in tc_stmt_dict.values()])
    hist_dict[('statements', 'readings')] = \
        np.array([len(stmts) for stmts in rd_stmt_dict.values()])
    hist_dict[('statements', 'readers')] = \
        np.array([len(stmts) for stmts in reader_stmts.values()])
    hist_dict[('text content', 'reader')] = \
        np.array([len(tcid_set) for tcid_set in reader_tcids.values()])
    hist_dict[('readings', 'reader')] = \
        np.array([len(rid_set) for rid_set in reader_rids.values()])

    # Produce the histograms
    for (agged, agg_over), data in hist_dict.items():
        plot_hist(agged, agg_over, data, s3, s3_base, bucket_name)
        label = '%s per %s' % (agged, agg_over)
        data_dict[label.capitalize()] = {'mean': data.mean(), 'std': data.std(),
                                         'median': np.median(data)}

    # Produce the summary report
    text_report_str = ''
    top_labels = ['Total readings', 'Content processed', 'Statements produced']
    for label in top_labels:
        text_report_str += '%s: %d\n' % (label, data_dict[label])

    for label, data in data_dict.items():
        if label in top_labels:
            continue
        if isinstance(data, dict):
            text_report_str += '%s:\n' % label
            text_report_str += '\n'.join(['\t%s: %d' % (k, v)
                                          for k, v in data.items()])
            text_report_str += '\n'
        else:
            text_report_str += '%s: %d\n' % (label, data)

    s3.put_object(Key=s3_base + 'summary.txt', Body=text_report_str,
                  Bucket=bucket_name)
    ends['stats'] = datetime.now()

    # Report on the timing
    timing_str = ''
    for step in ['reading', 'statement production', 'stats']:
        time_taken = ends[step] - starts[step]
        timing_str += ('%22s: start: %s, end: %s, duration: %s\n'
                       % (step, starts[step], ends[step], time_taken))

    s3.put_object(Key=s3_base + 'time.txt', Body=timing_str, Bucket=bucket_name)
    return

if __name__ == '__main__':
    from argparse import ArgumentParser
    parser = ArgumentParser(
        description=__doc__
        )
    parser.add_argument(
        dest='basename',
        help='The name of this run.'
        )
    parser.add_argument(
        dest='out_dir',
        help='The name of the temporary output directory'
        )
    parser.add_argument(
        dest='read_mode',
        choices=['all', 'unread', 'none'],
        help=("Set the reading mode. If 'all', read everything, if "
              "'unread', only read content that does not have pre-existing "
              "readings of the same reader and version, if 'none', only "
              "use pre-existing readings. Default is 'unread'.")
        )
    parser.add_argument(
        dest='stmt_mode',
        choices=['all', 'unread', 'none'],
        help=("Choose which readings should produce statements. If 'all', all "
              "readings that are produced or retrieved will be used to produce "
              "statements. If 'unread', only produce statements from "
              "previously unread content. If 'none', do not produce any "
              "statements (only readings will be produced).")
        )
    parser.add_argument(
        dest='num_cores',
        help='Select the number of cores on which to run.',
        type=int
        )
    parser.add_argument(
        dest='start_index',
        help='Select the index of the first pmid in the list to read.',
        type=int
        )
    parser.add_argument(
        dest='end_index',
        help='Select the index of the last pmid in the list to read.',
        type=int
        )
    parser.add_argument(
        '-r', '--readers',
        dest='readers',
        choices=['reach', 'sparser'],
        nargs='+',
        help='Choose which reader(s) to use.'
        )
    parser.add_argument(
        '--force_fulltext',
        action='store_true',
        help='Require that content be fulltext, skip anything that isn\'t.'
        )
    parser.add_argument(
        '--use_best_fulltext',
        action='store_true',
        help='Read only the "best" fulltext available for a given id.'
        )
    parser.add_argument(
        '--test',
        action='store_true',
        help="Use the test database."
        )
    args = parser.parse_args()

    logger = logging.getLogger('read_db_aws')
    logger.setLevel(logging.DEBUG)

    client = boto3.client('s3')
    bucket_name = 'bigmech'
    id_list_key = 'reading_inputs/%s/id_list' % args.basename
    readers = [reader_class(args.basename, args.num_cores)
               for reader_class in get_readers()
               if reader_class.name.lower() in args.readers]

    try:
        id_list_obj = client.get_object(
            Bucket=bucket_name,
            Key=id_list_key
            )
    # Handle a missing object gracefully
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == 'NoSuchKey':
            logger.info('Could not find PMID list file at %s, exiting' %
                        id_list_key)
            sys.exit(1)
        # If there was some other kind of problem, re-raise the exception
        else:
            raise e
    # Get the content from the object
    id_list_str = id_list_obj['Body'].read().decode('utf8').strip()
    id_str_list = id_list_str.splitlines()[args.start_index:args.end_index]
    random.shuffle(id_str_list)
    id_dict = get_id_dict([line.strip() for line in id_str_list])

    # Some combinations of options don't make sense:
    forbidden_combos = [('all', 'unread'), ('none', 'unread'), ('none', 'none')]
    assert (args.read_mode, args.stmt_mode) not in forbidden_combos, \
        ("The combination of reading mode %s and statement mode %s is not "
         "allowed." % (args.reading_mode, args.stmt_mode))

    # Init some timing dicts
    starts = {}
    ends = {}

    # Get a handle for the database
    if args.test:
        from indra.db import util as dbu
        db = dbu.get_test_db()
    else:
        db = None

    # Read everything ========================================
    starts['reading'] = datetime.now()
    outputs = produce_readings(id_dict, readers, verbose=True,
                               read_mode=args.read_mode,
                               get_preexisting=(args.stmt_mode == 'all'),
                               force_fulltext=args.force_fulltext,
                               prioritize=args.use_best_fulltext, db=db)
    ends['reading'] = datetime.now()

    # Preserve the sparser logs
    contents = os.listdir('.')
    sparser_logs = [fname for fname in contents
                    if fname.startswith('sparser') and fname.endswith('log')]
    sparser_log_dir = ('reading_results/%s/logs/run_db_reading_queue/'
                       'sparser_logs_%s/') % (
                           args.basename,
                           datetime.now().strftime('%Y%m%d_%H%M%S')
                           )
    for fname in sparser_logs:
        s3_key = sparser_log_dir + fname
        logger.info("Saving sparser logs to %s on s3 in %s."
                    % (s3_key, bucket_name))
        with open(fname, 'r') as f:
            client.put_object(Key=s3_key, Body=f.read(),
                              Bucket=bucket_name)

    # Convert the outputs to statements ==================================
    if args.stmt_mode != 'none':
        starts['statement production'] = datetime.now()
        stmt_data = produce_statements(outputs, n_proc=args.num_cores, db=db)
        ends['statement production'] = datetime.now()
    else:
        stmt_data = []

    report_statistics(outputs, stmt_data, starts, ends, args.basename, client,
                      bucket_name)

<|endoftext|>