<commit_before># Copyright 2018 Microsoft Corporation
#
# 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.
#
# Requires Python 2.6+ and Openssl 1.0+
#

import atexit
import datetime
import json
import os
import sys
import time
import traceback

from datetime import datetime

import azurelinuxagent.common.conf as conf
import azurelinuxagent.common.logger as logger

from azurelinuxagent.common.exception import EventError
from azurelinuxagent.common.future import ustr
from azurelinuxagent.common.protocol.restapi import TelemetryEventParam, \
    TelemetryEvent, \
    get_properties
from azurelinuxagent.common.version import CURRENT_VERSION

_EVENT_MSG = "Event: name={0}, op={1}, message={2}, duration={3}"

class WALAEventOperation:
    ActivateResourceDisk = "ActivateResourceDisk"
    AgentBlacklisted = "AgentBlacklisted"
    AgentEnabled = "AgentEnabled"
    AutoUpdate = "AutoUpdate"
    CustomData = "CustomData"
    Deploy = "Deploy"
    Disable = "Disable"
    Downgrade = "Downgrade"
    Download = "Download"
    Enable = "Enable"
    ExtensionProcessing = "ExtensionProcessing"
    Firewall = "Firewall"
    HealthCheck = "HealthCheck"
    HeartBeat = "HeartBeat"
    HostPlugin = "HostPlugin"
    HttpErrors = "HttpErrors"
    Install = "Install"
    InitializeHostPlugin = "InitializeHostPlugin"
    Log = "Log"
    Partition = "Partition"
    ProcessGoalState = "ProcessGoalState"
    Provision = "Provision"
    ReportStatus = "ReportStatus"
    Restart = "Restart"
    SkipUpdate = "SkipUpdate"
    UnhandledError = "UnhandledError"
    UnInstall = "UnInstall"
    Unknown = "Unknown"
    Upgrade = "Upgrade"
    Update = "Update"

class EventStatus(object):
    EVENT_STATUS_FILE = "event_status.json"

    def __init__(self):
        self._path = None
        self._status = {}

    def clear(self):
        self._status = {}
        self._save()

    def event_marked(self, name, version, op):
        return self._event_name(name, version, op) in self._status

    def event_succeeded(self, name, version, op):
        event = self._event_name(name, version, op)
        if event not in self._status:
            return True
        return self._status[event] is True

    def initialize(self, status_dir=conf.get_lib_dir()):
        self._path = os.path.join(status_dir, EventStatus.EVENT_STATUS_FILE)
        self._load()

    def mark_event_status(self, name, version, op, status):
        event = self._event_name(name, version, op)
        self._status[event] = (status is True)
        self._save()

    def _event_name(self, name, version, op):
        return "{0}-{1}-{2}".format(name, version, op)

    def _load(self):
        try:
            self._status = {}
            if os.path.isfile(self._path):
                with open(self._path, 'r') as f:
                    self._status = json.load(f)
        except Exception as e:
            logger.warn("Exception occurred loading event status: {0}".format(e))
            self._status = {}

    def _save(self):
        try:
            with open(self._path, 'w') as f:
                json.dump(self._status, f)
        except Exception as e:
            logger.warn("Exception occurred saving event status: {0}".format(e))

__event_status__ = EventStatus()
__event_status_operations__ = [
        WALAEventOperation.AutoUpdate,
        WALAEventOperation.ReportStatus
    ]

def _log_event(name, op, message, duration, is_success=True):
    global _EVENT_MSG

    if not is_success:
        logger.error(_EVENT_MSG, name, op, message, duration)
    else:
        logger.info(_EVENT_MSG, name, op, message, duration)

class EventLogger(object):
    def __init__(self):
        self.event_dir = None
        self.periodic_events = {}

    def save_event(self, data):
        if self.event_dir is None:
            logger.warn("Cannot save event -- Event reporter is not initialized.")
            return

        if not os.path.exists(self.event_dir):
            os.mkdir(self.event_dir)
            os.chmod(self.event_dir, 0o700)

        existing_events = os.listdir(self.event_dir)
        if len(existing_events) >= 1000:
            existing_events.sort()
            oldest_files = existing_events[:-999]
            logger.warn("Too many files under: {0}, removing oldest".format(self.event_dir))
            try:
                for f in oldest_files:
                    os.remove(os.path.join(self.event_dir, f))
            except IOError as e:
                raise EventError(e)

        filename = os.path.join(self.event_dir,
                                ustr(int(time.time() * 1000000)))
        try:
            with open(filename + ".tmp", 'wb+') as hfile:
                hfile.write(data.encode("utf-8"))
            os.rename(filename + ".tmp", filename + ".tld")
        except IOError as e:
            raise EventError("Failed to write events to file:{0}", e)

    def reset_periodic(self):
        self.periodic_events = {}

    def is_period_elapsed(self, delta, h):
        return h not in self.periodic_events or \
            (self.periodic_events[h] + delta) <= datetime.now()

    def add_periodic(self,
        delta, name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
        version=CURRENT_VERSION, message="", evt_type="",
        is_internal=False, log_event=True, force=False):

        h = hash(name+op+ustr(is_success)+message)

        if force or self.is_period_elapsed(delta, h):
            self.add_event(name,
                op=op, is_success=is_success, duration=duration,
                version=version, message=message, evt_type=evt_type,
                is_internal=is_internal, log_event=log_event)
            self.periodic_events[h] = datetime.now()

    def add_event(self,
                  name,
                  op=WALAEventOperation.Unknown,
                  is_success=True,
                  duration=0,
                  version=CURRENT_VERSION,
                  message="",
                  evt_type="",
                  is_internal=False,
                  log_event=True):

        if not is_success or log_event:
            _log_event(name, op, message, duration, is_success=is_success)

        event = TelemetryEvent(1, "69B669B9-4AF8-4C50-BDC4-6006FA76E975")
        event.parameters.append(TelemetryEventParam('Name', name))
        event.parameters.append(TelemetryEventParam('Version', str(version)))
        event.parameters.append(TelemetryEventParam('IsInternal', is_internal))
        event.parameters.append(TelemetryEventParam('Operation', op))
        event.parameters.append(TelemetryEventParam('OperationSuccess',
                                                    is_success))
        event.parameters.append(TelemetryEventParam('Message', message))
        event.parameters.append(TelemetryEventParam('Duration', duration))
        event.parameters.append(TelemetryEventParam('ExtensionType', evt_type))

        data = get_properties(event)
        try:
            self.save_event(json.dumps(data))
        except EventError as e:
            logger.error("{0}", e)

    def add_log_event(self, level, message):
        # By the time the message has gotten to this point it is formatted as
        #
        #   YYYY/MM/DD HH:mm:ss.fffffff LEVEL <text>.
        #
        # The timestamp and the level are redundant, and should be stripped.
        # The logging library does not schematize this data, so I am forced
        # to parse the message.  The format is regular, so the burden is low.

        parts = message.split(' ', 3)
        msg = parts[3] if len(parts) == 4 \
            else message

        event = TelemetryEvent(7, "FFF0196F-EE4C-4EAF-9AA5-776F622DEB4F")
        event.parameters.append(TelemetryEventParam('EventName', WALAEventOperation.Log))
        event.parameters.append(TelemetryEventParam('CapabilityUsed', logger.LogLevel.STRINGS[level]))
        event.parameters.append(TelemetryEventParam('Context1', msg))
        event.parameters.append(TelemetryEventParam('Context2', ''))
        event.parameters.append(TelemetryEventParam('Context3', ''))

        data = get_properties(event)
        try:
            self.save_event(json.dumps(data))
        except EventError:
            pass

__event_logger__ = EventLogger()

def elapsed_milliseconds(utc_start):
    now = datetime.utcnow()
    if now < utc_start:
        return 0

    d = now - utc_start
    return int(((d.days * 24 * 60 * 60 + d.seconds) * 1000) + \
                    (d.microseconds / 1000.0))

def report_event(op, is_success=True, message=''):
    from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
    add_event(AGENT_NAME,
              version=CURRENT_VERSION,
              is_success=is_success,
              message=message,
              op=op)

def report_periodic(delta, op, is_success=True, message=''):
    from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
    add_periodic(delta, AGENT_NAME,
              version=CURRENT_VERSION,
              is_success=is_success,
              message=message,
              op=op)

def add_event(name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
              version=CURRENT_VERSION,
              message="", evt_type="", is_internal=False, log_event=True,
              reporter=__event_logger__):
    if reporter.event_dir is None:
        logger.warn("Cannot add event -- Event reporter is not initialized.")
        _log_event(name, op, message, duration, is_success=is_success)
        return

    if should_emit_event(name, version, op, is_success):
        mark_event_status(name, version, op, is_success)
        reporter.add_event(
            name, op=op, is_success=is_success, duration=duration,
            version=str(version), message=message, evt_type=evt_type,
            is_internal=is_internal, log_event=log_event)

def add_log_event(level, message, reporter=__event_logger__):
    if reporter.event_dir is None:
        return

    reporter.add_log_event(level, message)

def add_periodic(
    delta, name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
    version=CURRENT_VERSION,
    message="", evt_type="", is_internal=False, log_event=True, force=False,
    reporter=__event_logger__):
    if reporter.event_dir is None:
        logger.warn("Cannot add periodic event -- Event reporter is not initialized.")
        _log_event(name, op, message, duration, is_success=is_success)
        return

    reporter.add_periodic(
        delta, name, op=op, is_success=is_success, duration=duration,
        version=str(version), message=message, evt_type=evt_type,
        is_internal=is_internal, log_event=log_event, force=force)

def mark_event_status(name, version, op, status):
    if op in __event_status_operations__:
        __event_status__.mark_event_status(name, version, op, status)

def should_emit_event(name, version, op, status):
    return \
        op not in __event_status_operations__ or \
        __event_status__ is None or \
        not __event_status__.event_marked(name, version, op) or \
        __event_status__.event_succeeded(name, version, op) != status

def init_event_logger(event_dir):
    __event_logger__.event_dir = event_dir

def init_event_status(status_dir):
    __event_status__.initialize(status_dir)

def dump_unhandled_err(name):
    if hasattr(sys, 'last_type') and hasattr(sys, 'last_value') and \
            hasattr(sys, 'last_traceback'):
        last_type = getattr(sys, 'last_type')
        last_value = getattr(sys, 'last_value')
        last_traceback = getattr(sys, 'last_traceback')
        error = traceback.format_exception(last_type, last_value,
                                           last_traceback)
        message = "".join(error)
        add_event(name, is_success=False, message=message,
                  op=WALAEventOperation.UnhandledError)

def enable_unhandled_err_dump(name):
    atexit.register(dump_unhandled_err, name)
<commit_msg>encode stdout/stderr from extension commands<commit_after># Copyright 2018 Microsoft Corporation
#
# 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.
#
# Requires Python 2.6+ and Openssl 1.0+
#

import atexit
import base64
import datetime
import json
import os
import sys
import time
import traceback
import zlib

from datetime import datetime

import azurelinuxagent.common.conf as conf
import azurelinuxagent.common.logger as logger

from azurelinuxagent.common.exception import EventError
from azurelinuxagent.common.future import ustr
from azurelinuxagent.common.protocol.restapi import TelemetryEventParam, \
    TelemetryEvent, \
    get_properties
from azurelinuxagent.common.version import CURRENT_VERSION

_EVENT_MSG = "Event: name={0}, op={1}, message={2}, duration={3}"

class WALAEventOperation:
    ActivateResourceDisk = "ActivateResourceDisk"
    AgentBlacklisted = "AgentBlacklisted"
    AgentEnabled = "AgentEnabled"
    AutoUpdate = "AutoUpdate"
    CustomData = "CustomData"
    Deploy = "Deploy"
    Disable = "Disable"
    Downgrade = "Downgrade"
    Download = "Download"
    Enable = "Enable"
    ExtensionProcessing = "ExtensionProcessing"
    Firewall = "Firewall"
    HealthCheck = "HealthCheck"
    HeartBeat = "HeartBeat"
    HostPlugin = "HostPlugin"
    HttpErrors = "HttpErrors"
    Install = "Install"
    InitializeHostPlugin = "InitializeHostPlugin"
    Log = "Log"
    Partition = "Partition"
    ProcessGoalState = "ProcessGoalState"
    Provision = "Provision"
    ReportStatus = "ReportStatus"
    Restart = "Restart"
    SkipUpdate = "SkipUpdate"
    UnhandledError = "UnhandledError"
    UnInstall = "UnInstall"
    Unknown = "Unknown"
    Upgrade = "Upgrade"
    Update = "Update"

SHOULD_ENCODE_MESSAGE_LEN = 80
SHOULD_ENCODE_MESSAGE_OP = [
    WALAEventOperation.Disable,
    WALAEventOperation.Enable,
    WALAEventOperation.Install,
    WALAEventOperation.UnInstall,
]

class EventStatus(object):
    EVENT_STATUS_FILE = "event_status.json"

    def __init__(self):
        self._path = None
        self._status = {}

    def clear(self):
        self._status = {}
        self._save()

    def event_marked(self, name, version, op):
        return self._event_name(name, version, op) in self._status

    def event_succeeded(self, name, version, op):
        event = self._event_name(name, version, op)
        if event not in self._status:
            return True
        return self._status[event] is True

    def initialize(self, status_dir=conf.get_lib_dir()):
        self._path = os.path.join(status_dir, EventStatus.EVENT_STATUS_FILE)
        self._load()

    def mark_event_status(self, name, version, op, status):
        event = self._event_name(name, version, op)
        self._status[event] = (status is True)
        self._save()

    def _event_name(self, name, version, op):
        return "{0}-{1}-{2}".format(name, version, op)

    def _load(self):
        try:
            self._status = {}
            if os.path.isfile(self._path):
                with open(self._path, 'r') as f:
                    self._status = json.load(f)
        except Exception as e:
            logger.warn("Exception occurred loading event status: {0}".format(e))
            self._status = {}

    def _save(self):
        try:
            with open(self._path, 'w') as f:
                json.dump(self._status, f)
        except Exception as e:
            logger.warn("Exception occurred saving event status: {0}".format(e))

__event_status__ = EventStatus()
__event_status_operations__ = [
        WALAEventOperation.AutoUpdate,
        WALAEventOperation.ReportStatus
    ]

def _encode_message(op, message):
    """
    Gzip and base64 encode a message based on the operation.

    The intent of this message is to make the logs human readable and include the
    stdout/stderr from extension operations.  Extension operations tend to generate
    a lot of noise, which makes it difficult to parse the line-oriented waagent.log.
    The compromise is to encode the stdout/stderr so we preserve the data and do
    not destroy the line oriented nature.

    The data can be recovered using the following command:

      $ echo '<encoded data>' | base64 -d | pigz -zd

    You may need to install the pigz command.

    :param op: Operation, e.g. Enable or Install
    :param message: Message to encode
    :return: gzip'ed and base64 encoded message, or the original message
    """

    if len(message) == 0:
        return message

    if op not in SHOULD_ENCODE_MESSAGE_OP:
        return message

    try:
        return base64.b64encode(zlib.compress(message))
    except Exception:
        # If the message could not be encoded and dummy message is returned.
        # The original message was still sent via telemetry, so all is not lost.
        return "<>"

def _log_event(name, op, message, duration, is_success=True):
    global _EVENT_MSG

    message = _encode_message(op, message)
    if not is_success:
        logger.error(_EVENT_MSG, name, op, message, duration)
    else:
        logger.info(_EVENT_MSG, name, op, message, duration)

class EventLogger(object):
    def __init__(self):
        self.event_dir = None
        self.periodic_events = {}

    def save_event(self, data):
        if self.event_dir is None:
            logger.warn("Cannot save event -- Event reporter is not initialized.")
            return

        if not os.path.exists(self.event_dir):
            os.mkdir(self.event_dir)
            os.chmod(self.event_dir, 0o700)

        existing_events = os.listdir(self.event_dir)
        if len(existing_events) >= 1000:
            existing_events.sort()
            oldest_files = existing_events[:-999]
            logger.warn("Too many files under: {0}, removing oldest".format(self.event_dir))
            try:
                for f in oldest_files:
                    os.remove(os.path.join(self.event_dir, f))
            except IOError as e:
                raise EventError(e)

        filename = os.path.join(self.event_dir,
                                ustr(int(time.time() * 1000000)))
        try:
            with open(filename + ".tmp", 'wb+') as hfile:
                hfile.write(data.encode("utf-8"))
            os.rename(filename + ".tmp", filename + ".tld")
        except IOError as e:
            raise EventError("Failed to write events to file:{0}", e)

    def reset_periodic(self):
        self.periodic_events = {}

    def is_period_elapsed(self, delta, h):
        return h not in self.periodic_events or \
            (self.periodic_events[h] + delta) <= datetime.now()

    def add_periodic(self,
        delta, name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
        version=CURRENT_VERSION, message="", evt_type="",
        is_internal=False, log_event=True, force=False):

        h = hash(name+op+ustr(is_success)+message)

        if force or self.is_period_elapsed(delta, h):
            self.add_event(name,
                op=op, is_success=is_success, duration=duration,
                version=version, message=message, evt_type=evt_type,
                is_internal=is_internal, log_event=log_event)
            self.periodic_events[h] = datetime.now()

    def add_event(self,
                  name,
                  op=WALAEventOperation.Unknown,
                  is_success=True,
                  duration=0,
                  version=CURRENT_VERSION,
                  message="",
                  evt_type="",
                  is_internal=False,
                  log_event=True):

        if not is_success or log_event:
            _log_event(name, op, message, duration, is_success=is_success)

        event = TelemetryEvent(1, "69B669B9-4AF8-4C50-BDC4-6006FA76E975")
        event.parameters.append(TelemetryEventParam('Name', name))
        event.parameters.append(TelemetryEventParam('Version', str(version)))
        event.parameters.append(TelemetryEventParam('IsInternal', is_internal))
        event.parameters.append(TelemetryEventParam('Operation', op))
        event.parameters.append(TelemetryEventParam('OperationSuccess',
                                                    is_success))
        event.parameters.append(TelemetryEventParam('Message', message))
        event.parameters.append(TelemetryEventParam('Duration', duration))
        event.parameters.append(TelemetryEventParam('ExtensionType', evt_type))

        data = get_properties(event)
        try:
            self.save_event(json.dumps(data))
        except EventError as e:
            logger.error("{0}", e)

    def add_log_event(self, level, message):
        # By the time the message has gotten to this point it is formatted as
        #
        #   YYYY/MM/DD HH:mm:ss.fffffff LEVEL <text>.
        #
        # The timestamp and the level are redundant, and should be stripped.
        # The logging library does not schematize this data, so I am forced
        # to parse the message.  The format is regular, so the burden is low.

        parts = message.split(' ', 3)
        msg = parts[3] if len(parts) == 4 \
            else message

        event = TelemetryEvent(7, "FFF0196F-EE4C-4EAF-9AA5-776F622DEB4F")
        event.parameters.append(TelemetryEventParam('EventName', WALAEventOperation.Log))
        event.parameters.append(TelemetryEventParam('CapabilityUsed', logger.LogLevel.STRINGS[level]))
        event.parameters.append(TelemetryEventParam('Context1', msg))
        event.parameters.append(TelemetryEventParam('Context2', ''))
        event.parameters.append(TelemetryEventParam('Context3', ''))

        data = get_properties(event)
        try:
            self.save_event(json.dumps(data))
        except EventError:
            pass

__event_logger__ = EventLogger()

def elapsed_milliseconds(utc_start):
    now = datetime.utcnow()
    if now < utc_start:
        return 0

    d = now - utc_start
    return int(((d.days * 24 * 60 * 60 + d.seconds) * 1000) + \
                    (d.microseconds / 1000.0))

def report_event(op, is_success=True, message=''):
    from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
    add_event(AGENT_NAME,
              version=CURRENT_VERSION,
              is_success=is_success,
              message=message,
              op=op)

def report_periodic(delta, op, is_success=True, message=''):
    from azurelinuxagent.common.version import AGENT_NAME, CURRENT_VERSION
    add_periodic(delta, AGENT_NAME,
              version=CURRENT_VERSION,
              is_success=is_success,
              message=message,
              op=op)

def add_event(name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
              version=CURRENT_VERSION,
              message="", evt_type="", is_internal=False, log_event=True,
              reporter=__event_logger__):
    if reporter.event_dir is None:
        logger.warn("Cannot add event -- Event reporter is not initialized.")
        _log_event(name, op, message, duration, is_success=is_success)
        return

    if should_emit_event(name, version, op, is_success):
        mark_event_status(name, version, op, is_success)
        reporter.add_event(
            name, op=op, is_success=is_success, duration=duration,
            version=str(version), message=message, evt_type=evt_type,
            is_internal=is_internal, log_event=log_event)

def add_log_event(level, message, reporter=__event_logger__):
    if reporter.event_dir is None:
        return

    reporter.add_log_event(level, message)

def add_periodic(
    delta, name, op=WALAEventOperation.Unknown, is_success=True, duration=0,
    version=CURRENT_VERSION,
    message="", evt_type="", is_internal=False, log_event=True, force=False,
    reporter=__event_logger__):
    if reporter.event_dir is None:
        logger.warn("Cannot add periodic event -- Event reporter is not initialized.")
        _log_event(name, op, message, duration, is_success=is_success)
        return

    reporter.add_periodic(
        delta, name, op=op, is_success=is_success, duration=duration,
        version=str(version), message=message, evt_type=evt_type,
        is_internal=is_internal, log_event=log_event, force=force)

def mark_event_status(name, version, op, status):
    if op in __event_status_operations__:
        __event_status__.mark_event_status(name, version, op, status)

def should_emit_event(name, version, op, status):
    return \
        op not in __event_status_operations__ or \
        __event_status__ is None or \
        not __event_status__.event_marked(name, version, op) or \
        __event_status__.event_succeeded(name, version, op) != status

def init_event_logger(event_dir):
    __event_logger__.event_dir = event_dir

def init_event_status(status_dir):
    __event_status__.initialize(status_dir)

def dump_unhandled_err(name):
    if hasattr(sys, 'last_type') and hasattr(sys, 'last_value') and \
            hasattr(sys, 'last_traceback'):
        last_type = getattr(sys, 'last_type')
        last_value = getattr(sys, 'last_value')
        last_traceback = getattr(sys, 'last_traceback')
        error = traceback.format_exception(last_type, last_value,
                                           last_traceback)
        message = "".join(error)
        add_event(name, is_success=False, message=message,
                  op=WALAEventOperation.UnhandledError)

def enable_unhandled_err_dump(name):
    atexit.register(dump_unhandled_err, name)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (C) 2007 Edgewall Software
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://babel.edgewall.org/wiki/License.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://babel.edgewall.org/log/.

import codecs
import doctest
from StringIO import StringIO
import sys
import unittest

from babel.messages import extract

class ExtractPythonTestCase(unittest.TestCase):

    def test_nested_calls(self):
        buf = StringIO("""\
msg1 = _(i18n_arg.replace(r'\"', '"'))
msg2 = ungettext(i18n_arg.replace(r'\"', '"'), multi_arg.replace(r'\"', '"'), 2)
msg3 = ungettext("Babel", multi_arg.replace(r'\"', '"'), 2)
msg4 = ungettext(i18n_arg.replace(r'\"', '"'), "Babels", 2)
msg5 = ungettext('bunny', 'bunnies', random.randint(1, 2))
msg6 = ungettext(arg0, 'bunnies', random.randint(1, 2))
msg7 = _(hello.there)
msg8 = gettext('Rabbit')
msg9 = dgettext('wiki', model.addPage())
msg10 = dngettext(getDomain(), 'Page', 'Pages', 3)
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([
                (1, '_', None, []),
                (2, 'ungettext', (None, None, None), []),
                (3, 'ungettext', (u'Babel', None, None), []),
                (4, 'ungettext', (None, u'Babels', None), []),
                (5, 'ungettext', (u'bunny', u'bunnies', None), []),
                (6, 'ungettext', (None, u'bunnies', None), []),
                (7, '_', None, []),
                (8, 'gettext', u'Rabbit', []),
                (9, 'dgettext', (u'wiki', None), []),
                (10, 'dngettext', (None, u'Page', u'Pages', None), [])],
                         messages)

    def test_nested_comments(self):
        buf = StringIO("""\
msg = ngettext('pylon',  # TRANSLATORS: shouldn't be
               'pylons', # TRANSLATORS: seeing this
               count)
""")
        messages = list(extract.extract_python(buf, ('ngettext',),
                                               ['TRANSLATORS:'], {}))
        self.assertEqual([(1, 'ngettext', (u'pylon', u'pylons', None), [])],
                         messages)

    def test_declarations(self):
        buf = StringIO("""\
class gettext(object):
    pass
def render_body(context,x,y=_('Page arg 1'),z=_('Page arg 2'),**pageargs):
    pass
def ngettext(y='arg 1',z='arg 2',**pageargs):
    pass
class Meta:
    verbose_name = _('log entry')
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([(3, '_', u'Page arg 1', []),
                          (3, '_', u'Page arg 2', []),
                          (8, '_', u'log entry', [])],
                         messages)

    def test_multiline(self):
        buf = StringIO("""\
msg1 = ngettext('pylon',
                'pylons', count)
msg2 = ngettext('elvis',
                'elvises',
                 count)
""")
        messages = list(extract.extract_python(buf, ('ngettext',), [], {}))
        self.assertEqual([(1, 'ngettext', (u'pylon', u'pylons', None), []),
                          (3, 'ngettext', (u'elvis', u'elvises', None), [])],
                         messages)

    def test_triple_quoted_strings(self):
        buf = StringIO("""\
msg1 = _('''pylons''')
msg2 = ngettext(r'''elvis''', \"\"\"elvises\"\"\", count)
msg2 = ngettext(\"\"\"elvis\"\"\", 'elvises', count)
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([(1, '_', (u'pylons'), []),
                          (2, 'ngettext', (u'elvis', u'elvises', None), []),
                          (3, 'ngettext', (u'elvis', u'elvises', None), [])],
                         messages)

    def test_multiline_strings(self):
        buf = StringIO("""\
_('''This module provides internationalization and localization
support for your Python programs by providing an interface to the GNU
gettext message catalog library.''')
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual(
            [(1, '_',
              u'This module provides internationalization and localization\n'
              'support for your Python programs by providing an interface to '
              'the GNU\ngettext message catalog library.', [])],
            messages)

    def test_unicode_string_arg(self):
        buf = StringIO("msg = _(u'Foo Bar')")
        messages = list(extract.extract_python(buf, ('_',), [], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])

    def test_comment_tag(self):
        buf = StringIO("""
# NOTE: A translation comment
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment'], messages[0][3])

    def test_comment_tag_multiline(self):
        buf = StringIO("""
# NOTE: A translation comment
# with a second line
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment', u'with a second line'],
                         messages[0][3])

    def test_translator_comments_with_previous_non_translator_comments(self):
        buf = StringIO("""
# This shouldn't be in the output
# because it didn't start with a comment tag
# NOTE: A translation comment
# with a second line
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment', u'with a second line'],
                         messages[0][3])

    def test_comment_tags_not_on_start_of_comment(self):
        buf = StringIO("""
# This shouldn't be in the output
# because it didn't start with a comment tag
# do NOTE: this will no be a translation comment
# NOTE: This one will be
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'This one will be'], messages[0][3])

    def test_multiple_comment_tags(self):
        buf = StringIO("""
# NOTE1: A translation comment for tag1
# with a second line
msg = _(u'Foo Bar1')

# NOTE2: A translation comment for tag2
msg = _(u'Foo Bar2')
""")
        messages = list(extract.extract_python(buf, ('_',),
                                               ['NOTE1:', 'NOTE2:'], {}))
        self.assertEqual(u'Foo Bar1', messages[0][2])
        self.assertEqual([u'A translation comment for tag1',
                          u'with a second line'], messages[0][3])
        self.assertEqual(u'Foo Bar2', messages[1][2])
        self.assertEqual([u'A translation comment for tag2'], messages[1][3])

    def test_two_succeeding_comments(self):
        buf = StringIO("""
# NOTE: one
# NOTE: two
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'one', u'NOTE: two'], messages[0][3])

    def test_invalid_translator_comments(self):
        buf = StringIO("""
# NOTE: this shouldn't apply to any messages
hello = 'there'

msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([], messages[0][3])

    def test_invalid_translator_comments2(self):
        buf = StringIO("""
# NOTE: Hi!
hithere = _('Hi there!')

# NOTE: you should not be seeing this in the .po
rows = [[v for v in range(0,10)] for row in range(0,10)]

# this (NOTE:) should not show up either
hello = _('Hello')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Hi there!', messages[0][2])
        self.assertEqual([u'Hi!'], messages[0][3])
        self.assertEqual(u'Hello', messages[1][2])
        self.assertEqual([], messages[1][3])

    def test_invalid_translator_comments3(self):
        buf = StringIO("""
# NOTE: Hi,

# there!
hithere = _('Hi there!')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Hi there!', messages[0][2])
        self.assertEqual([], messages[0][3])

    def test_utf8_message(self):
        buf = StringIO("""
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'],
                                               {'encoding': 'utf-8'}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_message_with_magic_comment(self):
        buf = StringIO("""# -*- coding: utf-8 -*-
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_message_with_utf8_bom(self):
        buf = StringIO(codecs.BOM_UTF8 + """
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_raw_strings_match_unicode_strings(self):
        buf = StringIO(codecs.BOM_UTF8 + """
msg = _('Bonjour à tous')
msgu = _(u'Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual(messages[0][2], messages[1][2])

class ExtractTestCase(unittest.TestCase):

    def test_invalid_filter(self):
        buf = StringIO("""\
msg1 = _(i18n_arg.replace(r'\"', '"'))
msg2 = ungettext(i18n_arg.replace(r'\"', '"'), multi_arg.replace(r'\"', '"'), 2)
msg3 = ungettext("Babel", multi_arg.replace(r'\"', '"'), 2)
msg4 = ungettext(i18n_arg.replace(r'\"', '"'), "Babels", 2)
msg5 = ungettext('bunny', 'bunnies', random.randint(1, 2))
msg6 = ungettext(arg0, 'bunnies', random.randint(1, 2))
msg7 = _(hello.there)
msg8 = gettext('Rabbit')
msg9 = dgettext('wiki', model.addPage())
msg10 = dngettext(domain, 'Page', 'Pages', 3)
""")
        messages = \
            list(extract.extract('python', buf, extract.DEFAULT_KEYWORDS, [],
                                 {}))
        self.assertEqual([(5, (u'bunny', u'bunnies'), []),
                          (8, u'Rabbit', []),
                          (10, (u'Page', u'Pages'), [])], messages)

    def test_empty_string_msgid(self):
        buf = StringIO("""\
msg = _('')
""")
        stderr = sys.stderr
        sys.stderr = StringIO()
        try:
            messages = \
                list(extract.extract('python', buf, extract.DEFAULT_KEYWORDS,
                                     [], {}))
            self.assertEqual([], messages)
            assert 'warning: Empty msgid.' in sys.stderr.getvalue()
        finally:
            sys.stderr = stderr

def suite():
    suite = unittest.TestSuite()
    suite.addTest(doctest.DocTestSuite(extract))
    suite.addTest(unittest.makeSuite(ExtractPythonTestCase))
    suite.addTest(unittest.makeSuite(ExtractTestCase))
    return suite

if __name__ == '__main__':
    unittest.main(defaultTest='suite')
<commit_msg>add extractor tests for some odder method signatures and concatenated strings<commit_after># -*- coding: utf-8 -*-
#
# Copyright (C) 2007 Edgewall Software
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://babel.edgewall.org/wiki/License.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://babel.edgewall.org/log/.

import codecs
import doctest
from StringIO import StringIO
import sys
import unittest

from babel.messages import extract

class ExtractPythonTestCase(unittest.TestCase):

    def test_nested_calls(self):
        buf = StringIO("""\
msg1 = _(i18n_arg.replace(r'\"', '"'))
msg2 = ungettext(i18n_arg.replace(r'\"', '"'), multi_arg.replace(r'\"', '"'), 2)
msg3 = ungettext("Babel", multi_arg.replace(r'\"', '"'), 2)
msg4 = ungettext(i18n_arg.replace(r'\"', '"'), "Babels", 2)
msg5 = ungettext('bunny', 'bunnies', random.randint(1, 2))
msg6 = ungettext(arg0, 'bunnies', random.randint(1, 2))
msg7 = _(hello.there)
msg8 = gettext('Rabbit')
msg9 = dgettext('wiki', model.addPage())
msg10 = dngettext(getDomain(), 'Page', 'Pages', 3)
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([
                (1, '_', None, []),
                (2, 'ungettext', (None, None, None), []),
                (3, 'ungettext', (u'Babel', None, None), []),
                (4, 'ungettext', (None, u'Babels', None), []),
                (5, 'ungettext', (u'bunny', u'bunnies', None), []),
                (6, 'ungettext', (None, u'bunnies', None), []),
                (7, '_', None, []),
                (8, 'gettext', u'Rabbit', []),
                (9, 'dgettext', (u'wiki', None), []),
                (10, 'dngettext', (None, u'Page', u'Pages', None), [])],
                         messages)

    def test_nested_comments(self):
        buf = StringIO("""\
msg = ngettext('pylon',  # TRANSLATORS: shouldn't be
               'pylons', # TRANSLATORS: seeing this
               count)
""")
        messages = list(extract.extract_python(buf, ('ngettext',),
                                               ['TRANSLATORS:'], {}))
        self.assertEqual([(1, 'ngettext', (u'pylon', u'pylons', None), [])],
                         messages)

    def test_declarations(self):
        buf = StringIO("""\
class gettext(object):
    pass
def render_body(context,x,y=_('Page arg 1'),z=_('Page arg 2'),**pageargs):
    pass
def ngettext(y='arg 1',z='arg 2',**pageargs):
    pass
class Meta:
    verbose_name = _('log entry')
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([(3, '_', u'Page arg 1', []),
                          (3, '_', u'Page arg 2', []),
                          (8, '_', u'log entry', [])],
                         messages)

    def test_multiline(self):
        buf = StringIO("""\
msg1 = ngettext('pylon',
                'pylons', count)
msg2 = ngettext('elvis',
                'elvises',
                 count)
""")
        messages = list(extract.extract_python(buf, ('ngettext',), [], {}))
        self.assertEqual([(1, 'ngettext', (u'pylon', u'pylons', None), []),
                          (3, 'ngettext', (u'elvis', u'elvises', None), [])],
                         messages)

    def test_triple_quoted_strings(self):
        buf = StringIO("""\
msg1 = _('''pylons''')
msg2 = ngettext(r'''elvis''', \"\"\"elvises\"\"\", count)
msg2 = ngettext(\"\"\"elvis\"\"\", 'elvises', count)
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual([(1, '_', (u'pylons'), []),
                          (2, 'ngettext', (u'elvis', u'elvises', None), []),
                          (3, 'ngettext', (u'elvis', u'elvises', None), [])],
                         messages)

    def test_multiline_strings(self):
        buf = StringIO("""\
_('''This module provides internationalization and localization
support for your Python programs by providing an interface to the GNU
gettext message catalog library.''')
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual(
            [(1, '_',
              u'This module provides internationalization and localization\n'
              'support for your Python programs by providing an interface to '
              'the GNU\ngettext message catalog library.', [])],
            messages)

    def test_concatenated_strings(self):
        buf = StringIO("""\
foobar = _('foo' 'bar')
""")
        messages = list(extract.extract_python(buf,
                                               extract.DEFAULT_KEYWORDS.keys(),
                                               [], {}))
        self.assertEqual(u'foobar', messages[0][2])

    def test_unicode_string_arg(self):
        buf = StringIO("msg = _(u'Foo Bar')")
        messages = list(extract.extract_python(buf, ('_',), [], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])

    def test_comment_tag(self):
        buf = StringIO("""
# NOTE: A translation comment
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment'], messages[0][3])

    def test_comment_tag_multiline(self):
        buf = StringIO("""
# NOTE: A translation comment
# with a second line
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment', u'with a second line'],
                         messages[0][3])

    def test_translator_comments_with_previous_non_translator_comments(self):
        buf = StringIO("""
# This shouldn't be in the output
# because it didn't start with a comment tag
# NOTE: A translation comment
# with a second line
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'A translation comment', u'with a second line'],
                         messages[0][3])

    def test_comment_tags_not_on_start_of_comment(self):
        buf = StringIO("""
# This shouldn't be in the output
# because it didn't start with a comment tag
# do NOTE: this will no be a translation comment
# NOTE: This one will be
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'This one will be'], messages[0][3])

    def test_multiple_comment_tags(self):
        buf = StringIO("""
# NOTE1: A translation comment for tag1
# with a second line
msg = _(u'Foo Bar1')

# NOTE2: A translation comment for tag2
msg = _(u'Foo Bar2')
""")
        messages = list(extract.extract_python(buf, ('_',),
                                               ['NOTE1:', 'NOTE2:'], {}))
        self.assertEqual(u'Foo Bar1', messages[0][2])
        self.assertEqual([u'A translation comment for tag1',
                          u'with a second line'], messages[0][3])
        self.assertEqual(u'Foo Bar2', messages[1][2])
        self.assertEqual([u'A translation comment for tag2'], messages[1][3])

    def test_two_succeeding_comments(self):
        buf = StringIO("""
# NOTE: one
# NOTE: two
msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([u'one', u'NOTE: two'], messages[0][3])

    def test_invalid_translator_comments(self):
        buf = StringIO("""
# NOTE: this shouldn't apply to any messages
hello = 'there'

msg = _(u'Foo Bar')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Foo Bar', messages[0][2])
        self.assertEqual([], messages[0][3])

    def test_invalid_translator_comments2(self):
        buf = StringIO("""
# NOTE: Hi!
hithere = _('Hi there!')

# NOTE: you should not be seeing this in the .po
rows = [[v for v in range(0,10)] for row in range(0,10)]

# this (NOTE:) should not show up either
hello = _('Hello')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Hi there!', messages[0][2])
        self.assertEqual([u'Hi!'], messages[0][3])
        self.assertEqual(u'Hello', messages[1][2])
        self.assertEqual([], messages[1][3])

    def test_invalid_translator_comments3(self):
        buf = StringIO("""
# NOTE: Hi,

# there!
hithere = _('Hi there!')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Hi there!', messages[0][2])
        self.assertEqual([], messages[0][3])

    def test_different_signatures(self):
        buf = StringIO("""
foo = _('foo', 'bar')
n = ngettext('hello', 'there', n=3)
n = ngettext(n=3, 'hello', 'there')
n = ngettext(n=3, *messages)
""")
        messages = list(extract.extract_python(buf, ('_', 'ngettext'), [], {}))
        self.assertEqual((u'foo', u'bar'), messages[0][2])
        self.assertEqual((u'hello', u'there', None), messages[1][2])
        self.assertEqual((None, u'hello', u'there'), messages[2][2])
        self.assertEqual((None, None), messages[3][2])

    def test_utf8_message(self):
        buf = StringIO("""
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'],
                                               {'encoding': 'utf-8'}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_message_with_magic_comment(self):
        buf = StringIO("""# -*- coding: utf-8 -*-
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_message_with_utf8_bom(self):
        buf = StringIO(codecs.BOM_UTF8 + """
# NOTE: hello
msg = _('Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual([u'hello'], messages[0][3])

    def test_utf8_raw_strings_match_unicode_strings(self):
        buf = StringIO(codecs.BOM_UTF8 + """
msg = _('Bonjour à tous')
msgu = _(u'Bonjour à tous')
""")
        messages = list(extract.extract_python(buf, ('_',), ['NOTE:'], {}))
        self.assertEqual(u'Bonjour à tous', messages[0][2])
        self.assertEqual(messages[0][2], messages[1][2])

class ExtractTestCase(unittest.TestCase):

    def test_invalid_filter(self):
        buf = StringIO("""\
msg1 = _(i18n_arg.replace(r'\"', '"'))
msg2 = ungettext(i18n_arg.replace(r'\"', '"'), multi_arg.replace(r'\"', '"'), 2)
msg3 = ungettext("Babel", multi_arg.replace(r'\"', '"'), 2)
msg4 = ungettext(i18n_arg.replace(r'\"', '"'), "Babels", 2)
msg5 = ungettext('bunny', 'bunnies', random.randint(1, 2))
msg6 = ungettext(arg0, 'bunnies', random.randint(1, 2))
msg7 = _(hello.there)
msg8 = gettext('Rabbit')
msg9 = dgettext('wiki', model.addPage())
msg10 = dngettext(domain, 'Page', 'Pages', 3)
""")
        messages = \
            list(extract.extract('python', buf, extract.DEFAULT_KEYWORDS, [],
                                 {}))
        self.assertEqual([(5, (u'bunny', u'bunnies'), []),
                          (8, u'Rabbit', []),
                          (10, (u'Page', u'Pages'), [])], messages)

    def test_different_signatures(self):
        buf = StringIO("""
foo = _('foo', 'bar')
n = ngettext('hello', 'there', n=3)
n = ngettext(n=3, 'hello', 'there')
n = ngettext(n=3, *messages)
""")
        messages = \
            list(extract.extract('python', buf, extract.DEFAULT_KEYWORDS, [],
                                 {}))
        self.assertEqual(len(messages), 2)
        self.assertEqual(u'foo', messages[0][1])
        self.assertEqual((u'hello', u'there'), messages[1][1])

    def test_empty_string_msgid(self):
        buf = StringIO("""\
msg = _('')
""")
        stderr = sys.stderr
        sys.stderr = StringIO()
        try:
            messages = \
                list(extract.extract('python', buf, extract.DEFAULT_KEYWORDS,
                                     [], {}))
            self.assertEqual([], messages)
            assert 'warning: Empty msgid.' in sys.stderr.getvalue()
        finally:
            sys.stderr = stderr

def suite():
    suite = unittest.TestSuite()
    suite.addTest(doctest.DocTestSuite(extract))
    suite.addTest(unittest.makeSuite(ExtractPythonTestCase))
    suite.addTest(unittest.makeSuite(ExtractTestCase))
    return suite

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


<commit_before>from django.core.validators import RegexValidator

validate_slug = RegexValidator(
    regex=r'[a-z0-9-]+',
    message='The slug may only contain lower case letters (a–z), numbers (0–9) and dashes (-).'
)

validate_path = RegexValidator(
    regex=r'[a-z0-9-/]+',
    message='The path may only contain lower case letters (a–z), numbers (0–9), dashes (-) and forward slashes (/).'
)

class CommonFields(object):
    path = dict(
        max_length=1023,
        validators=[validate_path],
        verbose_name='Path',
        help_text='Will be automatically determined from parent and slug.',
        unique=True,
    )

    slug = dict(
        blank=True,  # actually not, but autogenerated anyway
        max_length=255,
        validators=[validate_slug],
        verbose_name='Slug',
        help_text=(
            'The slug will be included as part of the URL. '
            'The slug may only contain lower case letters (a–z), numbers (0–9) and dashes (-). '
            'If you leave the slug blank, it will be generated from the title. '
            'If you change the slug after publishing, remember to create any redirects you may need.'
        ),
    )

    title = dict(
        max_length=1023,
        verbose_name='Title',
        help_text='The title will be displayed at the top of the view.',
    )

    description = dict(
        verbose_name='Description',
        help_text=(
            'While the description is not usually displayed in the user interface, '
            'it is visible to search engines and RSS readers. If your item contains '
            'text content, a good practice is to make the description a short summary '
            'of the content.'
        ),
        blank=True,
        default='',
    )

    body = dict(
        blank=True,
        default='',
        verbose_name='Text content',
        help_text='Will be displayed at the top of the album view before subalbums and pictures.',
    )

    order = dict(
        default=0,
        verbose_name='Ordering',
        help_text=(
            'The child items of a common parent are ordered in an ascending order by this number '
            '(smallest first).'
        )
    )

    is_public = dict(
        default=True,
        verbose_name='Public',
        help_text='Items that are not public are only accessible to admin users.',
    )

    is_visible = dict(
        default=True,
        verbose_name='Visible',
        help_text='Items that are not visible are not displayed in listings, but they can still be accessed via exact URL.',
    )
<commit_msg>backend: fix regex validators for slug and path<commit_after>from django.core.validators import RegexValidator

validate_slug = RegexValidator(
    regex=r'^[a-z0-9-]+$',
    message='The slug may only contain lower case letters (a–z), numbers (0–9) and dashes (-).',
    code='invalid_slug',
)

validate_path = RegexValidator(
    regex=r'^/[a-z0-9-/]+$',
    message='The path may only contain lower case letters (a–z), numbers (0–9), dashes (-) and forward slashes (/). The path must start with a slash (/).',
    code='invalid_path',
)

class CommonFields(object):
    path = dict(
        max_length=1023,
        validators=[validate_path],
        verbose_name='Path',
        help_text='Will be automatically determined from parent and slug.',
        unique=True,
    )

    slug = dict(
        blank=True,  # actually not, but autogenerated anyway
        max_length=255,
        validators=[validate_slug],
        verbose_name='Slug',
        help_text=(
            'The slug will be included as part of the URL. '
            'The slug may only contain lower case letters (a–z), numbers (0–9) and dashes (-). '
            'If you leave the slug blank, it will be generated from the title. '
            'If you change the slug after publishing, remember to create any redirects you may need.'
        ),
    )

    title = dict(
        max_length=1023,
        verbose_name='Title',
        help_text='The title will be displayed at the top of the view.',
    )

    description = dict(
        verbose_name='Description',
        help_text=(
            'While the description is not usually displayed in the user interface, '
            'it is visible to search engines and RSS readers. If your item contains '
            'text content, a good practice is to make the description a short summary '
            'of the content.'
        ),
        blank=True,
        default='',
    )

    body = dict(
        blank=True,
        default='',
        verbose_name='Text content',
        help_text='Will be displayed at the top of the album view before subalbums and pictures.',
    )

    order = dict(
        default=0,
        verbose_name='Ordering',
        help_text=(
            'The child items of a common parent are ordered in an ascending order by this number '
            '(smallest first).'
        )
    )

    is_public = dict(
        default=True,
        verbose_name='Public',
        help_text='Items that are not public are only accessible to admin users.',
    )

    is_visible = dict(
        default=True,
        verbose_name='Visible',
        help_text='Items that are not visible are not displayed in listings, but they can still be accessed via exact URL.',
    )
<|endoftext|>


<commit_before>import os
import logging
from threading import Semaphore, Thread
from unittest.mock import patch, Mock

import git

from app.core.logger import configure_logging
from app.core.messaging import Receiver
from app.core.servicemanager import ServiceManager
from app.services.updater.service import UpdaterService, UpdateScanner

configure_logging()

def make_receiver(count, obj, sem, r):
    def on_message(msg):
        obj.update(msg)
        sem.release()
        nonlocal count
        count -= 1
        if not count:
            r.stop()
    return on_message

class TestUpdateScanner(object):
    def setup_method(self):
        self.git_repo = git.Repo
        git.Repo = Mock()
        UpdateScanner.get_repos = lambda _: ["dir"]
        self.update_check_freq_backup = UpdateScanner.UPDATE_CHECK_FREQ
        UpdateScanner.UPDATE_CHECK_FREQ = 5

        os.environ["USE_FAKE_REDIS"] = "TRUE"
        self.service_manager = ServiceManager(None)
        self.service_manager.start_services(["messaging"])

    def teardown_method(self):
        git.Repo = self.git_repo
        del os.environ["USE_FAKE_REDIS"]
        self.service_manager.stop()

    def test_update(self):
        self.service = UpdaterService(None)
        self.service.on_service_start()

        obj = {}
        sem = Semaphore(0)

        receiver = Receiver("/services/shell/notifications")
        receiver.on_message = make_receiver(1, obj, sem, receiver)
        receiver.start()
        Thread(target=receiver.run).start()

        assert not sem.acquire(timeout=2)
        repo.asssert_called_with("dir")

<commit_msg>Bug fixes.<commit_after>import os
import logging
from threading import Semaphore, Thread
from unittest.mock import patch, Mock

import git

from app.core.logger import configure_logging
from app.core.messaging import Receiver
from app.core.servicemanager import ServiceManager
from app.services.updater.service import UpdaterService, UpdateScanner

configure_logging()

def make_receiver(count, obj, sem, r):
    def on_message(msg):
        obj.update(msg)
        sem.release()
        nonlocal count
        count -= 1
        if not count:
            r.stop()
    return on_message

class TestUpdateScanner(object):
    def setup_method(self):
        self.git_repo = git.Repo
        UpdateScanner.list_repos = lambda x, y: ["dir"]
        UpdateScanner.get_repo = lambda x, y: Mock()
        self.update_check_freq_backup = UpdateScanner.UPDATE_CHECK_FREQ
        UpdateScanner.UPDATE_CHECK_FREQ = 5

        os.environ["USE_FAKE_REDIS"] = "TRUE"
        self.service_manager = ServiceManager(None)
        self.service_manager.start_services(["messaging"])

    def teardown_method(self):
        git.Repo = self.git_repo
        del os.environ["USE_FAKE_REDIS"]
        self.service_manager.stop()

    def test_update(self):
        self.service = UpdaterService(None)
        self.service.on_service_start()

        obj = {}
        sem = Semaphore(0)

        receiver = Receiver("/services/shell/notifications")
        receiver.on_message = make_receiver(1, obj, sem, receiver)
        receiver.start()
        Thread(target=receiver.run).start()

        assert not sem.acquire(timeout=2)
        repo.asssert_called_with("dir")

<|endoftext|>


<commit_before>from utils import theyworkforyou
import json

class TWFYHansard:
    def __init__(self, link=False):
        self.twfy = theyworkforyou.TWFY("ESHFQ9AQWgTuFyCfx7DAWpDc")

    def get_mps(self):
        return json.loads(
            self.twfy.api.getMPs(output="js")
        )

    def get_mp_details(self, person_id):
        return json.loads(
            self.twfy.api.getMP(output="js", id=person_id)
        )

    def get_lords(self):
        return json.loads(
            self.twfy.api.getLords(output="js")
        )

    def get_lord_details(self, person_id):
        return json.loads(
            self.twfy.api.getLord(output="js", id=person_id)
        )

    def get_mp_debates(self, debate_type, person):
        return json.loads(
            self.twfy.api.getDebates(output="js", type=debate_type, person=person)
        )

    def get_debate(self, debate_id):
        debate_found = self._get_debate_detail(
            "commons",
            debate_id
        )
        if debate_found:
            starting_point = debate_found[0]
            mp_comment = debate_found[-1]
            start = {
                    "date": starting_point["hdate"],
                    "content_count": starting_point["contentcount"],
                    "body": starting_point["body"],
                    "debate_id": starting_point["gid"]
            }
            mp_comment = {
                    "date": mp_comment["hdate"],
                    "body": mp_comment["body"],
                    "debate_id": mp_comment["gid"]
            }
            if len(debate_found) == 3:
                sub = debate_found[1]
                sub_cat = {
                    "date": sub["hdate"],
                    "body": sub["body"],
                    "content_count": sub["contentcount"],
                    "debate_id": sub["gid"]
                }
            else:
                sub_cat = None
            return start, sub_cat, mp_comment
        else:
            return None

    def get_full_debate(self, full_debate_id, sub_category=True):
        debate_full = self._get_debate_detail(
            "commons",
            full_debate_id
        )
        #print "\n\n[o] Getting Full Debate Details [id: %s count:%s]" \
        #      % (full_debate_id, len(debate_full))
        main = {
            "debate_id": debate_full[0]["gid"],
            "date": debate_full[0]["hdate"],
            "topic": debate_full[0]["body"]
        }
        if sub_category:
            sub_cat = {
                "debate_id": debate_full[1]["gid"],
                "date": debate_full[1]["hdate"],
                "topic": debate_full[1]["body"]
            }
            full_debate = debate_full[2:]
        else:
            sub_cat = None
            full_debate = debate_full[1:]
        return main, sub_cat, full_debate

    def _get_debate_detail(self, debate_type, debate_id):
        try:
            return json.loads(
                self.twfy.api.getDebates(output="js", type=debate_type, gid=debate_id)
            )
        except json.scanner.JSONDecodeError:
            print "JSONDecodeError"
            return None
<commit_msg>Send kwargs on to TWFY API wrapper<commit_after>from utils import theyworkforyou
import json

class TWFYHansard:
    def __init__(self, link=False):
        self.twfy = theyworkforyou.TWFY("ESHFQ9AQWgTuFyCfx7DAWpDc")

    def get_mps(self, **kwargs):
        return json.loads(
            self.twfy.api.getMPs(output="js", **kwargs)
        )

    def get_mp_details(self, person_id, **kwargs):
        return json.loads(
            self.twfy.api.getMP(output="js", id=person_id, **kwargs)
        )

    def get_mp_info(self, person_id, **kwargs):
        return json.loads(
            self.twfy.api.getMPInfo(output="js", id=person_id, **kwargs)
        )

    def get_lords(self, **kwargs):
        return json.loads(
            self.twfy.api.getLords(output="js", **kwargs)
        )

    def get_lord_details(self, person_id, **kwargs):
        return json.loads(
            self.twfy.api.getLord(output="js", id=person_id, **kwargs)
        )

    def get_mp_debates(self, debate_type, person, **kwargs):
        return json.loads(
            self.twfy.api.getDebates(output="js", type=debate_type, person=person, **kwargs)
        )

    def get_debate(self, debate_id):
        debate_found = self._get_debate_detail(
            "commons",
            debate_id
        )
        if debate_found:
            starting_point = debate_found[0]
            mp_comment = debate_found[-1]
            start = {
                    "date": starting_point["hdate"],
                    "content_count": starting_point["contentcount"],
                    "body": starting_point["body"],
                    "debate_id": starting_point["gid"]
            }
            mp_comment = {
                    "date": mp_comment["hdate"],
                    "body": mp_comment["body"],
                    "debate_id": mp_comment["gid"]
            }
            if len(debate_found) == 3:
                sub = debate_found[1]
                sub_cat = {
                    "date": sub["hdate"],
                    "body": sub["body"],
                    "content_count": sub["contentcount"],
                    "debate_id": sub["gid"]
                }
            else:
                sub_cat = None
            return start, sub_cat, mp_comment
        else:
            return None

    def get_full_debate(self, full_debate_id, sub_category=True):
        debate_full = self._get_debate_detail(
            "commons",
            full_debate_id
        )
        #print "\n\n[o] Getting Full Debate Details [id: %s count:%s]" \
        #      % (full_debate_id, len(debate_full))
        main = {
            "debate_id": debate_full[0]["gid"],
            "date": debate_full[0]["hdate"],
            "topic": debate_full[0]["body"]
        }
        if sub_category:
            sub_cat = {
                "debate_id": debate_full[1]["gid"],
                "date": debate_full[1]["hdate"],
                "topic": debate_full[1]["body"]
            }
            full_debate = debate_full[2:]
        else:
            sub_cat = None
            full_debate = debate_full[1:]
        return main, sub_cat, full_debate

    def _get_debate_detail(self, debate_type, debate_id, **kwargs):
        try:
            return json.loads(
                self.twfy.api.getDebates(output="js", type=debate_type, gid=debate_id, **kwargs)
            )
        except json.scanner.JSONDecodeError:
            print "JSONDecodeError"
            return None
<|endoftext|>


<commit_before>from collections import OrderedDict
from sqlalchemy import or_, and_, func
from sqlalchemy.orm import outerjoin
from sqlalchemy.sql.expression import case

from dataactcore.models.domainModels import SF133, TASLookup, CGAC, FREC

gtas_model = SF133
tas_model = TASLookup

mapping = OrderedDict([
    ('allocation_transfer_agency', ['AllocationTransferAgencyIdentifier']),
    ('agency_identifier', ['AgencyIdentifier']),
    ('beginning_period_of_availa', ['BeginningPeriodOfAvailability']),
    ('ending_period_of_availabil', ['EndingPeriodOfAvailability']),
    ('availability_type_code', ['AvailabilityTypeCode']),
    ('main_account_code', ['MainAccountCode']),
    ('sub_account_code', ['SubAccountCode']),
    ('total_budgetary_resources_cpe', ['TotalBudgetaryResources_CPE']),
    ('budget_authority_appropria_cpe', ['BudgetAuthorityAppropriatedAmount_CPE']),
    ('budget_authority_unobligat_fyb', ['BudgetAuthorityUnobligatedBalanceBroughtForward_FYB']),
    ('adjustments_to_unobligated_cpe', ['AdjustmentsToUnobligatedBalanceBroughtForward_CPE']),
    ('other_budgetary_resources_cpe', ['OtherBudgetaryResourcesAmount_CPE']),
    ('contract_authority_amount_cpe', ['ContractAuthorityAmountTotal_CPE']),
    ('borrowing_authority_amount_cpe', ['BorrowingAuthorityAmountTotal_CPE']),
    ('spending_authority_from_of_cpe', ['SpendingAuthorityfromOffsettingCollectionsAmountTotal_CPE']),
    ('status_of_budgetary_resour_cpe', ['StatusOfBudgetaryResourcesTotal_CPE']),
    ('obligations_incurred_total_cpe', ['ObligationsIncurredTotalByTAS_CPE']),
    ('gross_outlay_amount_by_tas_cpe', ['GrossOutlayAmountByTAS_CPE']),
    ('unobligated_balance_cpe', ['UnobligatedBalance_CPE']),
    ('deobligations_recoveries_r_cpe', ['DeobligationsRecoveriesRefundsByTAS_CPE'])
])
db_columns = [key for key in mapping]

def query_data(session, agency_code, period, year):
    """ Request A file data

        Args:
            session: DB session
            agency_code: FREC or CGAC code for generation
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            The rows using the provided dates for the given agency.
    """
    # set a boolean to determine if the original agency code is frec or cgac
    frec_provided = len(agency_code) == 4
    tas_gtas = tas_gtas_combo(session, period, year)
    # Make a list of FRECs to compare to for 011 AID entries
    frec_list = []
    if not frec_provided:
        frec_list = session.query(FREC.frec_code).select_from(outerjoin(CGAC, FREC, CGAC.cgac_id == FREC.cgac_id)).\
            filter(CGAC.cgac_code == agency_code).all()
    # Group agencies together that need to be grouped
    agency_array = []
    if agency_code == '097':
        agency_array = ['017', '021', '057', '097']
    elif agency_code == '1601':
        agency_array = ['1601', '016']
    elif agency_code == '1125':
        agency_array = ['1125', '011']

    # Save the ATA filter
    agency_filters = []
    if not agency_array:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency == agency_code)
    else:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency.in_(agency_array))

    # Save the AID filter
    if agency_code == '097' and not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier.in_(agency_array)))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == agency_code))
    else:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.fr_entity_type == agency_code))

    # If we're checking a CGAC, we want to filter on all of the related FRECs for AID 011, otherwise just filter on
    # that FREC
    if frec_list:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == '011',
                                   tas_gtas.c.fr_entity_type.in_(frec_list)))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == '011',
                                   tas_gtas.c.fr_entity_type == agency_code))

    rows = initial_query(session, tas_gtas.c, year).\
        filter(func.coalesce(tas_gtas.c.financial_indicator2, '') != 'F').\
        filter(or_(*agency_filters)).\
        group_by(tas_gtas.c.allocation_transfer_agency,
                 tas_gtas.c.agency_identifier,
                 tas_gtas.c.beginning_period_of_availa,
                 tas_gtas.c.ending_period_of_availabil,
                 tas_gtas.c.availability_type_code,
                 tas_gtas.c.main_account_code,
                 tas_gtas.c.sub_account_code)

    return rows

def tas_gtas_combo(session, period, year):
    """ Creates a combined list of TAS and GTAS data filtered by the given period/year

        Args:
            session: DB session
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            A WITH clause to use with other queries
    """
    query = session.query(
        gtas_model.allocation_transfer_agency.label('allocation_transfer_agency'),
        gtas_model.agency_identifier.label('agency_identifier'),
        gtas_model.beginning_period_of_availa.label('beginning_period_of_availa'),
        gtas_model.ending_period_of_availabil.label('ending_period_of_availabil'),
        gtas_model.availability_type_code.label('availability_type_code'),
        gtas_model.main_account_code.label('main_account_code'),
        gtas_model.sub_account_code.label('sub_account_code'),
        gtas_model.amount.label('amount'),
        gtas_model.line.label('line'),
        tas_model.financial_indicator2.label('financial_indicator2'),
        tas_model.fr_entity_type.label('fr_entity_type')).\
        join(tas_model, gtas_model.tas == func.concat(func.coalesce(tas_model.allocation_transfer_agency, '000'),
                                                      func.coalesce(tas_model.agency_identifier, '000'),
                                                      func.coalesce(tas_model.beginning_period_of_availa, '0000'),
                                                      func.coalesce(tas_model.ending_period_of_availabil, '0000'),
                                                      func.coalesce(tas_model.availability_type_code, ' '),
                                                      func.coalesce(tas_model.main_account_code, '0000'),
                                                      func.coalesce(tas_model.sub_account_code, '000'))).\
        filter(gtas_model.period == period).\
        filter(gtas_model.fiscal_year == year)
    return query.cte('tas_gtas')

def initial_query(session, model, year):
    """ Creates the initial query for D2 files.

        Args:
            session: The current DB session
            model: subquery model to get data from
            year: the year for which to get data

        Returns:
            The base query (a select from the PublishedAwardFinancialAssistance table with the specified columns).
    """
    budget_authority_line_max = 1042 if year <= 2020 else 1065
    return session.query(
        model.allocation_transfer_agency,
        model.agency_identifier,
        model.beginning_period_of_availa,
        model.ending_period_of_availabil,
        model.availability_type_code,
        model.main_account_code,
        model.sub_account_code,
        func.sum(case([(model.line == 1910, model.amount)], else_=0)).label('total_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1160, 1180, 1260, 1280]), model.amount)],
                      else_=0)).label('budget_authority_appropria_cpe'),
        func.sum(case([(model.line == 1000, model.amount)], else_=0)).label('budget_authority_unobligat_fyb'),
        func.sum(case([(and_(model.line >= 1010, model.line <= budget_authority_line_max), model.amount)],
                      else_=0)).label('adjustments_to_unobligated_cpe'),
        func.sum(case([(model.line.in_([1540, 1640, 1340, 1440, 1750, 1850]), model.amount)],
                      else_=0)).label('other_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1540, 1640]), model.amount)], else_=0)).label('contract_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1340, 1440]), model.amount)], else_=0)).label('borrowing_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1750, 1850]), model.amount)], else_=0)).label('spending_authority_from_of_cpe'),
        func.sum(case([(model.line == 2500, model.amount)], else_=0)).label('status_of_budgetary_resour_cpe'),
        func.sum(case([(model.line == 2190, model.amount)], else_=0)).label('obligations_incurred_total_cpe'),
        func.sum(case([(model.line == 3020, model.amount)], else_=0)).label('gross_outlay_amount_by_tas_cpe'),
        func.sum(case([(model.line == 2490, model.amount)], else_=0)).label('unobligated_balance_cpe'),
        func.sum(case([(model.line.in_([1021, 1033]), model.amount)], else_=0)).label('deobligations_recoveries_r_cpe'))
<commit_msg>Fixing the actual FREC list<commit_after>from collections import OrderedDict
from sqlalchemy import or_, and_, func
from sqlalchemy.orm import outerjoin
from sqlalchemy.sql.expression import case

from dataactcore.models.domainModels import SF133, TASLookup, CGAC, FREC

gtas_model = SF133
tas_model = TASLookup

mapping = OrderedDict([
    ('allocation_transfer_agency', ['AllocationTransferAgencyIdentifier']),
    ('agency_identifier', ['AgencyIdentifier']),
    ('beginning_period_of_availa', ['BeginningPeriodOfAvailability']),
    ('ending_period_of_availabil', ['EndingPeriodOfAvailability']),
    ('availability_type_code', ['AvailabilityTypeCode']),
    ('main_account_code', ['MainAccountCode']),
    ('sub_account_code', ['SubAccountCode']),
    ('total_budgetary_resources_cpe', ['TotalBudgetaryResources_CPE']),
    ('budget_authority_appropria_cpe', ['BudgetAuthorityAppropriatedAmount_CPE']),
    ('budget_authority_unobligat_fyb', ['BudgetAuthorityUnobligatedBalanceBroughtForward_FYB']),
    ('adjustments_to_unobligated_cpe', ['AdjustmentsToUnobligatedBalanceBroughtForward_CPE']),
    ('other_budgetary_resources_cpe', ['OtherBudgetaryResourcesAmount_CPE']),
    ('contract_authority_amount_cpe', ['ContractAuthorityAmountTotal_CPE']),
    ('borrowing_authority_amount_cpe', ['BorrowingAuthorityAmountTotal_CPE']),
    ('spending_authority_from_of_cpe', ['SpendingAuthorityfromOffsettingCollectionsAmountTotal_CPE']),
    ('status_of_budgetary_resour_cpe', ['StatusOfBudgetaryResourcesTotal_CPE']),
    ('obligations_incurred_total_cpe', ['ObligationsIncurredTotalByTAS_CPE']),
    ('gross_outlay_amount_by_tas_cpe', ['GrossOutlayAmountByTAS_CPE']),
    ('unobligated_balance_cpe', ['UnobligatedBalance_CPE']),
    ('deobligations_recoveries_r_cpe', ['DeobligationsRecoveriesRefundsByTAS_CPE'])
])
db_columns = [key for key in mapping]

def query_data(session, agency_code, period, year):
    """ Request A file data

        Args:
            session: DB session
            agency_code: FREC or CGAC code for generation
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            The rows using the provided dates for the given agency.
    """
    # set a boolean to determine if the original agency code is frec or cgac
    frec_provided = len(agency_code) == 4
    tas_gtas = tas_gtas_combo(session, period, year)
    # Make a list of FRECs to compare to for 011 AID entries
    frec_list = []
    if not frec_provided:
        frec_list = session.query(FREC.frec_code).select_from(outerjoin(CGAC, FREC, CGAC.cgac_id == FREC.cgac_id)).\
            filter(CGAC.cgac_code == agency_code).all()
        # Put the frec list in a format that can be read by a filter
        frec_list = [frec.frec_code for frec in frec_list]
    # Group agencies together that need to be grouped
    agency_array = []
    if agency_code == '097':
        agency_array = ['017', '021', '057', '097']
    elif agency_code == '1601':
        agency_array = ['1601', '016']
    elif agency_code == '1125':
        agency_array = ['1125', '011']

    # Save the ATA filter
    agency_filters = []
    if not agency_array:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency == agency_code)
    else:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency.in_(agency_array))

    # Save the AID filter
    if agency_code == '097' and not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier.in_(agency_array)))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == agency_code))
    else:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.fr_entity_type == agency_code))

    # If we're checking a CGAC, we want to filter on all of the related FRECs for AID 011, otherwise just filter on
    # that FREC
    if frec_list:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == '011',
                                   tas_gtas.c.fr_entity_type.in_(frec_list)))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == '011',
                                   tas_gtas.c.fr_entity_type == agency_code))

    rows = initial_query(session, tas_gtas.c, year).\
        filter(func.coalesce(tas_gtas.c.financial_indicator2, '') != 'F').\
        filter(or_(*agency_filters)).\
        group_by(tas_gtas.c.allocation_transfer_agency,
                 tas_gtas.c.agency_identifier,
                 tas_gtas.c.beginning_period_of_availa,
                 tas_gtas.c.ending_period_of_availabil,
                 tas_gtas.c.availability_type_code,
                 tas_gtas.c.main_account_code,
                 tas_gtas.c.sub_account_code)

    return rows

def tas_gtas_combo(session, period, year):
    """ Creates a combined list of TAS and GTAS data filtered by the given period/year

        Args:
            session: DB session
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            A WITH clause to use with other queries
    """
    query = session.query(
        gtas_model.allocation_transfer_agency.label('allocation_transfer_agency'),
        gtas_model.agency_identifier.label('agency_identifier'),
        gtas_model.beginning_period_of_availa.label('beginning_period_of_availa'),
        gtas_model.ending_period_of_availabil.label('ending_period_of_availabil'),
        gtas_model.availability_type_code.label('availability_type_code'),
        gtas_model.main_account_code.label('main_account_code'),
        gtas_model.sub_account_code.label('sub_account_code'),
        gtas_model.amount.label('amount'),
        gtas_model.line.label('line'),
        tas_model.financial_indicator2.label('financial_indicator2'),
        tas_model.fr_entity_type.label('fr_entity_type')).\
        join(tas_model, gtas_model.tas == func.concat(func.coalesce(tas_model.allocation_transfer_agency, '000'),
                                                      func.coalesce(tas_model.agency_identifier, '000'),
                                                      func.coalesce(tas_model.beginning_period_of_availa, '0000'),
                                                      func.coalesce(tas_model.ending_period_of_availabil, '0000'),
                                                      func.coalesce(tas_model.availability_type_code, ' '),
                                                      func.coalesce(tas_model.main_account_code, '0000'),
                                                      func.coalesce(tas_model.sub_account_code, '000'))).\
        filter(gtas_model.period == period).\
        filter(gtas_model.fiscal_year == year)
    return query.cte('tas_gtas')

def initial_query(session, model, year):
    """ Creates the initial query for D2 files.

        Args:
            session: The current DB session
            model: subquery model to get data from
            year: the year for which to get data

        Returns:
            The base query (a select from the PublishedAwardFinancialAssistance table with the specified columns).
    """
    budget_authority_line_max = 1042 if year <= 2020 else 1065
    return session.query(
        model.allocation_transfer_agency,
        model.agency_identifier,
        model.beginning_period_of_availa,
        model.ending_period_of_availabil,
        model.availability_type_code,
        model.main_account_code,
        model.sub_account_code,
        func.sum(case([(model.line == 1910, model.amount)], else_=0)).label('total_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1160, 1180, 1260, 1280]), model.amount)],
                      else_=0)).label('budget_authority_appropria_cpe'),
        func.sum(case([(model.line == 1000, model.amount)], else_=0)).label('budget_authority_unobligat_fyb'),
        func.sum(case([(and_(model.line >= 1010, model.line <= budget_authority_line_max), model.amount)],
                      else_=0)).label('adjustments_to_unobligated_cpe'),
        func.sum(case([(model.line.in_([1540, 1640, 1340, 1440, 1750, 1850]), model.amount)],
                      else_=0)).label('other_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1540, 1640]), model.amount)], else_=0)).label('contract_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1340, 1440]), model.amount)], else_=0)).label('borrowing_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1750, 1850]), model.amount)], else_=0)).label('spending_authority_from_of_cpe'),
        func.sum(case([(model.line == 2500, model.amount)], else_=0)).label('status_of_budgetary_resour_cpe'),
        func.sum(case([(model.line == 2190, model.amount)], else_=0)).label('obligations_incurred_total_cpe'),
        func.sum(case([(model.line == 3020, model.amount)], else_=0)).label('gross_outlay_amount_by_tas_cpe'),
        func.sum(case([(model.line == 2490, model.amount)], else_=0)).label('unobligated_balance_cpe'),
        func.sum(case([(model.line.in_([1021, 1033]), model.amount)], else_=0)).label('deobligations_recoveries_r_cpe'))
<|endoftext|>