<commit_before># Invoked by: SNS Subscription
# Returns: Error or status message
#
# Triggered by messages sent to various SNS topics resulting from CloudWatch
# Alarms being triggered or changing states.

import boto3
import os
import json
import urllib.parse
from dateutil.parser import parse
import datetime
import re

sns = boto3.client('sns')
cloudwatch = boto3.client('cloudwatch')

SLACK_ICON = ':ops-cloudwatch-alarm:'
SLACK_USERNAME = 'Amazon CloudWatch Alarms'

def channel_for_topic(sns_topic_arn):
    if 'OpsFatal' in sns_topic_arn:
        return '#ops-fatal'
    elif 'OpsError' in sns_topic_arn:
        return '#ops-error'
    elif 'OpsWarn' in sns_topic_arn:
        return '#ops-warn'
    elif 'OpsInfo' in sns_topic_arn:
        return '#ops-info'
    else:
        return '#ops-debug'

def color_for_alarm(alarm):
    if alarm['NewStateValue'] == 'ALARM':
        return '#cc0000'
    elif alarm['NewStateValue'] == 'OK':
        return '#019933'
    else:
        return '#e07701'

def alarm_slack_attachment(alarm):
    # Extract datapoint values from the SNS alarm data
    # eg, "Threshold Crossed: 1 datapoint [10.0 (05/09/18 12:15:00)] was
    # greater than or equal to the threshold (10.0)."
    datapoints = re.findall(r'([0-9]+\.[0-9]+) ', alarm['NewStateReason'])
    datapoints_list = '`\n`'.join(datapoints)

    trigger = alarm['Trigger']

    # Get the complete alarm info for this alarm (Only partial data may have
    # been included in the SNS message)
    alarm_infos = cloudwatch.describe_alarms(AlarmNames=[alarm['AlarmName']])
    alarm_info = alarm_infos['MetricAlarms'][0]

    alarm_region = alarm_info['AlarmArn'].split(':', 4)[3]

    # If the alarm doesn't have DatapointsToAlarm defined, force it equal to
    # the EvaluationPeriods
    if 'DatapointsToAlarm' not in alarm_info:
        alarm_info['DatapointsToAlarm'] = alarm_info['EvaluationPeriods']

    # Get a count of how many times this alarm has alarmed in the last 24 hours
    now = datetime.datetime.now(datetime.timezone.utc)
    alarm_history = cloudwatch.describe_alarm_history(
        AlarmName=alarm['AlarmName'],
        HistoryItemType='StateUpdate',
        StartDate=now - datetime.timedelta(hours=24),
        EndDate=now,
        MaxRecords=100,
    )
    if 'AlarmHistoryItems' in alarm_history:
        items = alarm_history['AlarmHistoryItems']
        alarms = filter(lambda x: ('to ALARM' in x['HistorySummary']), items)
        alarms_count = len(list(alarms))
    else:
        alarms_count = 0

    # Construct a URL for this alarm in the Console
    cw_console_url = 'https://console.aws.amazon.com/cloudwatch/home'
    alarm_name_escaped = urllib.parse.quote(alarm['AlarmName'])
    alarm_console_url = f"{cw_console_url}?region={alarm_region}#alarm:alarmFilter=ANY;name={alarm_name_escaped}"

    cw_logs = 'n/a'

    # All periods are 10, 30, or a multiple of 60
    # Each datapoint is the aggregate (SUM, AVERAGE, etc) of one period
    if trigger['Period'] >= 60:
        each_datapoint = f"{round(trigger['Period'] / 60)} minute"
    else:
        each_datapoint = f"{trigger['Period']} second"

    # ExtendedStatistic is used for percentile statistics.
    # ExtendedStatistic and Statistic are mutually exclusive.
    if 'ExtendedStatistic' in trigger:
        stat = trigger['ExtendedStatistic']
    else:
        stat = trigger['Statistic']

    # eg "5 minute TargetResponseTime average"
    threshold_left = f"{each_datapoint} `{trigger['MetricName']}` {stat.lower()}"

    trigger_period = trigger['Period'] * trigger['EvaluationPeriods']
    trigger_period_label = "seconds"

    if trigger_period >= 60:
        trigger_period = round(trigger_period / 60)
        trigger_period_label = "minutes"

    if trigger['EvaluationPeriods'] == 1:
        # Entire threshold breach was a single period/datapoint
        threshold_right = ""
    elif alarm_info['DatapointsToAlarm'] == alarm_info['EvaluationPeriods']:
        # Threshold breach was multiple, consecutive periods
        threshold_right = f"for {trigger_period} consecutive {trigger_period_label}"
    else:
        # Threshold breach was "M of N" periods
        threshold_right = f"at least {alarm_info['DatapointsToAlarm']} times in {trigger_period} {trigger_period_label}"

    if trigger['ComparisonOperator'] == 'GreaterThanOrEqualToThreshold':
        comparison = '≥'
    elif trigger['ComparisonOperator'] == 'GreaterThanThreshold':
        comparison = '>'
    elif trigger['ComparisonOperator'] == 'LessThanThreshold':
        comparison = '<'
    elif trigger['ComparisonOperator'] == 'LessThanOrEqualToThreshold':
        comparison = '≤'

    threshold = f"{threshold_left} *{comparison}* `{trigger['Threshold']}` {threshold_right}"

    # Log URL handling
    # eg, alarm['StateChangeTime'] = 2019-08-03T01:46:44.418+0000
    state_change_time = parse(alarm['StateChangeTime'])
    log_end = state_change_time.strftime("%Y-%m-%dT%H:%M:%S")

    log_period = trigger['Period'] * trigger['EvaluationPeriods']

    log_start_time = state_change_time - datetime.timedelta(seconds=log_period)
    log_start = log_start_time.strftime("%Y-%m-%dT%H:%M:%S")

    if trigger['Namespace'] == 'AWS/Lambda':
        # Find the function name
        for dimension in trigger['Dimensions']:
            if dimension['name'] == 'FunctionName':
                function_name = dimension['value']

        cw_logs = f"<https://console.aws.amazon.com/cloudwatch/home?region={alarm_region}#logEventViewer:group=/aws/lambda/{function_name};start={log_start}Z;end={log_end}Z|CloudWatch Logs>"

    return {
        'color': color_for_alarm(alarm),
        'fallback': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title_link': alarm_console_url,
        'title': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'text': f"{alarm['AlarmDescription']}",
        'footer': f"{trigger['Namespace']} – {alarm['Region']}",
        'ts': round(parse(alarm['StateChangeTime']).timestamp()),
        'fields': [
            {
                'title': 'Last 24 Hours',
                'value': f"{alarms_count} alarms",
                'short': True,
            },
            {
                'title': 'Logs',
                'value': cw_logs,
                'short': True,
            },
            {
                'title': 'Threshold breach',
                'value': threshold,
                'short': False,
            },
            {
                'title': 'Datapoints',
                'value': f"`{datapoints_list}`",
                'short': False,
            }
        ]
    }

def ok_slack_attachment(alarm):
    duration = 'Unavailable'

    alarm_infos = cloudwatch.describe_alarms(AlarmNames=[alarm['AlarmName']])
    alarm_info = alarm_infos['MetricAlarms'][0]

    alarm_region = alarm_info['AlarmArn'].split(':', 4)[3]

    now = datetime.datetime.now(datetime.timezone.utc)
    alarm_history = cloudwatch.describe_alarm_history(
        AlarmName=alarm['AlarmName'],
        HistoryItemType='StateUpdate',
        StartDate=now - datetime.timedelta(hours=24),
        EndDate=now,
    )
    items = alarm_history['AlarmHistoryItems']

    if len(items) > 0:
        ok_item = items[0]
        history_data = json.loads(ok_item['HistoryData'])

        ok_time = ok_item['Timestamp']

        if 'oldState' in history_data:
            alarm_time = history_data['oldState']['stateReasonData']['startDate']
            alarm_time = parse(alarm_time)
        else:
            alarm_time = ok_time

        dif = ok_time - alarm_time
        duration = f"{round(dif.total_seconds() / 60)} min."

        cw_console_url = 'https://console.aws.amazon.com/cloudwatch/home'
        alarm_name_escaped = urllib.parse.quote(alarm['AlarmName'])
        alarm_console_url = f"{cw_console_url}?region={alarm_region}#alarm:alarmFilter=ANY;name={alarm_name_escaped}"

    return {
        'color': color_for_alarm(alarm),
        'fallback': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title_link': alarm_console_url,
        'text': f"Alarm duration: {duration}",
        'footer': f"{alarm['Trigger']['Namespace']} – {alarm['Region']}",
        'ts': round(parse(alarm['StateChangeTime']).timestamp())
    }

def slack_message(sns_payload):
    alarm = json.loads(sns_payload['Message'])

    if alarm['NewStateValue'] == 'OK':
        attachment = ok_slack_attachment(alarm)
    else:
        attachment = alarm_slack_attachment(alarm)

    return {
        'channel': channel_for_topic(sns_payload['TopicArn']),
        'username': SLACK_USERNAME,
        'icon_emoji': SLACK_ICON,
        'attachments': [attachment]
    }

def lambda_handler(event, context):
    sns_payload = event['Records'][0]['Sns']

    # TODO This isn't necessary once all stacks stop using the log-level SNS
    # topics for deployment notifications
    if sns_payload['Subject'] == 'AWS CloudFormation Notification':
        return

    sns.publish(
        TopicArn=os.environ['SLACK_MESSAGE_RELAY_TOPIC_ARN'],
        Message=json.dumps(slack_message(sns_payload))
    )
<commit_msg>Fix error when evaluating duration of newly created CloudWatch Alarms<commit_after># Invoked by: SNS Subscription
# Returns: Error or status message
#
# Triggered by messages sent to various SNS topics resulting from CloudWatch
# Alarms being triggered or changing states.

import boto3
import os
import json
import urllib.parse
from dateutil.parser import parse
import datetime
import re

sns = boto3.client('sns')
cloudwatch = boto3.client('cloudwatch')

SLACK_ICON = ':ops-cloudwatch-alarm:'
SLACK_USERNAME = 'Amazon CloudWatch Alarms'

def channel_for_topic(sns_topic_arn):
    if 'OpsFatal' in sns_topic_arn:
        return '#ops-fatal'
    elif 'OpsError' in sns_topic_arn:
        return '#ops-error'
    elif 'OpsWarn' in sns_topic_arn:
        return '#ops-warn'
    elif 'OpsInfo' in sns_topic_arn:
        return '#ops-info'
    else:
        return '#ops-debug'

def color_for_alarm(alarm):
    if alarm['NewStateValue'] == 'ALARM':
        return '#cc0000'
    elif alarm['NewStateValue'] == 'OK':
        return '#019933'
    else:
        return '#e07701'

def alarm_slack_attachment(alarm):
    # Extract datapoint values from the SNS alarm data
    # eg, "Threshold Crossed: 1 datapoint [10.0 (05/09/18 12:15:00)] was
    # greater than or equal to the threshold (10.0)."
    datapoints = re.findall(r'([0-9]+\.[0-9]+) ', alarm['NewStateReason'])
    datapoints_list = '`\n`'.join(datapoints)

    trigger = alarm['Trigger']

    # Get the complete alarm info for this alarm (Only partial data may have
    # been included in the SNS message)
    alarm_infos = cloudwatch.describe_alarms(AlarmNames=[alarm['AlarmName']])
    alarm_info = alarm_infos['MetricAlarms'][0]

    alarm_region = alarm_info['AlarmArn'].split(':', 4)[3]

    # If the alarm doesn't have DatapointsToAlarm defined, force it equal to
    # the EvaluationPeriods
    if 'DatapointsToAlarm' not in alarm_info:
        alarm_info['DatapointsToAlarm'] = alarm_info['EvaluationPeriods']

    # Get a count of how many times this alarm has alarmed in the last 24 hours
    now = datetime.datetime.now(datetime.timezone.utc)
    alarm_history = cloudwatch.describe_alarm_history(
        AlarmName=alarm['AlarmName'],
        HistoryItemType='StateUpdate',
        StartDate=now - datetime.timedelta(hours=24),
        EndDate=now,
        MaxRecords=100,
    )
    if 'AlarmHistoryItems' in alarm_history:
        items = alarm_history['AlarmHistoryItems']
        alarms = filter(lambda x: ('to ALARM' in x['HistorySummary']), items)
        alarms_count = len(list(alarms))
    else:
        alarms_count = 0

    # Construct a URL for this alarm in the Console
    cw_console_url = 'https://console.aws.amazon.com/cloudwatch/home'
    alarm_name_escaped = urllib.parse.quote(alarm['AlarmName'])
    alarm_console_url = f"{cw_console_url}?region={alarm_region}#alarm:alarmFilter=ANY;name={alarm_name_escaped}"

    cw_logs = 'n/a'

    # All periods are 10, 30, or a multiple of 60
    # Each datapoint is the aggregate (SUM, AVERAGE, etc) of one period
    if trigger['Period'] >= 60:
        each_datapoint = f"{round(trigger['Period'] / 60)} minute"
    else:
        each_datapoint = f"{trigger['Period']} second"

    # ExtendedStatistic is used for percentile statistics.
    # ExtendedStatistic and Statistic are mutually exclusive.
    if 'ExtendedStatistic' in trigger:
        stat = trigger['ExtendedStatistic']
    else:
        stat = trigger['Statistic']

    # eg "5 minute TargetResponseTime average"
    threshold_left = f"{each_datapoint} `{trigger['MetricName']}` {stat.lower()}"

    trigger_period = trigger['Period'] * trigger['EvaluationPeriods']
    trigger_period_label = "seconds"

    if trigger_period >= 60:
        trigger_period = round(trigger_period / 60)
        trigger_period_label = "minutes"

    if trigger['EvaluationPeriods'] == 1:
        # Entire threshold breach was a single period/datapoint
        threshold_right = ""
    elif alarm_info['DatapointsToAlarm'] == alarm_info['EvaluationPeriods']:
        # Threshold breach was multiple, consecutive periods
        threshold_right = f"for {trigger_period} consecutive {trigger_period_label}"
    else:
        # Threshold breach was "M of N" periods
        threshold_right = f"at least {alarm_info['DatapointsToAlarm']} times in {trigger_period} {trigger_period_label}"

    if trigger['ComparisonOperator'] == 'GreaterThanOrEqualToThreshold':
        comparison = '≥'
    elif trigger['ComparisonOperator'] == 'GreaterThanThreshold':
        comparison = '>'
    elif trigger['ComparisonOperator'] == 'LessThanThreshold':
        comparison = '<'
    elif trigger['ComparisonOperator'] == 'LessThanOrEqualToThreshold':
        comparison = '≤'

    threshold = f"{threshold_left} *{comparison}* `{trigger['Threshold']}` {threshold_right}"

    # Log URL handling
    # eg, alarm['StateChangeTime'] = 2019-08-03T01:46:44.418+0000
    state_change_time = parse(alarm['StateChangeTime'])
    log_end = state_change_time.strftime("%Y-%m-%dT%H:%M:%S")

    log_period = trigger['Period'] * trigger['EvaluationPeriods']

    log_start_time = state_change_time - datetime.timedelta(seconds=log_period)
    log_start = log_start_time.strftime("%Y-%m-%dT%H:%M:%S")

    if trigger['Namespace'] == 'AWS/Lambda':
        # Find the function name
        for dimension in trigger['Dimensions']:
            if dimension['name'] == 'FunctionName':
                function_name = dimension['value']

        cw_logs = f"<https://console.aws.amazon.com/cloudwatch/home?region={alarm_region}#logEventViewer:group=/aws/lambda/{function_name};start={log_start}Z;end={log_end}Z|CloudWatch Logs>"

    return {
        'color': color_for_alarm(alarm),
        'fallback': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title_link': alarm_console_url,
        'title': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'text': f"{alarm['AlarmDescription']}",
        'footer': f"{trigger['Namespace']} – {alarm['Region']}",
        'ts': round(parse(alarm['StateChangeTime']).timestamp()),
        'fields': [
            {
                'title': 'Last 24 Hours',
                'value': f"{alarms_count} alarms",
                'short': True,
            },
            {
                'title': 'Logs',
                'value': cw_logs,
                'short': True,
            },
            {
                'title': 'Threshold breach',
                'value': threshold,
                'short': False,
            },
            {
                'title': 'Datapoints',
                'value': f"`{datapoints_list}`",
                'short': False,
            }
        ]
    }

def ok_slack_attachment(alarm):
    # Get the complete alarm info for this alarm (Only partial data may have
    # been included in the SNS message)
    alarm_infos = cloudwatch.describe_alarms(AlarmNames=[alarm['AlarmName']])
    alarm_info = alarm_infos['MetricAlarms'][0]

    alarm_region = alarm_info['AlarmArn'].split(':', 4)[3]

    # Calculate the duration of the previous alarm state. The previous
    # state may not exist or may not be an alarm, so this needs to fail
    # gracefully
    try:
        # Retrieve the alarm history (only state updates)
        now = datetime.datetime.now(datetime.timezone.utc)
        alarm_history = cloudwatch.describe_alarm_history(
            AlarmName=alarm['AlarmName'],
            HistoryItemType='StateUpdate',
            StartDate=now - datetime.timedelta(hours=24),
            EndDate=now,
        )
        items = alarm_history['AlarmHistoryItems']

        # Since a state change triggered this, this should always be > 0
        if len(items) > 0:
            # The NewStateValue of the alarm was OK, so the most recent history
            # item should be an OK state (the inciting OK state, in fact)
            ok_item = items[0]

            # See history_data.json
            history_data = json.loads(ok_item['HistoryData'])

            if history_data['oldState']['stateValue'] == 'ALARM':
                alarm_time = history_data['oldState']['stateReasonData']['startDate']
                alarm_time = parse(alarm_time)

                ok_time = ok_item['Timestamp']

                dif = ok_time - alarm_time
                duration = f"{round(dif.total_seconds() / 60)} min."
    except Exception:
        duration = 'Unavailable'

    cw_console_url = 'https://console.aws.amazon.com/cloudwatch/home'
    alarm_name_escaped = urllib.parse.quote(alarm['AlarmName'])
    alarm_console_url = f"{cw_console_url}?region={alarm_region}#alarm:alarmFilter=ANY;name={alarm_name_escaped}"

    return {
        'color': color_for_alarm(alarm),
        'fallback': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title': f"{alarm['NewStateValue']} – {alarm['AlarmName']}",
        'title_link': alarm_console_url,
        'text': f"Alarm duration: {duration}",
        'footer': f"{alarm['Trigger']['Namespace']} – {alarm['Region']}",
        'ts': round(parse(alarm['StateChangeTime']).timestamp())
    }

def slack_message(sns_payload):
    alarm = json.loads(sns_payload['Message'])

    if alarm['NewStateValue'] == 'OK':
        attachment = ok_slack_attachment(alarm)
    else:
        attachment = alarm_slack_attachment(alarm)

    return {
        'channel': channel_for_topic(sns_payload['TopicArn']),
        'username': SLACK_USERNAME,
        'icon_emoji': SLACK_ICON,
        'attachments': [attachment]
    }

def lambda_handler(event, context):
    sns_payload = event['Records'][0]['Sns']

    # TODO This isn't necessary once all stacks stop using the log-level SNS
    # topics for deployment notifications
    if sns_payload['Subject'] == 'AWS CloudFormation Notification':
        return

    sns.publish(
        TopicArn=os.environ['SLACK_MESSAGE_RELAY_TOPIC_ARN'],
        Message=json.dumps(slack_message(sns_payload))
    )
<|endoftext|>


<commit_before>__version__ = '0.5.2'
__prog__ = "aetros"

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'<commit_msg>bump version<commit_after>__version__ = '0.5.3'
__prog__ = "aetros"

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'<|endoftext|>


<commit_before>#
# This file is 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.

"""DataSetMgr entry class
This module provides the python entry point to DataSetMgr on scala side
"""

from smv.utils import smv_copy_array, scala_seq_to_list, list_distinct, infer_full_name_from_part
from smv.datasetresolver import DataSetResolver

class DataSetMgr(object):
    """The Python representation of DataSetMgr.
    """

    def __init__(self, _jvm, smvconfig):
        self._jvm = _jvm

        self.smvconfig = smvconfig
        self.dsRepoFactories = []

        from py4j.java_gateway import java_import
        java_import(self._jvm, "org.tresamigos.smv.python.SmvPythonHelper")
        java_import(self._jvm, "org.tresamigos.smv.DataSetRepoFactoryPython")

        self.helper = self._jvm.SmvPythonHelper

    def stages(self):
        return self.smvconfig.stage_names()

    def tx(self):
        """Create a TXContext for multiple places, avoid the long TXContext line
        """
        return TXContext(self._jvm, self.dsRepoFactories, self.stages())

    def load(self, *urns):
        """Load SmvGenericModules for specified URNs

        Args:
            *urns (str): list of URNs as strings

        Returns:
            list(SmvGenericModules): list of Scala SmvGenericModules (j_ds)
        """
        with self.tx() as tx:
            return tx.load(urns)

    def inferDS(self, *partial_names):
        """Return DSs from a list of partial names

        Args:
            *partial_names (str): list of partial names

        Returns:
            list(SmvGenericModules): list of SmvGenericModules
        """
        with self.tx() as tx:
            return tx.inferDS(partial_names)

    def inferUrn(self, partial_name):
        """Return URN string from partial name
        """
        return self.inferDS(partial_name)[0].urn()

    def register(self, repo_factory):
        """Register python repo factory
        """
        self.dsRepoFactories.append(repo_factory)

    def allDataSets(self):
        """Return all the SmvGenericModules in the app
        """
        with self.tx() as tx:
            return tx.allDataSets()

    def modulesToRun(self, modPartialNames, stageNames, allMods):
        """Return a modules need to run
            Combine specified modules, (-m), stages, (-s) and if
            (--run-app) specified, all output modules
        """
        with self.tx() as tx:
            named_mods = tx.inferDS(modPartialNames)
            stage_mods = tx.outputModulesForStage(stageNames)
            app_mods = tx.allOutputModules() if allMods else []
            res = []
            res.extend(named_mods)
            res.extend(stage_mods)
            res.extend(app_mods)
            # Need to perserve the ordering
            return list_distinct(res)

class TXContext(object):
    """Create a TX context for "with tx() as tx" syntax
    """
    def __init__(self, _jvm, resourceFactories, stages):
        self._jvm = _jvm
        self.resourceFactories = resourceFactories
        self.stages = stages

    def __enter__(self):
        return TX(self._jvm, self.resourceFactories, self.stages)

    def __exit__(self, type, value, traceback):
        pass

class TX(object):
    """Abstraction of the transaction boundary for loading SmvGenericModules.
        A TX object

        * will instantiate a set of repos when itself instantiated and will
        * reuse the same repos for all queries. This means that each new TX object will
        * reload the SmvGenericModules from source **once** during its lifetime.

        NOTE: Once a new TX is created, the well-formedness of the SmvGenericModules provided
        by the previous TX is not guaranteed. Particularly it may become impossible
        to run modules from the previous TX.
    """
    def __init__(self, _jvm, resourceFactories, stages):
        self.repos = [rf.createRepo() for rf in resourceFactories]
        self.stages = stages
        self.resolver = DataSetResolver(self.repos[0])
        self.log = _jvm.org.apache.log4j.LogManager.getLogger("smv")

    def load(self, urns):
        fqns = [u[4:] for u in urns]
        return self.resolver.loadDataSet(fqns)

    def inferDS(self, partial_names):
        return self.load(self._inferUrn(partial_names))

    def allDataSets(self):
        return self.load(self._allUrns())

    def allOutputModules(self):
        return self._filterOutput(self.allDataSets())

    def outputModulesForStage(self, stageNames):
        return self._filterOutput(self._dsForStage(stageNames))

    def _dsForStage(self, stageNames):
        return self.load(self._urnsForStage(stageNames))

    def _urnsForStage(self, stageNames):
        return [u
            for repo in self.repos
            for s in stageNames
            for u in repo._dataSetsForStage(s)
        ]

    def _allUrns(self):
        if (len(self.stages) == 0):
            log.warn("No stage names configured. Unable to discover any modules.")
        return self._urnsForStage(self.stages)

    def _inferUrn(self, partial_names):
        def urn_str(pn):
            return infer_full_name_from_part(
                self._allUrns(),
                pn
            )

        return [urn_str(pn) for pn in partial_names]

    def _filterOutput(self, dss):
        return [ds for ds in dss if ds.isSmvOutput()]<commit_msg>smv.datasetmgr.TX should also use smv.logger (#1512)<commit_after>#
# This file is 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.

"""DataSetMgr entry class
This module provides the python entry point to DataSetMgr on scala side
"""
import smv

from smv.utils import smv_copy_array, scala_seq_to_list, list_distinct, infer_full_name_from_part
from smv.datasetresolver import DataSetResolver

class DataSetMgr(object):
    """The Python representation of DataSetMgr.
    """

    def __init__(self, _jvm, smvconfig):
        self._jvm = _jvm

        self.smvconfig = smvconfig
        self.dsRepoFactories = []

        from py4j.java_gateway import java_import
        java_import(self._jvm, "org.tresamigos.smv.python.SmvPythonHelper")
        java_import(self._jvm, "org.tresamigos.smv.DataSetRepoFactoryPython")

        self.helper = self._jvm.SmvPythonHelper

    def stages(self):
        return self.smvconfig.stage_names()

    def tx(self):
        """Create a TXContext for multiple places, avoid the long TXContext line
        """
        return TXContext(self._jvm, self.dsRepoFactories, self.stages())

    def load(self, *urns):
        """Load SmvGenericModules for specified URNs

        Args:
            *urns (str): list of URNs as strings

        Returns:
            list(SmvGenericModules): list of Scala SmvGenericModules (j_ds)
        """
        with self.tx() as tx:
            return tx.load(urns)

    def inferDS(self, *partial_names):
        """Return DSs from a list of partial names

        Args:
            *partial_names (str): list of partial names

        Returns:
            list(SmvGenericModules): list of SmvGenericModules
        """
        with self.tx() as tx:
            return tx.inferDS(partial_names)

    def inferUrn(self, partial_name):
        """Return URN string from partial name
        """
        return self.inferDS(partial_name)[0].urn()

    def register(self, repo_factory):
        """Register python repo factory
        """
        self.dsRepoFactories.append(repo_factory)

    def allDataSets(self):
        """Return all the SmvGenericModules in the app
        """
        with self.tx() as tx:
            return tx.allDataSets()

    def modulesToRun(self, modPartialNames, stageNames, allMods):
        """Return a modules need to run
            Combine specified modules, (-m), stages, (-s) and if
            (--run-app) specified, all output modules
        """
        with self.tx() as tx:
            named_mods = tx.inferDS(modPartialNames)
            stage_mods = tx.outputModulesForStage(stageNames)
            app_mods = tx.allOutputModules() if allMods else []
            res = []
            res.extend(named_mods)
            res.extend(stage_mods)
            res.extend(app_mods)
            # Need to perserve the ordering
            return list_distinct(res)

class TXContext(object):
    """Create a TX context for "with tx() as tx" syntax
    """
    def __init__(self, _jvm, resourceFactories, stages):
        self._jvm = _jvm
        self.resourceFactories = resourceFactories
        self.stages = stages

    def __enter__(self):
        return TX(self._jvm, self.resourceFactories, self.stages)

    def __exit__(self, type, value, traceback):
        pass

class TX(object):
    """Abstraction of the transaction boundary for loading SmvGenericModules.
        A TX object

        * will instantiate a set of repos when itself instantiated and will
        * reuse the same repos for all queries. This means that each new TX object will
        * reload the SmvGenericModules from source **once** during its lifetime.

        NOTE: Once a new TX is created, the well-formedness of the SmvGenericModules provided
        by the previous TX is not guaranteed. Particularly it may become impossible
        to run modules from the previous TX.
    """
    def __init__(self, _jvm, resourceFactories, stages):
        self.repos = [rf.createRepo() for rf in resourceFactories]
        self.stages = stages
        self.resolver = DataSetResolver(self.repos[0])
        self.log = smv.logger

    def load(self, urns):
        fqns = [u[4:] for u in urns]
        return self.resolver.loadDataSet(fqns)

    def inferDS(self, partial_names):
        return self.load(self._inferUrn(partial_names))

    def allDataSets(self):
        return self.load(self._allUrns())

    def allOutputModules(self):
        return self._filterOutput(self.allDataSets())

    def outputModulesForStage(self, stageNames):
        return self._filterOutput(self._dsForStage(stageNames))

    def _dsForStage(self, stageNames):
        return self.load(self._urnsForStage(stageNames))

    def _urnsForStage(self, stageNames):
        return [u
            for repo in self.repos
            for s in stageNames
            for u in repo._dataSetsForStage(s)
        ]

    def _allUrns(self):
        if (len(self.stages) == 0):
            log.warn("No stage names configured. Unable to discover any modules.")
        return self._urnsForStage(self.stages)

    def _inferUrn(self, partial_names):
        def urn_str(pn):
            return infer_full_name_from_part(
                self._allUrns(),
                pn
            )

        return [urn_str(pn) for pn in partial_names]

    def _filterOutput(self, dss):
        return [ds for ds in dss if ds.isSmvOutput()]<|endoftext|>


<commit_before>from datetime import datetime, timedelta

from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save

# this is the glue to the activity events framework, provided as a no-op here
def issue_update(kind, **kwargs):
    pass

class ForumCategory(models.Model):

    name = models.CharField(max_length=100)
    parent = models.ForeignKey("self", null=True, blank=True, related_name="subcategories")
    # @@@ total descendant forum count?
    # @@@ make group-aware

    class Meta:
        verbose_name_plural = "forum categories"

    def __unicode__(self):
        return self.name

class Forum(models.Model):

    title = models.CharField(max_length=100)
    description = models.TextField()
    creation_date = models.DateTimeField(default=datetime.now)

    # must only have one of these (or neither):
    parent = models.ForeignKey("self", null=True, blank=True, related_name="subforums")
    category = models.ForeignKey(ForumCategory, null=True, blank=True, related_name="forums")

    # @@@ make group-aware

    last_modified = models.DateTimeField(default=datetime.now)
    last_reply = models.ForeignKey("ForumReply", null=True, editable=False)

    view_count = models.IntegerField(default=0, editable=False)
    reply_count = models.IntegerField(default=0, editable=False)

    # this is what gets run normally
    def inc_views(self):
        self.view_count += 1
        self.save()

    # this can be used occasionally to get things back in sync
    def update_view_count(self):
        view_count = 0
        for thread in self.threads.all():
            view_count += thread.view_count
        self.view_count = view_count
        self.save()

    def update_reply_count(self):
        reply_count = 0
        for forum in self.subforums.all():
            forum.update_reply_count()
            reply_count += forum.reply_count
        for thread in self.threads.all():
            thread.update_reply_count()
            reply_count += thread.reply_count
        self.reply_count = reply_count
        self.save()

    def new_reply(self, reply):
        self.reply_count += 1 # if this gets out of sync run update_reply_count
        self.last_modified = reply.timestamp
        self.last_reply = reply
        self.save()
        if self.parent:
            self.parent.new_reply(reply)

    def __unicode__(self):
        return self.title

class ForumThread(models.Model):

    forum = models.ForeignKey(Forum, related_name="threads")

    title = models.CharField(max_length=100)
    creation_date = models.DateTimeField(default=datetime.now)
    author = models.ForeignKey(User, related_name="threads_started")
    last_modified = models.DateTimeField(default=datetime.now)
    last_reply = models.ForeignKey("ForumReply", null=True) # only temporarily null

    # @@@ support markup
    content = models.TextField()

    # @@@ sticky threads
    # @@@ closed threads

    view_count = models.IntegerField(default=0, editable=False)
    reply_count = models.IntegerField(default=0, editable=False)
    subscriber_count = models.IntegerField(default=0, editable=False)

    def inc_views(self):
        self.view_count += 1
        self.save()
        self.forum.inc_views()

    def update_reply_count(self):
        self.reply_count = self.replies.all().count()
        self.save()

    def update_subscriber_count(self):
        self.subscriber_count = self.subscriptions.count()
        self.save()

    def new_reply(self, reply):
        self.reply_count += 1
        self.last_modified = reply.timestamp
        self.last_reply = reply
        self.save()
        self.forum.new_reply(reply)

    def __unicode__(self):
        return self.title

class ForumReply(models.Model):

    thread = models.ForeignKey(ForumThread, related_name="replies")

    author = models.ForeignKey(User, related_name="replies")
    timestamp = models.DateTimeField(default=datetime.now)

    # @@@ support markup
    content = models.TextField()

    # allow editing for short period after posting
    def editable(self, user):
        if user == self.author:
            if datetime.now() < self.timestamp + timedelta(minutes=30): # @@@ factor out time interval
                return True
        return False

class UserPostCount(models.Model):

    user = models.ForeignKey(User, related_name="post_count")
    count = models.IntegerField(default=0)

class ThreadSubscription(models.Model):

    thread = models.ForeignKey(ForumThread, related_name="subscriptions")
    user = models.ForeignKey(User, related_name="forum_subscriptions")

    class Meta:
          unique_together = [("thread", "user")]

def forum_thread_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        issue_update("forum_thread", user=instance.author, forum_thread=instance)

def forum_reply_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        thread = instance.thread
        thread.new_reply(instance)
        # @@@ this next part could be manager method
        post_count, created = UserPostCount.objects.get_or_create(user=instance.author)
        post_count.count += 1
        post_count.save()
        issue_update("forum_reply", user=instance.author, forum_reply=instance)
        # @@@ where do subscribers get notified?

def forum_subscription_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        thread = instance.thread
        thread.update_subscriber_count()

post_save.connect(forum_thread_save, sender=ForumThread)
post_save.connect(forum_reply_save, sender=ForumReply)
post_save.connect(forum_subscription_save, sender=ThreadSubscription)

# @@@ handling unsubscribe?

# @@@ handling deletion? (e.g. counts, last_modified, last_reply)
<commit_msg>more consistent field naming<commit_after>from datetime import datetime, timedelta

from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save

# this is the glue to the activity events framework, provided as a no-op here
def issue_update(kind, **kwargs):
    pass

class ForumCategory(models.Model):

    title = models.CharField(max_length=100)

    parent = models.ForeignKey("self", null=True, blank=True, related_name="subcategories")

    # @@@ total descendant forum count?
    # @@@ make group-aware

    class Meta:
        verbose_name_plural = "forum categories"

    def __unicode__(self):
        return self.name

class Forum(models.Model):

    title = models.CharField(max_length=100)
    description = models.TextField()

    # must only have one of these (or neither):
    parent = models.ForeignKey("self", null=True, blank=True, related_name="subforums")
    category = models.ForeignKey(ForumCategory, null=True, blank=True, related_name="forums")

    # @@@ make group-aware

    last_modified = models.DateTimeField(default=datetime.now)
    last_reply = models.ForeignKey("ForumReply", null=True, editable=False)

    view_count = models.IntegerField(default=0, editable=False)
    reply_count = models.IntegerField(default=0, editable=False)

    # this is what gets run normally
    def inc_views(self):
        self.view_count += 1
        self.save()

    # this can be used occasionally to get things back in sync
    def update_view_count(self):
        view_count = 0
        for thread in self.threads.all():
            view_count += thread.view_count
        self.view_count = view_count
        self.save()

    def update_reply_count(self):
        reply_count = 0
        for forum in self.subforums.all():
            forum.update_reply_count()
            reply_count += forum.reply_count
        for thread in self.threads.all():
            thread.update_reply_count()
            reply_count += thread.reply_count
        self.reply_count = reply_count
        self.save()

    def new_reply(self, reply):
        self.reply_count += 1 # if this gets out of sync run update_reply_count
        self.last_modified = reply.timestamp
        self.last_reply = reply
        self.save()
        if self.parent:
            self.parent.new_reply(reply)

    def __unicode__(self):
        return self.title

class ForumThread(models.Model):

    forum = models.ForeignKey(Forum, related_name="threads")

    title = models.CharField(max_length=100)

    created = models.DateTimeField(default=datetime.now)
    author = models.ForeignKey(User, related_name="threads_started")

    last_modified = models.DateTimeField(default=datetime.now)
    last_reply = models.ForeignKey("ForumReply", null=True) # only temporarily null

    # @@@ support markup
    content = models.TextField()

    # @@@ sticky threads
    # @@@ closed threads

    view_count = models.IntegerField(default=0, editable=False)
    reply_count = models.IntegerField(default=0, editable=False)
    subscriber_count = models.IntegerField(default=0, editable=False)

    def inc_views(self):
        self.view_count += 1
        self.save()
        self.forum.inc_views()

    def update_reply_count(self):
        self.reply_count = self.replies.all().count()
        self.save()

    def update_subscriber_count(self):
        self.subscriber_count = self.subscriptions.count()
        self.save()

    def new_reply(self, reply):
        self.reply_count += 1
        self.last_modified = reply.timestamp
        self.last_reply = reply
        self.save()
        self.forum.new_reply(reply)

    def __unicode__(self):
        return self.title

class ForumReply(models.Model):

    thread = models.ForeignKey(ForumThread, related_name="replies")

    author = models.ForeignKey(User, related_name="replies")
    created = models.DateTimeField(default=datetime.now)

    # @@@ support markup
    content = models.TextField()

    # allow editing for short period after posting
    def editable(self, user):
        if user == self.author:
            if datetime.now() < self.timestamp + timedelta(minutes=30): # @@@ factor out time interval
                return True
        return False

class UserPostCount(models.Model):

    user = models.ForeignKey(User, related_name="post_count")
    count = models.IntegerField(default=0)

class ThreadSubscription(models.Model):

    thread = models.ForeignKey(ForumThread, related_name="subscriptions")
    user = models.ForeignKey(User, related_name="forum_subscriptions")

    class Meta:
          unique_together = [("thread", "user")]

def forum_thread_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        issue_update("forum_thread", user=instance.author, forum_thread=instance)

def forum_reply_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        thread = instance.thread
        thread.new_reply(instance)
        # @@@ this next part could be manager method
        post_count, created = UserPostCount.objects.get_or_create(user=instance.author)
        post_count.count += 1
        post_count.save()
        issue_update("forum_reply", user=instance.author, forum_reply=instance)
        # @@@ where do subscribers get notified?

def forum_subscription_save(sender, instance=None, created=False, **kwargs):
    if instance and created:
        thread = instance.thread
        thread.update_subscriber_count()

post_save.connect(forum_thread_save, sender=ForumThread)
post_save.connect(forum_reply_save, sender=ForumReply)
post_save.connect(forum_subscription_save, sender=ThreadSubscription)

# @@@ handling unsubscribe?

# @@@ handling deletion? (e.g. counts, last_modified, last_reply)
<|endoftext|>


<commit_before>from datetime import datetime

from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save

from django.contrib.auth.models import User
from django.contrib.contenttypes import generic
from django.contrib.contenttypes.models import ContentType

from voting.models import Vote

class Question(models.Model):

    group_content_type = models.ForeignKey(ContentType, null=True, blank=True)
    group_object_id = models.IntegerField(null=True, blank=True)
    group = generic.GenericForeignKey("group_content_type", "group_object_id")

    question = models.CharField(max_length=100)
    content = models.TextField()
    user = models.ForeignKey(User, related_name="questions")
    created = models.DateTimeField(default=datetime.now)

    score = models.IntegerField(editable=False, default=0)
    vote_count = models.IntegerField(editable=False, default=0)

    def update_score(self):
        result = Vote.objects.get_score(self)
        self.score = result["score"]
        self.vote_count = result["num_votes"]
        self.save()

    def get_absolute_url(self, group=None):
        kwargs = {
            "question_id": self.pk,
        }
        if group:
            return group.content_bridge.reverse("questions_question_detail", kwargs=kwargs)
        return reverse("questions_question_detail", kwargs=kwargs)

class Response(models.Model):

    question = models.ForeignKey(Question, related_name="responses")
    content = models.TextField()
    accepted = models.BooleanField(default=False)
    user = models.ForeignKey(User, related_name="responses")
    created = models.DateTimeField(default=datetime.now)

    score = models.IntegerField(editable=False, default=0)
    vote_count = models.IntegerField(editable=False, default=0)

    def update_score(self):
        result = Vote.objects.get_score(self)
        self.score = result["score"]
        self.vote_count = result["num_votes"]
        self.save()

    def accept(self):
        # check for another active one and mark it inactive
        try:
            response = Response.objects.get(question=self.question, accepted=True)
        except Response.DoesNotExist:
            pass
        else:
            if self != response:
                response.accepted = False
                response.save()
        self.accepted = True
        self.save()

    def get_absolute_url(self, group=None):
        return "%s#response-%d" % (self.question.get_absolute_url(group), self.pk)

def vote_save(sender, instance=None, **kwargs):
    if instance:
        # GFK "join" -- issues a query
        obj = instance.object
        if isinstance(obj, (Question, Response)):
            obj.update_score()

post_save.connect(vote_save, sender=Vote)
<commit_msg>always reverse group attached to instance instead of requiring a group kwarg<commit_after>from datetime import datetime

from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save

from django.contrib.auth.models import User
from django.contrib.contenttypes import generic
from django.contrib.contenttypes.models import ContentType

from voting.models import Vote

class Question(models.Model):

    group_content_type = models.ForeignKey(ContentType, null=True, blank=True)
    group_object_id = models.IntegerField(null=True, blank=True)
    group = generic.GenericForeignKey("group_content_type", "group_object_id")

    question = models.CharField(max_length=100)
    content = models.TextField()
    user = models.ForeignKey(User, related_name="questions")
    created = models.DateTimeField(default=datetime.now)

    score = models.IntegerField(editable=False, default=0)
    vote_count = models.IntegerField(editable=False, default=0)

    def update_score(self):
        result = Vote.objects.get_score(self)
        self.score = result["score"]
        self.vote_count = result["num_votes"]
        self.save()

    def get_absolute_url(self):
        kwargs = {
            "question_id": self.pk,
        }
        if self.group:
            reverse = self.group.content_bridge.reverse
        return reverse("questions_question_detail", kwargs=kwargs)

class Response(models.Model):

    question = models.ForeignKey(Question, related_name="responses")
    content = models.TextField()
    accepted = models.BooleanField(default=False)
    user = models.ForeignKey(User, related_name="responses")
    created = models.DateTimeField(default=datetime.now)

    score = models.IntegerField(editable=False, default=0)
    vote_count = models.IntegerField(editable=False, default=0)

    def update_score(self):
        result = Vote.objects.get_score(self)
        self.score = result["score"]
        self.vote_count = result["num_votes"]
        self.save()

    def accept(self):
        # check for another active one and mark it inactive
        try:
            response = Response.objects.get(question=self.question, accepted=True)
        except Response.DoesNotExist:
            pass
        else:
            if self != response:
                response.accepted = False
                response.save()
        self.accepted = True
        self.save()

    def get_absolute_url(self):
        return "%s#response-%d" % (self.question.get_absolute_url(), self.pk)

def vote_save(sender, instance=None, **kwargs):
    if instance:
        # GFK "join" -- issues a query
        obj = instance.object
        if isinstance(obj, (Question, Response)):
            obj.update_score()

post_save.connect(vote_save, sender=Vote)
<|endoftext|>


<commit_before><commit_msg>Bug fix.<commit_after><|endoftext|>


<commit_before><commit_msg>remove unusued function<commit_after><|endoftext|>


<commit_before># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Copyright (C) 2017, James R. Barlow (https://github.com/jbarlow83/)
import os
from abc import ABC, abstractmethod
from pathlib import Path
from subprocess import DEVNULL, PIPE, CalledProcessError, run
from tempfile import TemporaryDirectory
from typing import Optional, cast

import deprecation
from packaging.version import Version
from PIL import Image

import pikepdf
from pikepdf._exceptions import DependencyError

@deprecation.deprecated(
    deprecated_in="5.1.5", removed_in="6.0", details="Use extract_jbig2_bytes instead"
)
def extract_jbig2(
    im_obj: pikepdf.Object, globals_obj: Optional[pikepdf.Object] = None
) -> Image.Image:  # pragma: no cover

    with TemporaryDirectory(prefix='pikepdf-', suffix='.jbig2') as tmpdir:
        image_path = Path(tmpdir) / "image"
        global_path = Path(tmpdir) / "global"
        output_path = Path(tmpdir) / "outfile"

        args = [
            "jbig2dec",
            "--embedded",
            "--format",
            "png",
            "--output",
            os.fspath(output_path),
        ]

        # Get the raw stream, because we can't decode im_obj - that is why we are here
        # (Strictly speaking we should remove any non-JBIG2 filters if double encoded)
        image_path.write_bytes(cast(memoryview, im_obj.get_raw_stream_buffer()))

        if globals_obj is not None:
            # For globals, we do want to remove any encoding since it's just a binary
            # blob and won't be marked with /JBIG2Decode
            global_path.write_bytes(cast(memoryview, globals_obj.get_stream_buffer()))
            args.append(os.fspath(global_path))

        args.append(os.fspath(image_path))

        run(args, stdout=DEVNULL, check=True)
        im = Image.open(output_path)
        im.load()  # Load pixel data into memory so file/tempdir can be closed
        return im

def _extract_jbig2_bytes(jbig2: bytes, jbig2_globals: bytes) -> bytes:
    with TemporaryDirectory(prefix='pikepdf-', suffix='.jbig2') as tmpdir:
        image_path = Path(tmpdir) / "image"
        global_path = Path(tmpdir) / "global"
        output_path = Path(tmpdir) / "outfile"

        args = [
            "jbig2dec",
            "--embedded",
            "--format",
            "png",
            "--output",
            os.fspath(output_path),
        ]

        # Get the raw stream, because we can't decode im_obj - that is why we are here
        # (Strictly speaking we should remove any non-JBIG2 filters if double encoded)
        image_path.write_bytes(jbig2)

        if len(jbig2_globals) > 0:
            global_path.write_bytes(jbig2_globals)
            args.append(os.fspath(global_path))

        args.append(os.fspath(image_path))

        run(args, stdout=DEVNULL, check=True)
        with Image.open(output_path) as im:
            return im.tobytes()

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def extract_jbig2_bytes(
    jbig2: bytes, jbig2_globals: bytes
) -> bytes:  # pragma: no cover
    return _extract_jbig2_bytes(jbig2, jbig2_globals)

def _check_jbig2dec_available() -> None:
    try:
        proc = run(['jbig2dec', '--version'], stdout=PIPE, check=True, encoding='ascii')
    except (CalledProcessError, FileNotFoundError) as e:
        raise DependencyError("jbig2dec - not installed or not found") from e
    else:
        result = proc.stdout
        version_str = result.replace('jbig2dec', '').strip()  # returns "jbig2dec 0.xx"
        version = Version(version_str)
        if version < Version('0.15'):
            raise DependencyError("jbig2dec is too old (older than version 0.15)")

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def check_jbig2dec_available() -> None:  # pragma: no cover
    _check_jbig2dec_available()

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def jbig2dec_available() -> bool:  # pragma: no cover
    try:
        _check_jbig2dec_available()
    except (DependencyError, CalledProcessError, FileNotFoundError):
        return False
    else:
        return True

class JBIG2DecoderInterface(ABC):
    @abstractmethod
    def check_available(self) -> None:
        ...

    @abstractmethod
    def decode_jbig2(self, jbig2: bytes, jbig2_globals: bytes) -> bytes:
        ...

    def available(self) -> bool:
        try:
            self.check_available()
        except DependencyError:
            return False
        else:
            return True

class JBIG2Decoder(JBIG2DecoderInterface):
    def check_available(self) -> None:
        version = self._version()
        if version < Version('0.15'):
            raise DependencyError("jbig2dec is too old (older than version 0.15)")

    def decode_jbig2(self, jbig2: bytes, jbig2_globals: bytes) -> bytes:
        return _extract_jbig2_bytes(jbig2, jbig2_globals)

    def _version(self) -> Version:
        try:
            proc = run(
                ['jbig2dec', '--version'], stdout=PIPE, check=True, encoding='ascii'
            )
        except (CalledProcessError, FileNotFoundError) as e:
            raise DependencyError("jbig2dec - not installed or not found") from e
        else:
            result = proc.stdout
            version_str = result.replace(
                'jbig2dec', ''
            ).strip()  # returns "jbig2dec 0.xx"
            return Version(version_str)

_jbig2_decoder = JBIG2Decoder()

def get_decoder():
    return _jbig2_decoder
<commit_msg>JBIG2: skip coverage on deprecrated code<commit_after># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Copyright (C) 2017, James R. Barlow (https://github.com/jbarlow83/)
import os
from abc import ABC, abstractmethod
from pathlib import Path
from subprocess import DEVNULL, PIPE, CalledProcessError, run
from tempfile import TemporaryDirectory
from typing import Optional, cast

import deprecation
from packaging.version import Version
from PIL import Image

import pikepdf
from pikepdf._exceptions import DependencyError

@deprecation.deprecated(
    deprecated_in="5.1.5", removed_in="6.0", details="Use extract_jbig2_bytes instead"
)
def extract_jbig2(
    im_obj: pikepdf.Object, globals_obj: Optional[pikepdf.Object] = None
) -> Image.Image:  # pragma: no cover

    with TemporaryDirectory(prefix='pikepdf-', suffix='.jbig2') as tmpdir:
        image_path = Path(tmpdir) / "image"
        global_path = Path(tmpdir) / "global"
        output_path = Path(tmpdir) / "outfile"

        args = [
            "jbig2dec",
            "--embedded",
            "--format",
            "png",
            "--output",
            os.fspath(output_path),
        ]

        # Get the raw stream, because we can't decode im_obj - that is why we are here
        # (Strictly speaking we should remove any non-JBIG2 filters if double encoded)
        image_path.write_bytes(cast(memoryview, im_obj.get_raw_stream_buffer()))

        if globals_obj is not None:
            # For globals, we do want to remove any encoding since it's just a binary
            # blob and won't be marked with /JBIG2Decode
            global_path.write_bytes(cast(memoryview, globals_obj.get_stream_buffer()))
            args.append(os.fspath(global_path))

        args.append(os.fspath(image_path))

        run(args, stdout=DEVNULL, check=True)
        im = Image.open(output_path)
        im.load()  # Load pixel data into memory so file/tempdir can be closed
        return im

def _extract_jbig2_bytes(jbig2: bytes, jbig2_globals: bytes) -> bytes:
    with TemporaryDirectory(prefix='pikepdf-', suffix='.jbig2') as tmpdir:
        image_path = Path(tmpdir) / "image"
        global_path = Path(tmpdir) / "global"
        output_path = Path(tmpdir) / "outfile"

        args = [
            "jbig2dec",
            "--embedded",
            "--format",
            "png",
            "--output",
            os.fspath(output_path),
        ]

        # Get the raw stream, because we can't decode im_obj - that is why we are here
        # (Strictly speaking we should remove any non-JBIG2 filters if double encoded)
        image_path.write_bytes(jbig2)

        if len(jbig2_globals) > 0:
            global_path.write_bytes(jbig2_globals)
            args.append(os.fspath(global_path))

        args.append(os.fspath(image_path))

        run(args, stdout=DEVNULL, check=True)
        with Image.open(output_path) as im:
            return im.tobytes()

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def extract_jbig2_bytes(
    jbig2: bytes, jbig2_globals: bytes
) -> bytes:  # pragma: no cover
    return _extract_jbig2_bytes(jbig2, jbig2_globals)

def _check_jbig2dec_available() -> None:  # pragma: no cover
    try:
        proc = run(['jbig2dec', '--version'], stdout=PIPE, check=True, encoding='ascii')
    except (CalledProcessError, FileNotFoundError) as e:
        raise DependencyError("jbig2dec - not installed or not found") from e
    else:
        result = proc.stdout
        version_str = result.replace('jbig2dec', '').strip()  # returns "jbig2dec 0.xx"
        version = Version(version_str)
        if version < Version('0.15'):
            raise DependencyError("jbig2dec is too old (older than version 0.15)")

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def check_jbig2dec_available() -> None:  # pragma: no cover
    _check_jbig2dec_available()

@deprecation.deprecated(
    deprecated_in="5.2.0",
    removed_in="6.0",
    details="Use jbig2.get_decoder() interface instead",
)
def jbig2dec_available() -> bool:  # pragma: no cover
    try:
        _check_jbig2dec_available()
    except (DependencyError, CalledProcessError, FileNotFoundError):
        return False
    else:
        return True

class JBIG2DecoderInterface(ABC):
    """pikepdf's C++ expects this Python interface to be available for JBIG2."""

    @abstractmethod
    def check_available(self) -> None:
        """Check if decoder is available. Throws DependencyError if not."""

    @abstractmethod
    def decode_jbig2(self, jbig2: bytes, jbig2_globals: bytes) -> bytes:
        """Decode JBIG2 from jbig2 and globals, returning decoded bytes."""

    def available(self) -> bool:
        """Returns True if decoder is available."""
        try:
            self.check_available()
        except DependencyError:
            return False
        else:
            return True

class JBIG2Decoder(JBIG2DecoderInterface):
    def check_available(self) -> None:
        version = self._version()
        if version < Version('0.15'):
            raise DependencyError("jbig2dec is too old (older than version 0.15)")

    def decode_jbig2(self, jbig2: bytes, jbig2_globals: bytes) -> bytes:
        return _extract_jbig2_bytes(jbig2, jbig2_globals)

    def _version(self) -> Version:
        try:
            proc = run(
                ['jbig2dec', '--version'], stdout=PIPE, check=True, encoding='ascii'
            )
        except (CalledProcessError, FileNotFoundError) as e:
            raise DependencyError("jbig2dec - not installed or not found") from e
        else:
            result = proc.stdout
            version_str = result.replace(
                'jbig2dec', ''
            ).strip()  # returns "jbig2dec 0.xx"
            return Version(version_str)

_jbig2_decoder = JBIG2Decoder()

def get_decoder():
    return _jbig2_decoder
<|endoftext|>


<commit_before>'''Declarative ORM event decorators and event registration.
'''

from functools import wraps, partial
from collections import namedtuple

import sqlalchemy
import six

Event = namedtuple('Event', ['name', 'attribute', 'listener', 'kargs'])

def register(cls, dct):
    '''Register events defined on a class during metaclass creation.'''

    events = []

    # append class attribute defined events
    if dct.get('__events__'):
        # events defined on __events__ can have many forms (e.g. string based, list of tuples, etc)
        # so we need to iterate over them and parse into standardized Event object
        for event_name, listeners in six.iteritems(dct['__events__']):
            if not isinstance(listeners, list):
                listeners = [listeners]

            for listener in listeners:
                if isinstance(listener, tuple):
                    # listener definition includes event.listen keyword args
                    listener, kargs = listener
                else:
                    kargs = {}

                if not callable(listener):
                    # assume listener is a string reference to class method
                    listener = dct[listener]

                events.append(Event(event_name, kargs.pop('attribute', None), listener, kargs))

    # add events which were added via @event decorator
    for value in dct.values():
        if hasattr(value, '__event__'):
            if not isinstance(value.__event__, list):  # pragma: no cover
                value.__event__ = [value.__event__]
            events.extend(value.__event__)

    if events:
        # reassemble events dict into consistent form using Event objects as values
        events_dict = {}
        for evt in events:
            obj = cls if evt.attribute is None else getattr(cls, evt.attribute)
            event_name = evt.name.replace('on_', '', 1) if evt.name.startswith('on_') else evt.name

            sqlalchemy.event.listen(obj, event_name, evt.listener, **evt.kargs)
            events_dict.setdefault(evt.name, []).append(evt)

        dct['__events__'].update(events_dict)

#####
## Event decorators
####

def event(event_names, attribute=None, f=None, **event_kargs):
    '''
    Generic event decorator maker which attaches metadata to function object
    so that `register()` can find the event definition.
    '''
    if callable(f):
        # being called as a decorator with no arguments
        return event(event_names)(f)

    def decorator(f):
        if not hasattr(f, '__event__'):
            # set initial value to list so a function can handle multiple events
            f.__event__ = []

        # @note: have to assign to a separately variable name due to global name access issues
        if not isinstance(event_names, (list, tuple)):
            _event_names = [event_names]
        else:
            _event_names = event_names

        # Attach event object to function which will be picked up in `register()`.
        f.__event__ += [Event(event_name, attribute, f, event_kargs) for event_name in _event_names]

        # Return function as-is since method definition should be compatible with sqlalchemy.event.listen().
        return f
    return decorator

def make_attribute_event(event_names):
    '''Event decorator maker for attribute events.'''
    return partial(event, event_names)

def make_event(event_names):
    '''Event decorator maker for mapper or instance events which don't require an attribute.'''
    # bind `None` to attribute argument
    return partial(event, event_names, None)

##
# Attribute Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#attribute-events
##

on_set = make_attribute_event('set')
on_append = make_attribute_event('append')
on_remove = make_attribute_event('remove')

##
# Mapper Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#mapper-events
##

before_delete = make_event('before_delete')
before_insert = make_event('before_insert')
before_update = make_event('before_update')
before_insert_update = make_event(['before_insert', 'before_update'])

after_delete = make_event('after_delete')
after_insert = make_event('after_insert')
after_update = make_event('after_update')
after_insert_update = make_event(['after_insert', 'after_update'])

on_append_result = make_event('append_result')
on_create_instance = make_event('create_instance')
on_instrument_class = make_event('instrument_class')
before_configured = make_event('before_configured')
after_configured = make_event('after_configured')
on_mapper_configured = make_event('mapper_configured')
on_populate_instance = make_event('populate_instance')
on_translate_row = make_event('translate_row')

##
# Instance Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#instance-events
##

on_expire = make_event('expire')
on_load = make_event('load')
on_refresh = make_event('refresh')

# The following events work as intended, but they don't seem like
# candidates for supporting their definition on the model class.

# @why: Having an on-some-init event defined on the model class
# seems inefficient since whatever logic they contain should be
# handled in model.__init__() anyway.
#on_first_init = make_event('first_init')
#on_init = make_event('init')
#on_init_failure = make_event('init_failure')

# @why: Again model class would already define pickle support
# so logic should be contained there and not in a separate event handler.
#on_pickle = make_event('pickle')
#on_unpickle = make_event('unpickle')

# @why: Well, not really sure how to actually trigger this event
# so don't want to support it if it doesn't have a test.
# If someone really wants this event, then it can be enabled.
#on_resurrect = make_event('resurrect')
<commit_msg>Rename function arguments for clarity.<commit_after>'''Declarative ORM event decorators and event registration.
'''

from functools import partial
from collections import namedtuple

import sqlalchemy
import six

Event = namedtuple('Event', ['name', 'attribute', 'listener', 'kargs'])

def register(cls, dct):
    '''Register events defined on a class during metaclass creation.'''

    events = []

    # append class attribute defined events
    if dct.get('__events__'):
        # events defined on __events__ can have many forms (e.g. string based, list of tuples, etc)
        # so we need to iterate over them and parse into standardized Event object
        for event_name, listeners in six.iteritems(dct['__events__']):
            if not isinstance(listeners, list):
                listeners = [listeners]

            for listener in listeners:
                if isinstance(listener, tuple):
                    # listener definition includes event.listen keyword args
                    listener, kargs = listener
                else:
                    kargs = {}

                if not callable(listener):
                    # assume listener is a string reference to class method
                    listener = dct[listener]

                events.append(Event(event_name, kargs.pop('attribute', None), listener, kargs))

    # add events which were added via @event decorator
    for value in dct.values():
        if hasattr(value, '__event__'):
            if not isinstance(value.__event__, list):  # pragma: no cover
                value.__event__ = [value.__event__]
            events.extend(value.__event__)

    if events:
        # reassemble events dict into consistent form using Event objects as values
        events_dict = {}
        for evt in events:
            obj = cls if evt.attribute is None else getattr(cls, evt.attribute)
            event_name = evt.name.replace('on_', '', 1) if evt.name.startswith('on_') else evt.name

            sqlalchemy.event.listen(obj, event_name, evt.listener, **evt.kargs)
            events_dict.setdefault(evt.name, []).append(evt)

        dct['__events__'].update(events_dict)

#####
## Event decorators
####

def event(event_names, attribute=None, no_args_func=None, **event_kargs):
    '''Generic event decorator maker which attaches metadata to function object
    so that `register()` can find the event definition.
    '''
    if callable(no_args_func):
        # being called as a decorator with no arguments
        return event(event_names)(no_args_func)

    def decorator(func):
        '''Function decorator that attaches an `__event__` attribute hook which is
        expected when registering a method as an event handler. See `register()` in
        this module for details on how this is implemented.
        '''
        if not hasattr(func, '__event__'):
            # set initial value to list so a function can handle multiple events
            func.__event__ = []

        # @note: have to assign to a separately variable name due to global name access issues
        if not isinstance(event_names, (list, tuple)):
            _event_names = [event_names]
        else:
            _event_names = event_names

        # Attach event object to function which will be picked up in `register()`.
        func.__event__ += [Event(event_name, attribute, func, event_kargs) for event_name in _event_names]

        # Return function as-is since method definition should be compatible with sqlalchemy.event.listen().
        return func
    return decorator

def make_attribute_event(event_names):
    '''Event decorator maker for attribute events.'''
    return partial(event, event_names)

def make_event(event_names):
    '''Event decorator maker for mapper or instance events which don't require an attribute.'''
    # bind `None` to attribute argument
    return partial(event, event_names, None)

##
# Attribute Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#attribute-events
##

on_set = make_attribute_event('set')
on_append = make_attribute_event('append')
on_remove = make_attribute_event('remove')

##
# Mapper Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#mapper-events
##

before_delete = make_event('before_delete')
before_insert = make_event('before_insert')
before_update = make_event('before_update')
before_insert_update = make_event(['before_insert', 'before_update'])

after_delete = make_event('after_delete')
after_insert = make_event('after_insert')
after_update = make_event('after_update')
after_insert_update = make_event(['after_insert', 'after_update'])

on_append_result = make_event('append_result')
on_create_instance = make_event('create_instance')
on_instrument_class = make_event('instrument_class')
before_configured = make_event('before_configured')
after_configured = make_event('after_configured')
on_mapper_configured = make_event('mapper_configured')
on_populate_instance = make_event('populate_instance')
on_translate_row = make_event('translate_row')

##
# Instance Events
# http://docs.sqlalchemy.org/en/latest/orm/events.html#instance-events
##

on_expire = make_event('expire')
on_load = make_event('load')
on_refresh = make_event('refresh')

# The following events work as intended, but they don't seem like
# candidates for supporting their definition on the model class.

# @why: Having an on-some-init event defined on the model class
# seems inefficient since whatever logic they contain should be
# handled in model.__init__() anyway.
#on_first_init = make_event('first_init')
#on_init = make_event('init')
#on_init_failure = make_event('init_failure')

# @why: Again model class would already define pickle support
# so logic should be contained there and not in a separate event handler.
#on_pickle = make_event('pickle')
#on_unpickle = make_event('unpickle')

# @why: Well, not really sure how to actually trigger this event
# so don't want to support it if it doesn't have a test.
# If someone really wants this event, then it can be enabled.
#on_resurrect = make_event('resurrect')
<|endoftext|>


<commit_before># coding: utf-8
import os
import logging
from normality import slugify

from flask_script import Manager, commands as flask_script_commands
from flask_script.commands import ShowUrls
from flask_migrate import MigrateCommand

from aleph.core import create_app, archive, datasets
from aleph.model import db, upgrade_db, Collection, Document, Entity
from aleph.views import mount_app_blueprints
from aleph.analyze import install_analyzers
from aleph.ingest import reingest_collection, ingest_document
from aleph.index import init_search, delete_index, upgrade_search
from aleph.index import index_document_id, delete_dataset
from aleph.logic import reindex_entities, delete_collection, analyze_collection
from aleph.logic import load_dataset, update_entity_full, delete_pending
from aleph.logic import update_collection
from aleph.logic.alerts import check_alerts
from aleph.ext import get_crawlers

log = logging.getLogger('aleph')
flask_script_commands.text_type = str

app = create_app()
mount_app_blueprints(app)
manager = Manager(app)
manager.add_command('db', MigrateCommand)
manager.add_command('routes', ShowUrls)

@manager.command
def collections():
    """List all collections."""
    for collection in Collection.all():
        print collection.id, collection.foreign_id, collection.label

@manager.command
def alerts():
    """Generate alert notifications."""
    check_alerts.delay()

@manager.command
def crawl(name):
    """Execute the given crawler."""
    log.info('Crawling %r...', name)
    crawlers = get_crawlers()
    if name not in crawlers:
        log.info('No such crawler: %r', name)
    else:
        crawler = crawlers.get(name)()
        crawler.execute()
    db.session.commit()

@manager.command
@manager.option('-l', '--language', dest='language', nargs='*')
@manager.option('-c', '--country', dest='country', nargs='*')
@manager.option('-f', '--foreign_id', dest='foreign_id')
def crawldir(directory, language=None, country=None, foreign_id=None):
    """Crawl the given directory."""
    if directory is None or not os.path.exists(directory):
        log.error("Invalid directory: %r", directory)
        return
    directory = os.path.abspath(os.path.normpath(directory))

    if foreign_id is None:
        foreign_id = 'directory:%s' % slugify(directory)
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        collection = Collection.create({
            'foreign_id': foreign_id,
            'label': directory,
            'managed': True
        })

    if language is not None:
        collection.languages = [language]
    if country is not None:
        collection.countries = [country]
    db.session.commit()
    update_collection(collection)

    log.info('Crawling %r to %r...', directory, collection.foreign_id)
    document = Document.by_keys(collection_id=collection.id,
                                foreign_id=directory)
    ingest_document(document, directory)

@manager.command
def flush(foreign_id):
    """Reset the crawler state for a given collecton."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    delete_collection(collection.id)

@manager.command
def deletepending(foreign_id=None):
    """Deletes any pending entities and related items."""
    collection_id = None
    if foreign_id is None:
        collection = Collection.by_foreign_id(foreign_id)
        if collection is None:
            raise ValueError("No such collection: %r" % foreign_id)
        collection_id = collection.id
    delete_pending(collection_id=collection_id)

@manager.command
def analyze(foreign_id):
    """Re-analyze documents in the given collection (or throughout)."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    analyze_collection.delay(collection.id)

@manager.command
def reingest(foreign_id):
    """Re-ingest documents in the given collection."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    reingest_collection(collection)

@manager.command
def index(foreign_id=None):
    """Index documents in the given collection (or throughout)."""
    q = Document.all_ids()
    # re-index newest document first.
    q = q.order_by(Document.id.desc())
    if foreign_id:
        collection = Collection.by_foreign_id(foreign_id)
        if collection is None:
            raise ValueError("No such collection: %r" % foreign_id)
        q = q.filter(Document.collection_id == collection.id)
    for idx, (doc_id,) in enumerate(q.yield_per(10000), 1):
        index_document_id.delay(doc_id)
        if idx % 1000 == 0:
            log.info("Index: %s documents...", idx)
    if foreign_id is None:
        reindex_entities()

@manager.command
def loaddataset(name):
    """Index all the entities in a given dataset."""
    dataset = datasets.get(name)
    load_dataset(dataset)

@manager.command
def deletedataset(name):
    delete_dataset(name)

@manager.command
def resetindex():
    """Re-create the ES index configuration, dropping all data."""
    delete_index()
    init_search()

@manager.command
def indexentities():
    """Re-index all the entities."""
    reindex_entities()

@manager.command
def updateentities():
    """Re-index all the entities."""
    q = db.session.query(Entity.id)
    for (entity_id,) in q:
        update_entity_full.delay(entity_id)

@manager.command
@manager.option('-s', '--skip-downloads', dest='skip', default='')
def init(skip=''):
    """Create or upgrade the search index and database."""
    upgrade_db()
    init_search()
    upgrade_search()
    if 'analyzers' not in skip:
        install_analyzers()
    archive.upgrade()

@manager.command
def upgrade():
    """Create or upgrade the search index and database."""
    upgrade_db()
    upgrade_search()
    archive.upgrade()

@manager.command
def installdata():
    """Create or upgrade the search index and database."""
    install_analyzers()

@manager.command
def evilshit():
    """EVIL: Delete all data and recreate the database."""
    delete_index()
    from sqlalchemy import MetaData, inspect
    from sqlalchemy.exc import InternalError
    from sqlalchemy.dialects.postgresql import ENUM
    metadata = MetaData()
    metadata.bind = db.engine
    metadata.reflect()
    tables = list(metadata.sorted_tables)
    while len(tables):
        for table in tables:
            try:
                table.drop(checkfirst=True)
                tables.remove(table)
            except InternalError:
                pass
    for enum in inspect(db.engine).get_enums():
        enum = ENUM(name=enum['name'])
        enum.drop(bind=db.engine, checkfirst=True)
    init()

def main():
    manager.run()

if __name__ == "__main__":
    main()
<commit_msg>extra safe<commit_after># coding: utf-8
import os
import logging
from normality import slugify

from flask_script import Manager, commands as flask_script_commands
from flask_script.commands import ShowUrls
from flask_migrate import MigrateCommand

from aleph.core import create_app, archive, datasets
from aleph.model import db, upgrade_db, Collection, Document, Entity
from aleph.views import mount_app_blueprints
from aleph.analyze import install_analyzers
from aleph.ingest import reingest_collection, ingest_document
from aleph.index import init_search, delete_index, upgrade_search
from aleph.index import index_document_id, delete_dataset
from aleph.logic import reindex_entities, delete_collection, analyze_collection
from aleph.logic import load_dataset, update_entity_full, delete_pending
from aleph.logic import update_collection
from aleph.logic.alerts import check_alerts
from aleph.ext import get_crawlers

log = logging.getLogger('aleph')
flask_script_commands.text_type = str

app = create_app()
mount_app_blueprints(app)
manager = Manager(app)
manager.add_command('db', MigrateCommand)
manager.add_command('routes', ShowUrls)

@manager.command
def collections():
    """List all collections."""
    for collection in Collection.all():
        print collection.id, collection.foreign_id, collection.label

@manager.command
def alerts():
    """Generate alert notifications."""
    check_alerts.delay()

@manager.command
def crawl(name):
    """Execute the given crawler."""
    log.info('Crawling %r...', name)
    crawlers = get_crawlers()
    if name not in crawlers:
        log.info('No such crawler: %r', name)
    else:
        crawler = crawlers.get(name)()
        crawler.execute()
    db.session.commit()

@manager.command
@manager.option('-l', '--language', dest='language', nargs='*')
@manager.option('-c', '--country', dest='country', nargs='*')
@manager.option('-f', '--foreign_id', dest='foreign_id')
def crawldir(directory, language=None, country=None, foreign_id=None):
    """Crawl the given directory."""
    if directory is None or not os.path.exists(directory):
        log.error("Invalid directory: %r", directory)
        return
    directory = os.path.abspath(os.path.normpath(directory))

    if foreign_id is None:
        foreign_id = 'directory:%s' % slugify(directory)
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        collection = Collection.create({
            'foreign_id': foreign_id,
            'label': directory,
            'managed': True
        })

    if language is not None:
        collection.languages = [language]
    if country is not None:
        collection.countries = [country]
    db.session.commit()
    update_collection(collection)

    log.info('Crawling %r to %r...', directory, collection.foreign_id)
    document = Document.by_keys(collection_id=collection.id,
                                foreign_id=directory)
    document.collection = collection
    ingest_document(document, directory)

@manager.command
def flush(foreign_id):
    """Reset the crawler state for a given collecton."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    delete_collection(collection.id)

@manager.command
def deletepending(foreign_id=None):
    """Deletes any pending entities and related items."""
    collection_id = None
    if foreign_id is None:
        collection = Collection.by_foreign_id(foreign_id)
        if collection is None:
            raise ValueError("No such collection: %r" % foreign_id)
        collection_id = collection.id
    delete_pending(collection_id=collection_id)

@manager.command
def analyze(foreign_id):
    """Re-analyze documents in the given collection (or throughout)."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    analyze_collection.delay(collection.id)

@manager.command
def reingest(foreign_id):
    """Re-ingest documents in the given collection."""
    collection = Collection.by_foreign_id(foreign_id)
    if collection is None:
        raise ValueError("No such collection: %r" % foreign_id)
    reingest_collection(collection)

@manager.command
def index(foreign_id=None):
    """Index documents in the given collection (or throughout)."""
    q = Document.all_ids()
    # re-index newest document first.
    q = q.order_by(Document.id.desc())
    if foreign_id:
        collection = Collection.by_foreign_id(foreign_id)
        if collection is None:
            raise ValueError("No such collection: %r" % foreign_id)
        q = q.filter(Document.collection_id == collection.id)
    for idx, (doc_id,) in enumerate(q.yield_per(10000), 1):
        index_document_id.delay(doc_id)
        if idx % 1000 == 0:
            log.info("Index: %s documents...", idx)
    if foreign_id is None:
        reindex_entities()

@manager.command
def loaddataset(name):
    """Index all the entities in a given dataset."""
    dataset = datasets.get(name)
    load_dataset(dataset)

@manager.command
def deletedataset(name):
    delete_dataset(name)

@manager.command
def resetindex():
    """Re-create the ES index configuration, dropping all data."""
    delete_index()
    init_search()

@manager.command
def indexentities():
    """Re-index all the entities."""
    reindex_entities()

@manager.command
def updateentities():
    """Re-index all the entities."""
    q = db.session.query(Entity.id)
    for (entity_id,) in q:
        update_entity_full.delay(entity_id)

@manager.command
@manager.option('-s', '--skip-downloads', dest='skip', default='')
def init(skip=''):
    """Create or upgrade the search index and database."""
    upgrade_db()
    init_search()
    upgrade_search()
    if 'analyzers' not in skip:
        install_analyzers()
    archive.upgrade()

@manager.command
def upgrade():
    """Create or upgrade the search index and database."""
    upgrade_db()
    upgrade_search()
    archive.upgrade()

@manager.command
def installdata():
    """Create or upgrade the search index and database."""
    install_analyzers()

@manager.command
def evilshit():
    """EVIL: Delete all data and recreate the database."""
    delete_index()
    from sqlalchemy import MetaData, inspect
    from sqlalchemy.exc import InternalError
    from sqlalchemy.dialects.postgresql import ENUM
    metadata = MetaData()
    metadata.bind = db.engine
    metadata.reflect()
    tables = list(metadata.sorted_tables)
    while len(tables):
        for table in tables:
            try:
                table.drop(checkfirst=True)
                tables.remove(table)
            except InternalError:
                pass
    for enum in inspect(db.engine).get_enums():
        enum = ENUM(name=enum['name'])
        enum.drop(bind=db.engine, checkfirst=True)
    init()

def main():
    manager.run()

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


<commit_before>#!/usr/bin/python

##################################################################

# alerts_slack.py by Jeff Johnston <jj358mhz@gmail.com>
# Alerting script (Slack) for Scanner Listeners
# THIS FILE: /usr/local/bin/alerts_slack.py
# No Warranty is implied, promised or permitted.

##################################################################

import json
import requests
import urllib2

# The Broadcastify endpoint URL ** DO NOT ALTER **
broadcastifyURL = 'https://api.broadcastify.com/owner/?a=feed&feedId='

# Enter the account data for your Broadcastify feed
feedID = ''  # ENTER YOUR BROADCASTIFY FEED ID HERE
username = ''  # ENTER YOUR BROADCASTIFY USERNAME HERE
password = ''  # ENTER YOUR BROADCASTIFY PASSWORD HERE

# This threshold amount is the number of listeners that need to be exceeded before Slack alerts are sent out
alertThreshold = 20  # ENTER YOUR DESIRED ALERT LISTENER THRESHOLD HERE

# The Slack endpoint URL
webhook_url = ''  # ENTER YOUR SLACK WEBHOOK URL HERE

url = broadcastifyURL + feedID + '&type=json&u=' + username + '&p=' + password
response = urllib2.urlopen(url)
data = json.load(response)
descr = data['Feed'][0]['descr']
listeners = data['Feed'][0]['listeners']
status = data['Feed'][0]['status']

slack_payload = {"text": "*{} Broadcastify Alert* :ghost:\n"
                         "Listener threshold *{}* exceeded. Listeners = *{}*\n"
                         "\n"
                         "Listen to the feed here: <http://www.broadcastify.com/listen/feed/{}>\n"
                         "Manage the feed here: <http://www.broadcastify.com/manage/feed/{}>".format(descr,
                                                                                alertThreshold,
                                                                                listeners, feedID,
                                                                                feedID)}

slack_payload_feed_down = {"text": "*{} Broadcastify Alert* :ghost:\n"
                         "*FEED IS DOWN*\n"
                         "\n"
                         "Manage the feed here: <http://www.broadcastify.com/manage/feed/{}>".format(descr, feedID)}

def slack_alert(slack_payload, webhook_url):
    response = requests.post(webhook_url, data=json.dumps(slack_payload), headers={'Content-Type': 'application/json'})

    if response.status_code != 200:
        raise ValueError(
            'Request to Slack returned an error %s, the response is:\n'
            '%s' % (response.status_code, response.text))

if status == False:
    slack_alert(slack_payload_feed_down, webhook_url)
else:
    if listeners >= alertThreshold:
        slack_alert(slack_payload, webhook_url)
<commit_msg>Add emoji & code format<commit_after>#!/usr/bin/python

##################################################################

# alerts_slack.py by Jeff Johnston <jj358mhz@gmail.com>
# Alerting script (Slack) for Scanner Listeners
# THIS FILE: /usr/local/bin/alerts_slack.py
# No Warranty is implied, promised or permitted.

##################################################################

import json
import requests
import urllib2

# The Broadcastify endpoint URL ** DO NOT ALTER **
broadcastifyURL = 'https://api.broadcastify.com/owner/?a=feed&feedId='

# Enter the account data for your Broadcastify feed
feedID = ''  # ENTER YOUR BROADCASTIFY FEED ID HERE
username = ''  # ENTER YOUR BROADCASTIFY USERNAME HERE
password = ''  # ENTER YOUR BROADCASTIFY PASSWORD HERE

# This threshold amount is the number of listeners that need to be exceeded before Slack alerts are sent out
alertThreshold = 20  # ENTER YOUR DESIRED ALERT LISTENER THRESHOLD HERE

# The Slack endpoint URL
webhook_url = ''  # ENTER YOUR SLACK WEBHOOK URL HERE

url = broadcastifyURL + feedID + '&type=json&u=' + username + '&p=' + password
response = urllib2.urlopen(url)
data = json.load(response)
descr = data['Feed'][0]['descr']
listeners = data['Feed'][0]['listeners']
status = data['Feed'][0]['status']

slack_payload = {"text": "*{} Broadcastify Alert* :cop::fire:\n"
                         "Listener threshold *{}* exceeded. Listeners = *{}*\n"
                         "\n"
                         "Listen to the feed here: <http://www.broadcastify.com/listen/feed/{}>\n"
                         "Manage the feed here: <http://www.broadcastify.com/manage/feed/{}>".format(descr,
                                                                                alertThreshold,
                                                                                listeners, feedID,
                                                                                feedID)}

slack_payload_feed_down = {"text": "*{} Broadcastify Alert* :ghost:\n"
                                   "*FEED IS DOWN*\n"
                                   "\n"
                                   "Manage the feed here: <http://www.broadcastify.com/manage/feed/{}>".format(descr,
                                                                                feedID)}

def slack_alert(slack_payload, webhook_url):
    response = requests.post(webhook_url, data=json.dumps(slack_payload), headers={'Content-Type': 'application/json'})

    if response.status_code != 200:
        raise ValueError(
            'Request to Slack returned an error %s, the response is:\n'
            '%s' % (response.status_code, response.text))

if status == False:
    slack_alert(slack_payload_feed_down, webhook_url)
else:
    if listeners >= alertThreshold:
        slack_alert(slack_payload, webhook_url)
<|endoftext|>