<commit_before>#!/usr/bin/env python
########################################
#
# alert-query.py - Alert query tool
#
########################################

import os
import sys
from optparse import OptionParser
import datetime
import pymongo
import operator
import pytz

__version__ = '1.0.0'

SEV = {
    'CRITICAL': 'Crit',
    'MAJOR':    'Majr',
    'MINOR':    'Minr',
    'WARNING':  'Warn',
    'NORMAL':   'Norm',
    'INFORM':   'Info',
    'DEBUG':    'Dbug',
}

SEVERITY_CODE = {
    # ITU RFC5674 -> Syslog RFC5424
    'CRITICAL':       1, # Alert
    'MAJOR':          2, # Crtical
    'MINOR':          3, # Error
    'WARNING':        4, # Warning
    'NORMAL':         5, # Notice
    'INFORM':         6, # Informational
    'DEBUG':          7, # Debug
}

COLOR = {
    'CRITICAL': '\033[91m',
    'MAJOR':    '\033[95m',
    'MINOR':    '\033[93m',
    'WARNING':  '\033[96m',
    'NORMAL':   '\033[92m',
    'INFORM':   '\033[92m',
    'DEBUG':    '\033[90m',
}
ENDC     = '\033[0m'

# Defaults
PGM='alert-query.py'
SERVER='localhost'
TIMEZONE='Europe/London'
DATE_FORMAT = '%d/%m/%y %H:%M:%S'

def main():

    # Command-line options
    parser = OptionParser(
                      version="%prog " + __version__,
                      description="Alert database query tool - show alerts filtered by attributes",
                      epilog="alert-query.py --color --env 'QA|REL' --group Puppet --count 10 --show all")
    parser.add_option("-m",
                      "--server",
                      dest="server",
                      help="MongoDB server (default: localhost)")
    parser.add_option("-z",
                      "--timezone",
                      dest="timezone",
                      help="Set timezone (default: Europe/London)")
    parser.add_option("--minutes",
                      "--mins",
                      type="int",
                      dest="minutes",
                      help="Show alerts for last <x> minutes")
    parser.add_option("--hours",
                      "--hrs",
                      type="int",
                      dest="hours",
                      help="Show alerts for last <x> hours")
    parser.add_option("--days",
                      type="int",
                      dest="days",
                      help="Show alerts for last <x> days")
    parser.add_option("-i",
                      "--id",
                      dest="alertid",
                      help="Alert ID (can use 8-char abbrev id)")
    parser.add_option("-E",
                      "--environment",
                      dest="environment",
                      help="Environment eg. PROD, REL, QA, TEST, CODE, STAGE, DEV, LWP, INFRA")
    parser.add_option("-S",
                      "--svc",
                      "--service",
                      dest="service",
                      help="Service eg. R1, R2, Discussion, Soulmates, ContentAPI, MicroApp, FlexibleContent, Mutualisation, SharedSvcs")
    parser.add_option("-r",
                      "--resource",
                      dest="resource",
                      help="Resource under alarm eg. hostname, network device, application")
    parser.add_option("-s",
                      "--severity",
                      dest="severity",
                      help="Severity or range eg. major, critical..warning")
    parser.add_option("-e",
                      "--event",
                      dest="event",
                      help="Event name eg. HostAvail, PingResponse, AppStatus")
    parser.add_option("-g",
                      "--group",
                      dest="group",
                      help="Event group eg. Application, Backup, Database, HA, Hardware, Job, Network, OS, Performance, Security")
    parser.add_option("-v",
                      "--value",
                      dest="value",
                      help="Event value eg. 100%, Down, PingFail, 55tps, ORA-1664")
    parser.add_option("-t",
                      "--text",
                      dest="text")
    parser.add_option("--show",
                      action="append",
                      dest="show",
                      default=[],
                      help="Show 'text', 'times', 'details', 'tags' and 'color'")
    parser.add_option("-c",
                      "--count",
                      "--limit",
                      type="int",
                      dest="limit",
                      default=0)
    parser.add_option( "--no-header",
                      action="store_true",
                      dest="noheader")
    parser.add_option( "--no-footer",
                      action="store_true",
                      dest="nofooter")
    parser.add_option("--color",
                      "--colour",
                      action="store_true",
                      default=False,
                      help="Synonym for --show=color")
    parser.add_option("-d",
                      "--dry-run",
                      action="store_true",
                      default=False,
                      help="Do not run. Output query and filter.")
    options, args = parser.parse_args()

    if options.server:
        server = options.server
    else:
        server = SERVER

    # Connect to MongoDB
    try:
        mongo = pymongo.Connection(server)
        db = mongo.monitoring
        alerts = db.alerts
    except pymongo.errors.ConnectionFailure, e:
        print >>sys.stderr, "ERROR: Connection to MongoDB %s failed" % server
        sys.exit(1)

    query = dict()
    minutes = 0
    hours = 0
    days = 0
    num_regex = 0
    if options.minutes:
        minutes = options.minutes
    if options.hours:
        hours = options.minutes
    if options.days:
        days = options.days
    if options.alertid:
        query['history.id'] = { '$regex': '^'+options.alertid }
        num_regex += 1
    if options.environment:
        query['environment'] = { '$regex': options.environment }
        num_regex += 1
    if options.service:
        query['service'] = { '$regex': options.service }
        num_regex += 1
    if options.resource:
        query['resource'] = { '$regex': options.resource }
        num_regex += 1
    if options.severity:
        m = options.severity.split('..')
        if len(m) > 1:
            query['severityCode'] = { '$gte': SEVERITY_CODE[m[0].upper()], '$lte': SEVERITY_CODE[m[1].upper()] }
        else:
            query['severityCode'] = SEVERITY_CODE[options.severity.upper()]
    if options.event:
        query['event'] = { '$regex': options.event }
        num_regex += 1
    if options.group:
        query['group'] = { '$regex': options.group }
        num_regex += 1
    if options.value:
        query['history.value'] = { '$regex': options.value }
        num_regex += 1
    if options.text:
        query['history.text'] = { '$regex': options.text }
        num_regex += 1

    if num_regex > 4:
        print "ERROR: Too many regexes in query (limit is 4 until v2.1.0)"
        # see https://jira.mongodb.org/browse/SERVER-969
        sys.exit(1)

    if minutes or hours or days:
        end = datetime.datetime.utcnow()
        end = end.replace(tzinfo=pytz.utc)
        start = end - datetime.timedelta(days=days, minutes=minutes+hours*60)
        start = start.replace(tzinfo=pytz.utc)
        query['createTime'] = {'$gte': start, '$lt': end}

    fields = { "environment": 1, "service": 1,
        "resource": 1, "event": 1, "group": 1,
        "type": 1, "tags": 1, "origin": 1,
        "thresholdInfo": 1 }
    if 'history' in options.show:
        fields['history'] = 1
    else:
        fields['history'] = { '$slice': -1 }

    if options.limit:
        LIMIT = options.limit
    else:
        LIMIT = 0

    orderby = "'history.createTime': pymongo.DESCENDING"
    # TODO - make orderby user-configurable

    if options.dry_run:
        print "DEBUG: monitoring.alerts { $query: %s, $orderby: { %s } }" % (query, orderby)
        sys.exit(0)

    results = list()
    for alert in alerts.find(query, fields).sort(orderby).limit(LIMIT):
        for hist in alert['history']:
            results.append((alert['_id'],
                hist['id'],
                alert['environment'],
                alert['service'],
                alert['resource'],
                alert['event'],
                alert['group'],
                hist.get('value', 'n/a'),
                hist['severity'],
                hist['text'],
                alert['type'],
                alert['tags'],
                hist['createTime'].replace(tzinfo=pytz.utc),
                hist['receiveTime'].replace(tzinfo=pytz.utc),
                alert['origin'],
                alert.get('thresholdInfo', 'n/a')))

    if not options.timezone:
        user_tz = TIMEZONE
    else:
        user_tz = options.timezone
    tz=pytz.timezone(user_tz)

    if not options.noheader:
        print "Alerta Report Tool"
        print "    database: %s" % server
        print "    timezone: %s" % user_tz
        if minutes or hours or days:
            print "    interval: %s - %s" % (start.astimezone(tz).strftime(DATE_FORMAT), end.astimezone(tz).strftime(DATE_FORMAT))
        if options.show:
            print "        show: %s" % ', '.join(options.show)
        if options.alertid:
            print "    alert id: ^%s" % options.alertid
        if options.environment:
            print " environment: %s" % options.environment
        if options.service:
            print "     service: %s" % options.service
        if options.resource:
            print "    resource: %s" % options.resource
        if options.severity:
            print "    severity: %s" % options.severity
        if options.event:
            print "       event: %s" % options.event
        if options.group:
            print "       group: %s" % options.group
        if options.value:
            print "       value: %s" % options.value
        if options.text:
            print "        text: %s" % options.text
        if options.limit:
            print "       count: %d" % LIMIT
        print

    if 'some' in options.show:
        options.show.append('text')
        options.show.append('details')
    elif 'all' in options.show:
        options.show.append('text')
        options.show.append('times')
        options.show.append('details')
        options.show.append('tags')

    line_color = ''
    end_color = ''
    if 'color' in options.show or options.color:
        end_color = ENDC

    count = 0
    results = sorted(results, key=operator.itemgetter(12)) # sort by createTime ie. results[12]
    for r in results:
        objectid      = r[0]
        alertid       = r[1]
        environment   = r[2]
        service       = r[3]
        resource      = r[4]
        event         = r[5]
        group         = r[6]
        value         = r[7]
        severity      = r[8]
        severityCode  = SEVERITY_CODE[severity]
        text          = r[9]
        type          = r[10]
        tags          = r[11]
        createTime    = r[12]
        receiveTime   = r[13]
        latency       = receiveTime - createTime
        origin        = r[14]
        thresholdInfo = r[15]
        count += 1

        if 'color' in options.show or options.color:
            line_color = COLOR[severity]
        print(line_color + '%s|%s|%s|%-15s|%12s|%11s|%12s' % (alertid[0:8],
            createTime.astimezone(tz).strftime(DATE_FORMAT),
            SEV[severity],
            resource.split('.')[-1],
            group,
            event,
            value) + end_color)

        if 'text' in options.show:
            print(line_color + '   |%s' % (text) + end_color)

        if 'times' in options.show:
            print(line_color + '    time created  | %s' % (createTime.astimezone(tz).strftime(DATE_FORMAT)) + end_color)
            print(line_color + '    time received | %s' % (receiveTime.astimezone(tz).strftime(DATE_FORMAT)) + end_color)
            print(line_color + '    latency       | %s' % (latency) + end_color)

        if 'details' in options.show:
            print(line_color + '        object id   | %s' % (objectid) + end_color)
            print(line_color + '        alert id    | %s' % (alertid) + end_color)
            print(line_color + '        environment | %s' % (environment) + end_color)
            print(line_color + '        service     | %s' % (service) + end_color)
            print(line_color + '        resource    | %s' % (resource) + end_color)
            print(line_color + '        type        | %s' % (type) + end_color)
            print(line_color + '        origin      | %s' % (origin) + end_color)
            print(line_color + '        threshold   | %s' % (thresholdInfo) + end_color)

        if 'tags' in options.show and tags:
            for t in tags:
                print(line_color + '            tag | %s' % (t) + end_color)

    if not options.nofooter:
        now = datetime.datetime.utcnow()
        print
        print "Total: %d (produced on %s at %s by %s,%s on %s)" % (count, now.strftime("%d/%m/%y"), now.strftime("%H:%M:%S"), PGM, __version__, os.uname()[1])

if __name__ == '__main__':
    main()
<commit_msg>Adjust column widths<commit_after>#!/usr/bin/env python
########################################
#
# alert-query.py - Alert query tool
#
########################################

import os
import sys
from optparse import OptionParser
import datetime
import pymongo
import operator
import pytz

__version__ = '1.0.1'

SEV = {
    'CRITICAL': 'Crit',
    'MAJOR':    'Majr',
    'MINOR':    'Minr',
    'WARNING':  'Warn',
    'NORMAL':   'Norm',
    'INFORM':   'Info',
    'DEBUG':    'Dbug',
}

SEVERITY_CODE = {
    # ITU RFC5674 -> Syslog RFC5424
    'CRITICAL':       1, # Alert
    'MAJOR':          2, # Crtical
    'MINOR':          3, # Error
    'WARNING':        4, # Warning
    'NORMAL':         5, # Notice
    'INFORM':         6, # Informational
    'DEBUG':          7, # Debug
}

COLOR = {
    'CRITICAL': '\033[91m',
    'MAJOR':    '\033[95m',
    'MINOR':    '\033[93m',
    'WARNING':  '\033[96m',
    'NORMAL':   '\033[92m',
    'INFORM':   '\033[92m',
    'DEBUG':    '\033[90m',
}
ENDC     = '\033[0m'

# Defaults
PGM='alert-query.py'
SERVER='localhost'
TIMEZONE='Europe/London'
DATE_FORMAT = '%d/%m/%y %H:%M:%S'

def main():

    # Command-line options
    parser = OptionParser(
                      version="%prog " + __version__,
                      description="Alert database query tool - show alerts filtered by attributes",
                      epilog="alert-query.py --color --env 'QA|REL' --group Puppet --count 10 --show all")
    parser.add_option("-m",
                      "--server",
                      dest="server",
                      help="MongoDB server (default: localhost)")
    parser.add_option("-z",
                      "--timezone",
                      dest="timezone",
                      help="Set timezone (default: Europe/London)")
    parser.add_option("--minutes",
                      "--mins",
                      type="int",
                      dest="minutes",
                      help="Show alerts for last <x> minutes")
    parser.add_option("--hours",
                      "--hrs",
                      type="int",
                      dest="hours",
                      help="Show alerts for last <x> hours")
    parser.add_option("--days",
                      type="int",
                      dest="days",
                      help="Show alerts for last <x> days")
    parser.add_option("-i",
                      "--id",
                      dest="alertid",
                      help="Alert ID (can use 8-char abbrev id)")
    parser.add_option("-E",
                      "--environment",
                      dest="environment",
                      help="Environment eg. PROD, REL, QA, TEST, CODE, STAGE, DEV, LWP, INFRA")
    parser.add_option("-S",
                      "--svc",
                      "--service",
                      dest="service",
                      help="Service eg. R1, R2, Discussion, Soulmates, ContentAPI, MicroApp, FlexibleContent, Mutualisation, SharedSvcs")
    parser.add_option("-r",
                      "--resource",
                      dest="resource",
                      help="Resource under alarm eg. hostname, network device, application")
    parser.add_option("-s",
                      "--severity",
                      dest="severity",
                      help="Severity or range eg. major, critical..warning")
    parser.add_option("-e",
                      "--event",
                      dest="event",
                      help="Event name eg. HostAvail, PingResponse, AppStatus")
    parser.add_option("-g",
                      "--group",
                      dest="group",
                      help="Event group eg. Application, Backup, Database, HA, Hardware, Job, Network, OS, Performance, Security")
    parser.add_option("-v",
                      "--value",
                      dest="value",
                      help="Event value eg. 100%, Down, PingFail, 55tps, ORA-1664")
    parser.add_option("-t",
                      "--text",
                      dest="text")
    parser.add_option("--show",
                      action="append",
                      dest="show",
                      default=[],
                      help="Show 'text', 'times', 'details', 'tags' and 'color'")
    parser.add_option("-c",
                      "--count",
                      "--limit",
                      type="int",
                      dest="limit",
                      default=0)
    parser.add_option( "--no-header",
                      action="store_true",
                      dest="noheader")
    parser.add_option( "--no-footer",
                      action="store_true",
                      dest="nofooter")
    parser.add_option("--color",
                      "--colour",
                      action="store_true",
                      default=False,
                      help="Synonym for --show=color")
    parser.add_option("-d",
                      "--dry-run",
                      action="store_true",
                      default=False,
                      help="Do not run. Output query and filter.")
    options, args = parser.parse_args()

    if options.server:
        server = options.server
    else:
        server = SERVER

    # Connect to MongoDB
    try:
        mongo = pymongo.Connection(server)
        db = mongo.monitoring
        alerts = db.alerts
    except pymongo.errors.ConnectionFailure, e:
        print >>sys.stderr, "ERROR: Connection to MongoDB %s failed" % server
        sys.exit(1)

    query = dict()
    minutes = 0
    hours = 0
    days = 0
    num_regex = 0
    if options.minutes:
        minutes = options.minutes
    if options.hours:
        hours = options.minutes
    if options.days:
        days = options.days
    if options.alertid:
        query['history.id'] = { '$regex': '^'+options.alertid }
        num_regex += 1
    if options.environment:
        query['environment'] = { '$regex': options.environment }
        num_regex += 1
    if options.service:
        query['service'] = { '$regex': options.service }
        num_regex += 1
    if options.resource:
        query['resource'] = { '$regex': options.resource }
        num_regex += 1
    if options.severity:
        m = options.severity.split('..')
        if len(m) > 1:
            query['severityCode'] = { '$gte': SEVERITY_CODE[m[0].upper()], '$lte': SEVERITY_CODE[m[1].upper()] }
        else:
            query['severityCode'] = SEVERITY_CODE[options.severity.upper()]
    if options.event:
        query['event'] = { '$regex': options.event }
        num_regex += 1
    if options.group:
        query['group'] = { '$regex': options.group }
        num_regex += 1
    if options.value:
        query['history.value'] = { '$regex': options.value }
        num_regex += 1
    if options.text:
        query['history.text'] = { '$regex': options.text }
        num_regex += 1

    if num_regex > 4:
        print "ERROR: Too many regexes in query (limit is 4 until v2.1.0)"
        # see https://jira.mongodb.org/browse/SERVER-969
        sys.exit(1)

    if minutes or hours or days:
        end = datetime.datetime.utcnow()
        end = end.replace(tzinfo=pytz.utc)
        start = end - datetime.timedelta(days=days, minutes=minutes+hours*60)
        start = start.replace(tzinfo=pytz.utc)
        query['createTime'] = {'$gte': start, '$lt': end}

    fields = { "environment": 1, "service": 1,
        "resource": 1, "event": 1, "group": 1,
        "type": 1, "tags": 1, "origin": 1,
        "thresholdInfo": 1 }
    if 'history' in options.show:
        fields['history'] = 1
    else:
        fields['history'] = { '$slice': -1 }

    if options.limit:
        LIMIT = options.limit
    else:
        LIMIT = 0

    orderby = "'history.createTime': pymongo.DESCENDING"
    # TODO - make orderby user-configurable

    if options.dry_run:
        print "DEBUG: monitoring.alerts { $query: %s, $orderby: { %s } }" % (query, orderby)
        sys.exit(0)

    results = list()
    for alert in alerts.find(query, fields).sort(orderby).limit(LIMIT):
        for hist in alert['history']:
            results.append((alert['_id'],
                hist['id'],
                alert['environment'],
                alert['service'],
                alert['resource'],
                alert['event'],
                alert['group'],
                hist.get('value', 'n/a'),
                hist['severity'],
                hist['text'],
                alert['type'],
                alert['tags'],
                hist['createTime'].replace(tzinfo=pytz.utc),
                hist['receiveTime'].replace(tzinfo=pytz.utc),
                alert['origin'],
                alert.get('thresholdInfo', 'n/a')))

    if not options.timezone:
        user_tz = TIMEZONE
    else:
        user_tz = options.timezone
    tz=pytz.timezone(user_tz)

    if not options.noheader:
        print "Alerta Report Tool"
        print "    database: %s" % server
        print "    timezone: %s" % user_tz
        if minutes or hours or days:
            print "    interval: %s - %s" % (start.astimezone(tz).strftime(DATE_FORMAT), end.astimezone(tz).strftime(DATE_FORMAT))
        if options.show:
            print "        show: %s" % ', '.join(options.show)
        if options.alertid:
            print "    alert id: ^%s" % options.alertid
        if options.environment:
            print " environment: %s" % options.environment
        if options.service:
            print "     service: %s" % options.service
        if options.resource:
            print "    resource: %s" % options.resource
        if options.severity:
            print "    severity: %s" % options.severity
        if options.event:
            print "       event: %s" % options.event
        if options.group:
            print "       group: %s" % options.group
        if options.value:
            print "       value: %s" % options.value
        if options.text:
            print "        text: %s" % options.text
        if options.limit:
            print "       count: %d" % LIMIT
        print

    if 'some' in options.show:
        options.show.append('text')
        options.show.append('details')
    elif 'all' in options.show:
        options.show.append('text')
        options.show.append('times')
        options.show.append('details')
        options.show.append('tags')

    line_color = ''
    end_color = ''
    if 'color' in options.show or options.color:
        end_color = ENDC

    count = 0
    results = sorted(results, key=operator.itemgetter(12)) # sort by createTime ie. results[12]
    for r in results:
        objectid      = r[0]
        alertid       = r[1]
        environment   = r[2]
        service       = r[3]
        resource      = r[4]
        event         = r[5]
        group         = r[6]
        value         = r[7]
        severity      = r[8]
        severityCode  = SEVERITY_CODE[severity]
        text          = r[9]
        type          = r[10]
        tags          = r[11]
        createTime    = r[12]
        receiveTime   = r[13]
        latency       = receiveTime - createTime
        origin        = r[14]
        thresholdInfo = r[15]
        count += 1

        if 'color' in options.show or options.color:
            line_color = COLOR[severity]
        print(line_color + '%s|%s|%s|%-18s|%12s|%16s|%12s' % (alertid[0:8],
            createTime.astimezone(tz).strftime(DATE_FORMAT),
            SEV[severity],
            resource.split('.')[-1],
            group,
            event,
            value) + end_color)

        if 'text' in options.show:
            print(line_color + '   |%s' % (text) + end_color)

        if 'times' in options.show:
            print(line_color + '    time created  | %s' % (createTime.astimezone(tz).strftime(DATE_FORMAT)) + end_color)
            print(line_color + '    time received | %s' % (receiveTime.astimezone(tz).strftime(DATE_FORMAT)) + end_color)
            print(line_color + '    latency       | %s' % (latency) + end_color)

        if 'details' in options.show:
            print(line_color + '        object id   | %s' % (objectid) + end_color)
            print(line_color + '        alert id    | %s' % (alertid) + end_color)
            print(line_color + '        environment | %s' % (environment) + end_color)
            print(line_color + '        service     | %s' % (service) + end_color)
            print(line_color + '        resource    | %s' % (resource) + end_color)
            print(line_color + '        type        | %s' % (type) + end_color)
            print(line_color + '        origin      | %s' % (origin) + end_color)
            print(line_color + '        threshold   | %s' % (thresholdInfo) + end_color)

        if 'tags' in options.show and tags:
            for t in tags:
                print(line_color + '            tag | %s' % (t) + end_color)

    if not options.nofooter:
        now = datetime.datetime.utcnow()
        print
        print "Total: %d (produced on %s at %s by %s,v%s on %s)" % (count, now.strftime("%d/%m/%y"), now.strftime("%H:%M:%S"), PGM, __version__, os.uname()[1])

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


<commit_before>import json
import logging
from copy import deepcopy

import duckdb
import requests
from cloudpathlib import AnyPath
from did_resolver import Resolver
from python_did_resolver import get_resolver

from .utils import get_files, get_latest_details, run_job

class RewardProgram:

    processed_cycles = set()
    last_update_block = 0

    def __init__(
        self,
        reward_program_id,
        reward_manager,
        subgraph_url,
        result_file_root,
        subgraph_extract_location=None,
    ) -> None:
        self.reward_manager = reward_manager
        self.reward_program_id = reward_program_id
        self.subgraph_url = subgraph_url
        self.reward_program_output_location = AnyPath(result_file_root).joinpath(
            f"rewardProgramID={reward_program_id}"
        )
        if subgraph_extract_location is not None:
            self.load_submitted_from_subgraph_extraction(subgraph_extract_location)
        self.update_processed()

    def load_submitted_from_subgraph_extraction(self, subgraph_extract_location):
        logging.info(
            f"Loading processed payment cycles for {self.reward_program_id} from {subgraph_extract_location}"
        )
        subgraph_export_files = get_files(
            subgraph_extract_location, "merkle_root_submission"
        )
        con = duckdb.connect(database=":memory:", read_only=False)
        con.execute(
            f"""
        select distinct payment_cycle_uint64 from parquet_scan({subgraph_export_files})
        """
        )
        self.processed_cycles.update(r[0] for r in con.fetchall())
        con.execute(
            f"select max(_block_number) from parquet_scan({subgraph_export_files})"
        )
        self.last_update_block = con.fetchall()[0][0]

    def update_processed(self):
        """
        Update the set of processed payment cycles from subgraph
        """
        logging.info(
            f"Updating processed payment cycles for {self.reward_program_id} from subgraph"
        )
        # graphql check
        query = f"""query {{
            merkleRootSubmissions(orderBy:blockNumber, orderDirection:asc,
                where:{{rewardProgram:"{self.reward_program_id}", blockNumber_gt:{self.last_update_block}}}) {{
                blockNumber
                paymentCycle
            }}
        }}"""
        try:
            r = requests.post(
                self.subgraph_url,
                json={"query": query},
            )
            if r.ok:
                json_data = r.json()
                for submission in json_data["data"]["merkleRootSubmissions"]:
                    self.processed_cycles.add(int(submission["paymentCycle"]))
                    self.last_update_block = max(
                        self.last_update_block, int(submission["blockNumber"])
                    )
            else:
                raise (r.raise_for_status())
        except requests.exceptions.ConnectionError:
            logging.warn("Connection error during query subgraph")
        except Exception as e:
            logging.warn("Error when querying subgraph")
            raise (e)

    def get_latest_data_block(self, rule):
        """
        Get the latest block that has data for this program
        based on the subgraph exports
        """
        latest_block = None
        for subgraph_config in rule["subgraph_config_locations"].values():
            latest_data = get_latest_details(subgraph_config)
            if latest_block is None:
                latest_block = latest_data["latest_block"]
            else:
                latest_block = min(latest_block, latest_data["latest_block"])
        return latest_block

    def is_locked(self):
        return self.reward_manager.caller.rewardProgramLocked(self.reward_program_id)

    def get_all_payment_cycles(self, rule):
        """
        Return a set of payment cycles that are valid for this program
        where the data is available
        """
        core_config = rule["core"]
        self.update_processed()
        start_block = core_config["start_block"]
        latest_data_block = self.get_latest_data_block(core_config)
        # If there is no dependency on data then all cycles from the start
        # to end block are valid to process
        if latest_data_block is None:
            end_block = core_config["end_block"]
        else:
            end_block = min(core_config["end_block"], latest_data_block)
        payment_cycle_length = core_config["payment_cycle_length"]
        return set(range(start_block, end_block, payment_cycle_length))

    def get_rules(self):
        blob = self.reward_manager.caller.rule(self.reward_program_id)
        if blob and blob != b"":
            try:
                did = blob.decode("utf-8")  # new blob format: hex encodes a did string
                rules = resolve_rule(did)
            except Exception:
                # old blob format: our rule blobs were hex encoded json
                # try..except maintains backward compatibality with our old blob format
                # TODO: remove code once all reward programs have the old rule /blob removed
                rules = json.loads(blob)
            if type(rules) == list:
                return rules
            else:
                return [rules]
        else:
            return []

    def raise_on_payment_cycle_overlap(self, rules):
        """
        Given a list of rules, raise an exception if there are overlapping
        payment cycles.

        For example, with two rules
        1. A rule that starts on block 0, ends on block 200 and runs every 10 blocks
        2. A rule that starts on block 20, ends on block 70 and runs every 20 blocks

        There would be a clash for payment cycles 20, 40, 60
        """
        current_cycles = set()
        for rule in rules:
            rule_cycles = self.get_all_payment_cycles(rule)
            if rule_cycles.intersection(current_cycles):
                raise Exception(
                    f"Reward program {self.reward_program_id} has overlapping payment cycles {rule_cycles.intersection(current_cycles)}"
                )
            current_cycles.update(rule_cycles)

    def run(self, payment_cycle, rule):
        """
        Trigger a job in AWS batch for a single payment cycle and single rule
        """
        logging.info(f"Running {payment_cycle} for {self.reward_program_id}")
        submission_data = deepcopy(rule)
        docker_image = submission_data["core"]["docker_image"]
        submission_data["run"] = {
            "reward_program_id": self.reward_program_id,
            "payment_cycle": payment_cycle,
        }
        payment_cycle_output = self.reward_program_output_location.joinpath(
            f"paymentCycle={payment_cycle}"
        )

        if payment_cycle_output.joinpath("results.parquet").exists():
            logging.info(
                f"Already processed: {payment_cycle} for {self.reward_program_id}"
            )
            return None

        # We can't pass much data directly in AWS batch, so write the
        # parameters we want to use to S3 and pass the location into the task
        parameters_location = payment_cycle_output.joinpath("parameters.json")
        with parameters_location.open("w") as params_out:
            json.dump(submission_data, params_out)

        job = run_job(
            docker_image,
            parameters_location.as_uri(),
            payment_cycle_output.as_uri(),
            tags={
                "reward_program_id": self.reward_program_id,
                "payment_cycle": payment_cycle,
            },
        )
        return job

    def run_all_payment_cycles(self) -> None:
        """
        Run all rules for all payment cycles that can be processed at this time
        """
        logging.info(f"Running all for {self.reward_program_id}")
        if self.is_locked():
            logging.info(f"Reward program {self.reward_program_id} is locked, skipping")
            return
        rules = self.get_rules()
        self.raise_on_payment_cycle_overlap(rules)
        for rule in rules:
            processable_payment_cycles = (
                self.get_all_payment_cycles(rule) - self.processed_cycles
            )
            for payment_cycle in sorted(processable_payment_cycles):
                self.run(payment_cycle, rule)

def resolve_rule(did: str):
    url = Resolver(get_resolver()).resolve(did)["didDocument"]["alsoKnownAs"][0]
    return requests.get(url).json()
<commit_msg>Handle case when there are no merkle submissions in the parquet data (#3181)<commit_after>import json
import logging
from copy import deepcopy

import duckdb
import requests
from cloudpathlib import AnyPath
from did_resolver import Resolver
from python_did_resolver import get_resolver

from .utils import get_files, get_latest_details, run_job

class RewardProgram:

    processed_cycles = set()
    last_update_block = 0

    def __init__(
        self,
        reward_program_id,
        reward_manager,
        subgraph_url,
        result_file_root,
        subgraph_extract_location=None,
    ) -> None:
        self.reward_manager = reward_manager
        self.reward_program_id = reward_program_id
        self.subgraph_url = subgraph_url
        self.reward_program_output_location = AnyPath(result_file_root).joinpath(
            f"rewardProgramID={reward_program_id}"
        )
        if subgraph_extract_location is not None:
            self.load_submitted_from_subgraph_extraction(subgraph_extract_location)
        self.update_processed()

    def load_submitted_from_subgraph_extraction(self, subgraph_extract_location):
        logging.info(
            f"Loading processed payment cycles for {self.reward_program_id} from {subgraph_extract_location}"
        )
        subgraph_export_files = get_files(
            subgraph_extract_location, "merkle_root_submission"
        )
        con = duckdb.connect(database=":memory:", read_only=False)
        con.execute(
            f"""
        select distinct payment_cycle_uint64 from parquet_scan({subgraph_export_files})
        """
        )
        self.processed_cycles.update(r[0] for r in con.fetchall())
        con.execute(
            f"select max(_block_number), count(*) from parquet_scan({subgraph_export_files})"
        )
        last_update_block, total_payment_cycles = con.fetchall()[0]
        if total_payment_cycles == 0:
            # There are no payment cycles in the parquet files
            self.last_update_block = 0
        else:
            self.last_update_block = last_update_block

    def update_processed(self):
        """
        Update the set of processed payment cycles from subgraph
        """
        logging.info(
            f"Updating processed payment cycles for {self.reward_program_id} from subgraph"
        )
        # graphql check
        query = f"""query {{
            merkleRootSubmissions(orderBy:blockNumber, orderDirection:asc,
                where:{{rewardProgram:"{self.reward_program_id}", blockNumber_gt:{self.last_update_block}}}) {{
                blockNumber
                paymentCycle
            }}
        }}"""
        try:
            r = requests.post(
                self.subgraph_url,
                json={"query": query},
            )
            if r.ok:
                json_data = r.json()
                for submission in json_data["data"]["merkleRootSubmissions"]:
                    self.processed_cycles.add(int(submission["paymentCycle"]))
                    self.last_update_block = max(
                        self.last_update_block, int(submission["blockNumber"])
                    )
            else:
                raise (r.raise_for_status())
        except requests.exceptions.ConnectionError:
            logging.warn("Connection error during query subgraph")
        except Exception as e:
            logging.warn("Error when querying subgraph")
            raise (e)

    def get_latest_data_block(self, rule):
        """
        Get the latest block that has data for this program
        based on the subgraph exports
        """
        latest_block = None
        for subgraph_config in rule["subgraph_config_locations"].values():
            latest_data = get_latest_details(subgraph_config)
            if latest_block is None:
                latest_block = latest_data["latest_block"]
            else:
                latest_block = min(latest_block, latest_data["latest_block"])
        return latest_block

    def is_locked(self):
        return self.reward_manager.caller.rewardProgramLocked(self.reward_program_id)

    def get_all_payment_cycles(self, rule):
        """
        Return a set of payment cycles that are valid for this program
        where the data is available
        """
        core_config = rule["core"]
        self.update_processed()
        start_block = core_config["start_block"]
        latest_data_block = self.get_latest_data_block(core_config)
        # If there is no dependency on data then all cycles from the start
        # to end block are valid to process
        if latest_data_block is None:
            end_block = core_config["end_block"]
        else:
            end_block = min(core_config["end_block"], latest_data_block)
        payment_cycle_length = core_config["payment_cycle_length"]
        return set(range(start_block, end_block, payment_cycle_length))

    def get_rules(self):
        blob = self.reward_manager.caller.rule(self.reward_program_id)
        if blob and blob != b"":
            try:
                did = blob.decode("utf-8")  # new blob format: hex encodes a did string
                rules = resolve_rule(did)
            except Exception:
                # old blob format: our rule blobs were hex encoded json
                # try..except maintains backward compatibality with our old blob format
                # TODO: remove code once all reward programs have the old rule /blob removed
                rules = json.loads(blob)
            if type(rules) == list:
                return rules
            else:
                return [rules]
        else:
            return []

    def raise_on_payment_cycle_overlap(self, rules):
        """
        Given a list of rules, raise an exception if there are overlapping
        payment cycles.

        For example, with two rules
        1. A rule that starts on block 0, ends on block 200 and runs every 10 blocks
        2. A rule that starts on block 20, ends on block 70 and runs every 20 blocks

        There would be a clash for payment cycles 20, 40, 60
        """
        current_cycles = set()
        for rule in rules:
            rule_cycles = self.get_all_payment_cycles(rule)
            if rule_cycles.intersection(current_cycles):
                raise Exception(
                    f"Reward program {self.reward_program_id} has overlapping payment cycles {rule_cycles.intersection(current_cycles)}"
                )
            current_cycles.update(rule_cycles)

    def run(self, payment_cycle, rule):
        """
        Trigger a job in AWS batch for a single payment cycle and single rule
        """
        logging.info(f"Running {payment_cycle} for {self.reward_program_id}")
        submission_data = deepcopy(rule)
        docker_image = submission_data["core"]["docker_image"]
        submission_data["run"] = {
            "reward_program_id": self.reward_program_id,
            "payment_cycle": payment_cycle,
        }
        payment_cycle_output = self.reward_program_output_location.joinpath(
            f"paymentCycle={payment_cycle}"
        )

        if payment_cycle_output.joinpath("results.parquet").exists():
            logging.info(
                f"Already processed: {payment_cycle} for {self.reward_program_id}"
            )
            return None

        # We can't pass much data directly in AWS batch, so write the
        # parameters we want to use to S3 and pass the location into the task
        parameters_location = payment_cycle_output.joinpath("parameters.json")
        with parameters_location.open("w") as params_out:
            json.dump(submission_data, params_out)

        job = run_job(
            docker_image,
            parameters_location.as_uri(),
            payment_cycle_output.as_uri(),
            tags={
                "reward_program_id": self.reward_program_id,
                "payment_cycle": payment_cycle,
            },
        )
        return job

    def run_all_payment_cycles(self) -> None:
        """
        Run all rules for all payment cycles that can be processed at this time
        """
        logging.info(f"Running all for {self.reward_program_id}")
        if self.is_locked():
            logging.info(f"Reward program {self.reward_program_id} is locked, skipping")
            return
        rules = self.get_rules()
        self.raise_on_payment_cycle_overlap(rules)
        for rule in rules:
            processable_payment_cycles = (
                self.get_all_payment_cycles(rule) - self.processed_cycles
            )
            for payment_cycle in sorted(processable_payment_cycles):
                self.run(payment_cycle, rule)
        logging.info(f"All new payment cycles triggered for {self.reward_program_id}")

def resolve_rule(did: str):
    url = Resolver(get_resolver()).resolve(did)["didDocument"]["alsoKnownAs"][0]
    return requests.get(url).json()
<|endoftext|>


<commit_before>import sublime, sublime_plugin
import subprocess

class PhpSimpleRefactorCommand(sublime_plugin.TextCommand):
	def run(self, edit):
		(rowBegin,col) = self.view.rowcol(self.view.sel()[0].begin())
		(rowEnd,col) = self.view.rowcol(self.view.sel()[0].end())

		self.rowBegin = rowBegin + 1
		self.rowEnd = rowEnd + 1
		view = sublime.Window.active_view(sublime.active_window())
		if(view.file_name() is None):
			sublime.status_message('File not saved yet, please save first')
		else:
			self.file_name = view.file_name();
			sublime.active_window().show_input_panel('Function name', '', self.on_filled_info, None, None)

	def on_filled_info(self, functionName):
		view = sublime.Window.active_view(sublime.active_window())
		# save if dirty
		if view.is_dirty():
			view.run_command('save')
			sublime.status_message('File saved')

		settings = sublime.load_settings('PHPSimpleRefactor.sublime-settings')
		self.php_path = settings.get('php_path')
		self.refactor_path = settings.get('refactor_path')
		rows = ''.join([str(self.rowBegin), "-", str(self.rowEnd)])
		cmd = ''.join([self.php_path, ' ', self.refactor_path,' ',  'extract-method', ' ', self.file_name, ' ', rows, ' ', functionName, ' | patch -p1'])
		p = subprocess.Popen(cmd, shell=True, bufsize=-1, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
		output, error = p.communicate()
		if error:
			sublime.error_message(error.decode('utf-8'))
		else:
			sublime.status_message('Refactor completed')
<commit_msg>move 'patch' to another subprocess<commit_after>import sublime, sublime_plugin
import subprocess
import tempfile
import os

class PhpSimpleRefactorCommand(sublime_plugin.TextCommand):
    def run(self, edit):
        (rowBegin,col) = self.view.rowcol(self.view.sel()[0].begin())
        (rowEnd,col) = self.view.rowcol(self.view.sel()[0].end())

        self.rowBegin = rowBegin + 1
        self.rowEnd = rowEnd + 1
        view = sublime.Window.active_view(sublime.active_window())
        if(view.file_name() is None):
            sublime.status_message('File not saved yet, please save first')
        else:
            self.file_name = view.file_name();
            sublime.active_window().show_input_panel('Function name', '', self.on_filled_info, None, None)

    def on_filled_info(self, functionName):
        view = sublime.Window.active_view(sublime.active_window())
        # save if dirty
        if view.is_dirty():
            view.run_command('save')
            sublime.status_message('File saved')

        settings = sublime.load_settings('PHPSimpleRefactor.sublime-settings')
        self.php_path = settings.get('php_path')
        self.refactor_path = settings.get('refactor_path')
        rows = ''.join([str(self.rowBegin), "-", str(self.rowEnd)])
        cmd = ''.join([self.php_path, ' "', self.refactor_path,'" ',  'extract-method', ' "', self.file_name, '" ', rows, ' ', functionName])
        p = subprocess.Popen(cmd, shell=True, bufsize=-1, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
        output, error = p.communicate()
        if error:
            sublime.error_message(error.decode('utf-8'))
        else:
            fp = tempfile.NamedTemporaryFile(delete=False)
            fp.write(output)
            fp.close()
            cmd = ''.join(['patch -p1 "',self.file_name,'" "',fp.name,'"'])
            pr = subprocess.Popen(cmd, shell=True, bufsize=-1, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
            output, error = pr.communicate()
            os.remove(fp.name)
            if error:
                sublime.error_message(error.decode('utf-8'))
            else:
                sublime.status_message('Refactor completed')
<|endoftext|>


<commit_before>"""
Module to read / write wav files using numpy arrays

Functions
---------
`read`: Return the sample rate (in samples/sec) and data from a WAV file.

`write`: Write a numpy array as a WAV file.

"""
from __future__ import division, print_function, absolute_import

import sys
import numpy
import struct
import warnings

__all__ = [
    'WavFileWarning',
    'read',
    'write'
]

class WavFileWarning(UserWarning):
    pass

WAVE_FORMAT_PCM = 0x0001
WAVE_FORMAT_IEEE_FLOAT = 0x0003
WAVE_FORMAT_EXTENSIBLE = 0xfffe
KNOWN_WAVE_FORMATS = (WAVE_FORMAT_PCM, WAVE_FORMAT_IEEE_FLOAT)

# assumes file pointer is immediately
#  after the 'fmt ' id

def _read_fmt_chunk(fid, is_big_endian):
    if is_big_endian:
        fmt = '>'
    else:
        fmt = '<'

    size = res = struct.unpack(fmt+'I', fid.read(4))[0]
    bytes_read = 0

    if size < 16:
        raise ValueError("Binary structure of wave file is not compliant")

    res = struct.unpack(fmt+'HHIIHH', fid.read(16))
    bytes_read += 16
    comp, noc, rate, sbytes, ba, bits = res

    if comp == WAVE_FORMAT_EXTENSIBLE and size >= (16+2):
        ext_chunk_size = struct.unpack(fmt+'H', fid.read(2))[0]
        bytes_read += 2
        if ext_chunk_size >= 22:
            extensible_chunk_data = fid.read(22)
            bytes_read += 22
            raw_guid = extensible_chunk_data[2+4:2+4+16]
            # GUID template {XXXXXXXX-0000-0010-8000-00AA00389B71} (RFC-2361)
            # MS GUID byte order: first three groups are native byte order,
            # rest is Big Endian
            if is_big_endian:
                tail = b'\x00\x00\x00\x10\x80\x00\x00\xAA\x00\x38\x9B\x71'
            else:
                tail = b'\x00\x00\x10\x00\x80\x00\x00\xAA\x00\x38\x9B\x71'
            if raw_guid.endswith(tail):
                comp = struct.unpack(fmt+'I', raw_guid[:4])[0]
        else:
            raise ValueError("Binary structure of wave file is not compliant")

    if comp not in KNOWN_WAVE_FORMATS:
        raise ValueError("Unknown wave file format")

    # move file pointer to next chunk
    if size > (bytes_read):
        fid.read(size - bytes_read)

    return size, comp, noc, rate, sbytes, ba, bits

# assumes file pointer is immediately
#   after the 'data' id
def _read_data_chunk(fid, comp, noc, bits, is_big_endian, mmap=False):
    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'
    size = struct.unpack(fmt, fid.read(4))[0]

    bytes = bits//8
    if bits == 8:
        dtype = 'u1'
    else:
        if is_big_endian:
            dtype = '>'
        else:
            dtype = '<'
        if comp == 1:
            dtype += 'i%d' % bytes
        else:
            dtype += 'f%d' % bytes
    if not mmap:
        data = numpy.fromstring(fid.read(size), dtype=dtype)
    else:
        start = fid.tell()
        data = numpy.memmap(fid, dtype=dtype, mode='c', offset=start,
                            shape=(size//bytes,))
        fid.seek(start + size)

    if noc > 1:
        data = data.reshape(-1, noc)
    return data

def _skip_unknown_chunk(fid, is_big_endian):
    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'

    data = fid.read(4)
    # call unpack() and seek() only if we have really read data from file
    # otherwise empty read at the end of the file would trigger
    # unnecessary exception at unpack() call
    # in case data equals somehow to 0, there is no need for seek() anyway
    if data:
        size = struct.unpack(fmt, data)[0]
        fid.seek(size, 1)

def _read_riff_chunk(fid):
    str1 = fid.read(4)
    if str1 == b'RIFF':
        is_big_endian = False
    elif str1 == b'RIFX':
        is_big_endian = True
    else:
        raise ValueError("Not a WAV file.")

    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'
    fsize = struct.unpack(fmt, fid.read(4))[0] + 8
    str2 = fid.read(4)
    if (str2 != b'WAVE'):
        raise ValueError("Not a WAV file.")

    return fsize, is_big_endian

# open a wave-file

def read(filename, mmap=False):
    """
    Return the sample rate (in samples/sec) and data from a WAV file.

    Parameters
    ----------
    filename : string or open file handle
        Input wav file.
    mmap : bool, optional
        Whether to read data as memory mapped.
        Only to be used on real files (Default: False).

        .. versionadded:: 0.12.0

    Returns
    -------
    rate : int
        Sample rate of wav file.
    data : numpy array
        Data read from wav file.

    Notes
    -----
    * The file can be an open file or a filename.
    * The returned sample rate is a Python integer.
    * The data is returned as a numpy array with a data-type determined
      from the file.
    * This function cannot read wav files with 24 bit data.

    """
    if hasattr(filename, 'read'):
        fid = filename
        mmap = False
    else:
        fid = open(filename, 'rb')

    try:
        fsize, is_big_endian = _read_riff_chunk(fid)
        fmt_chunk_received = False
        noc = 1
        bits = 8
        comp = WAVE_FORMAT_PCM
        while (fid.tell() < fsize):
            # read the next chunk
            chunk_id = fid.read(4)
            if chunk_id == b'fmt ':
                fmt_chunk_received = True
                size, comp, noc, rate, sbytes, ba, bits = _read_fmt_chunk(
                                        fid, is_big_endian=is_big_endian)
                if bits not in (8, 16, 32, 64, 128):
                    raise ValueError("Unsupported bit depth: the wav file "
                                     "has {}-bit data.".format(bits))
            elif chunk_id == b'fact':
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            elif chunk_id == b'data':
                if not fmt_chunk_received:
                    raise ValueError("No fmt chunk before data")
                data = _read_data_chunk(fid, comp, noc, bits,
                                        is_big_endian=is_big_endian, mmap=mmap)
            elif chunk_id == b'LIST':
                # Someday this could be handled properly but for now skip it
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            elif chunk_id in (b'JUNK', b'Fake'):
                # Skip alignment chunks without warning
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            else:
                warnings.warn("Chunk (non-data) not understood, skipping it.",
                              WavFileWarning)
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
    finally:
        if not hasattr(filename, 'read'):
            fid.close()
        else:
            fid.seek(0)

    return rate, data

# Write a wave-file
# sample rate, data

def write(filename, rate, data):
    """
    Write a numpy array as a WAV file.

    Parameters
    ----------
    filename : string or open file handle
        Output wav file.
    rate : int
        The sample rate (in samples/sec).
    data : ndarray
        A 1-D or 2-D numpy array of either integer or float data-type.

    Notes
    -----
    * The file can be an open file or a filename.

    * Writes a simple uncompressed WAV file.
    * The bits-per-sample will be determined by the data-type.
    * To write multiple-channels, use a 2-D array of shape
      (Nsamples, Nchannels).

    """
    if hasattr(filename, 'write'):
        fid = filename
    else:
        fid = open(filename, 'wb')

    try:
        dkind = data.dtype.kind
        if not (dkind == 'i' or dkind == 'f' or (dkind == 'u' and
                                                 data.dtype.itemsize == 1)):
            raise ValueError("Unsupported data type '%s'" % data.dtype)

        header_data = b''

        header_data += b'RIFF'
        header_data += b'\x00\x00\x00\x00'
        header_data += b'WAVE'

        # fmt chunk
        header_data += b'fmt '
        if dkind == 'f':
            comp = 3
        else:
            comp = 1
        if data.ndim == 1:
            noc = 1
        else:
            noc = data.shape[1]
        bits = data.dtype.itemsize * 8
        sbytes = rate*(bits // 8)*noc
        ba = noc * (bits // 8)

        fmt_chunk_data = struct.pack('<HHIIHH', comp, noc, rate, sbytes,
                                     ba, bits)
        if not (dkind == 'i' or dkind == 'u'):
            # add cbSize field for non-PCM files
            fmt_chunk_data += b'\x00\x00'

        header_data += struct.pack('<I', len(fmt_chunk_data))
        header_data += fmt_chunk_data

        # fact chunk (non-PCM files)
        if not (dkind == 'i' or dkind == 'u'):
            header_data += b'fact'
            header_data += struct.pack('<II', 4, data.shape[0])

        # check data size (needs to be immediately before the data chunk)
        if ((len(header_data)-4-4) + (4+4+data.nbytes)) > 0xFFFFFFFF:
            raise ValueError("Data exceeds wave file size limit")

        fid.write(header_data)

        # data chunk
        fid.write(b'data')
        fid.write(struct.pack('<I', data.nbytes))
        if data.dtype.byteorder == '>' or (data.dtype.byteorder == '=' and
                                           sys.byteorder == 'big'):
            data = data.byteswap()
        _array_tofile(fid, data)

        # Determine file size and place it in correct
        #  position at start of the file.
        size = fid.tell()
        fid.seek(4)
        fid.write(struct.pack('<I', size-8))

    finally:
        if not hasattr(filename, 'write'):
            fid.close()
        else:
            fid.seek(0)

if sys.version_info[0] >= 3:
    def _array_tofile(fid, data):
        # ravel gives a c-contiguous buffer
        fid.write(data.ravel().view('b').data)
else:
    def _array_tofile(fid, data):
        fid.write(data.tostring())
<commit_msg>DOC: Cleanup read/write docstrings<commit_after>"""
Module to read / write wav files using numpy arrays

Functions
---------
`read`: Return the sample rate (in samples/sec) and data from a WAV file.

`write`: Write a numpy array as a WAV file.

"""
from __future__ import division, print_function, absolute_import

import sys
import numpy
import struct
import warnings

__all__ = [
    'WavFileWarning',
    'read',
    'write'
]

class WavFileWarning(UserWarning):
    pass

WAVE_FORMAT_PCM = 0x0001
WAVE_FORMAT_IEEE_FLOAT = 0x0003
WAVE_FORMAT_EXTENSIBLE = 0xfffe
KNOWN_WAVE_FORMATS = (WAVE_FORMAT_PCM, WAVE_FORMAT_IEEE_FLOAT)

# assumes file pointer is immediately
#  after the 'fmt ' id

def _read_fmt_chunk(fid, is_big_endian):
    if is_big_endian:
        fmt = '>'
    else:
        fmt = '<'

    size = res = struct.unpack(fmt+'I', fid.read(4))[0]
    bytes_read = 0

    if size < 16:
        raise ValueError("Binary structure of wave file is not compliant")

    res = struct.unpack(fmt+'HHIIHH', fid.read(16))
    bytes_read += 16
    comp, noc, rate, sbytes, ba, bits = res

    if comp == WAVE_FORMAT_EXTENSIBLE and size >= (16+2):
        ext_chunk_size = struct.unpack(fmt+'H', fid.read(2))[0]
        bytes_read += 2
        if ext_chunk_size >= 22:
            extensible_chunk_data = fid.read(22)
            bytes_read += 22
            raw_guid = extensible_chunk_data[2+4:2+4+16]
            # GUID template {XXXXXXXX-0000-0010-8000-00AA00389B71} (RFC-2361)
            # MS GUID byte order: first three groups are native byte order,
            # rest is Big Endian
            if is_big_endian:
                tail = b'\x00\x00\x00\x10\x80\x00\x00\xAA\x00\x38\x9B\x71'
            else:
                tail = b'\x00\x00\x10\x00\x80\x00\x00\xAA\x00\x38\x9B\x71'
            if raw_guid.endswith(tail):
                comp = struct.unpack(fmt+'I', raw_guid[:4])[0]
        else:
            raise ValueError("Binary structure of wave file is not compliant")

    if comp not in KNOWN_WAVE_FORMATS:
        raise ValueError("Unknown wave file format")

    # move file pointer to next chunk
    if size > (bytes_read):
        fid.read(size - bytes_read)

    return size, comp, noc, rate, sbytes, ba, bits

# assumes file pointer is immediately
#   after the 'data' id
def _read_data_chunk(fid, comp, noc, bits, is_big_endian, mmap=False):
    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'
    size = struct.unpack(fmt, fid.read(4))[0]

    bytes = bits//8
    if bits == 8:
        dtype = 'u1'
    else:
        if is_big_endian:
            dtype = '>'
        else:
            dtype = '<'
        if comp == 1:
            dtype += 'i%d' % bytes
        else:
            dtype += 'f%d' % bytes
    if not mmap:
        data = numpy.fromstring(fid.read(size), dtype=dtype)
    else:
        start = fid.tell()
        data = numpy.memmap(fid, dtype=dtype, mode='c', offset=start,
                            shape=(size//bytes,))
        fid.seek(start + size)

    if noc > 1:
        data = data.reshape(-1, noc)
    return data

def _skip_unknown_chunk(fid, is_big_endian):
    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'

    data = fid.read(4)
    # call unpack() and seek() only if we have really read data from file
    # otherwise empty read at the end of the file would trigger
    # unnecessary exception at unpack() call
    # in case data equals somehow to 0, there is no need for seek() anyway
    if data:
        size = struct.unpack(fmt, data)[0]
        fid.seek(size, 1)

def _read_riff_chunk(fid):
    str1 = fid.read(4)
    if str1 == b'RIFF':
        is_big_endian = False
    elif str1 == b'RIFX':
        is_big_endian = True
    else:
        raise ValueError("Not a WAV file.")

    if is_big_endian:
        fmt = '>I'
    else:
        fmt = '<I'
    fsize = struct.unpack(fmt, fid.read(4))[0] + 8
    str2 = fid.read(4)
    if (str2 != b'WAVE'):
        raise ValueError("Not a WAV file.")

    return fsize, is_big_endian

def read(filename, mmap=False):
    """
    Open a WAV file

    Return the sample rate (in samples/sec) and data from a WAV file.

    Parameters
    ----------
    filename : string or open file handle
        Input wav file.
    mmap : bool, optional
        Whether to read data as memory-mapped.
        Only to be used on real files (Default: False).

        .. versionadded:: 0.12.0

    Returns
    -------
    rate : int
        Sample rate of wav file as a Python integer.
    data : numpy array
        Data read from wav file.  Data-type is determined from the file;
        see Notes.

    Notes
    -----
    This function cannot read wav files with 24-bit data.

    Common data types: [1]_

    =====================  ===========  ===========  =============
         WAV format            Min          Max       NumPy dtype
    =====================  ===========  ===========  =============
    32-bit floating-point  -1.0         +1.0         float32
    32-bit PCM             -2147483648  +2147483647  int32
    16-bit PCM             -32768       +32767       int16
    8-bit PCM              0            255          uint8
    =====================  ===========  ===========  =============

    Note that 8-bit PCM is unsigned.

    References
    ----------
    .. [1] IBM Corporation and Microsoft Corporation, "Multimedia Programming
       Interface and Data Specifications 1.0", section "Data Format of the
       Samples", August 1991
       http://www-mmsp.ece.mcgill.ca/documents/audioformats/wave/Docs/riffmci.pdf

    """
    if hasattr(filename, 'read'):
        fid = filename
        mmap = False
    else:
        fid = open(filename, 'rb')

    try:
        fsize, is_big_endian = _read_riff_chunk(fid)
        fmt_chunk_received = False
        noc = 1
        bits = 8
        comp = WAVE_FORMAT_PCM
        while (fid.tell() < fsize):
            # read the next chunk
            chunk_id = fid.read(4)
            if chunk_id == b'fmt ':
                fmt_chunk_received = True
                size, comp, noc, rate, sbytes, ba, bits = _read_fmt_chunk(
                                        fid, is_big_endian=is_big_endian)
                if bits not in (8, 16, 32, 64, 128):
                    raise ValueError("Unsupported bit depth: the wav file "
                                     "has {}-bit data.".format(bits))
            elif chunk_id == b'fact':
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            elif chunk_id == b'data':
                if not fmt_chunk_received:
                    raise ValueError("No fmt chunk before data")
                data = _read_data_chunk(fid, comp, noc, bits,
                                        is_big_endian=is_big_endian, mmap=mmap)
            elif chunk_id == b'LIST':
                # Someday this could be handled properly but for now skip it
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            elif chunk_id in (b'JUNK', b'Fake'):
                # Skip alignment chunks without warning
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
            else:
                warnings.warn("Chunk (non-data) not understood, skipping it.",
                              WavFileWarning)
                _skip_unknown_chunk(fid, is_big_endian=is_big_endian)
    finally:
        if not hasattr(filename, 'read'):
            fid.close()
        else:
            fid.seek(0)

    return rate, data

def write(filename, rate, data):
    """
    Write a numpy array as a WAV file.

    Parameters
    ----------
    filename : string or open file handle
        Output wav file.
    rate : int
        The sample rate (in samples/sec).
    data : ndarray
        A 1-D or 2-D numpy array of either integer or float data-type.

    Notes
    -----
    * Writes a simple uncompressed WAV file.
    * To write multiple-channels, use a 2-D array of shape
      (Nsamples, Nchannels).
    * The bits-per-sample and PCM/float will be determined by the data-type.

    Common data types: [1]_

    =====================  ===========  ===========  =============
         WAV format            Min          Max       NumPy dtype
    =====================  ===========  ===========  =============
    32-bit floating-point  -1.0         +1.0         float32
    32-bit PCM             -2147483648  +2147483647  int32
    16-bit PCM             -32768       +32767       int16
    8-bit PCM              0            255          uint8
    =====================  ===========  ===========  =============

    Note that 8-bit PCM is unsigned.

    References
    ----------
    .. [1] IBM Corporation and Microsoft Corporation, "Multimedia Programming
       Interface and Data Specifications 1.0", section "Data Format of the
       Samples", August 1991
       http://www-mmsp.ece.mcgill.ca/documents/audioformats/wave/Docs/riffmci.pdf

    """
    if hasattr(filename, 'write'):
        fid = filename
    else:
        fid = open(filename, 'wb')

    try:
        dkind = data.dtype.kind
        if not (dkind == 'i' or dkind == 'f' or (dkind == 'u' and
                                                 data.dtype.itemsize == 1)):
            raise ValueError("Unsupported data type '%s'" % data.dtype)

        header_data = b''

        header_data += b'RIFF'
        header_data += b'\x00\x00\x00\x00'
        header_data += b'WAVE'

        # fmt chunk
        header_data += b'fmt '
        if dkind == 'f':
            comp = 3
        else:
            comp = 1
        if data.ndim == 1:
            noc = 1
        else:
            noc = data.shape[1]
        bits = data.dtype.itemsize * 8
        sbytes = rate*(bits // 8)*noc
        ba = noc * (bits // 8)

        fmt_chunk_data = struct.pack('<HHIIHH', comp, noc, rate, sbytes,
                                     ba, bits)
        if not (dkind == 'i' or dkind == 'u'):
            # add cbSize field for non-PCM files
            fmt_chunk_data += b'\x00\x00'

        header_data += struct.pack('<I', len(fmt_chunk_data))
        header_data += fmt_chunk_data

        # fact chunk (non-PCM files)
        if not (dkind == 'i' or dkind == 'u'):
            header_data += b'fact'
            header_data += struct.pack('<II', 4, data.shape[0])

        # check data size (needs to be immediately before the data chunk)
        if ((len(header_data)-4-4) + (4+4+data.nbytes)) > 0xFFFFFFFF:
            raise ValueError("Data exceeds wave file size limit")

        fid.write(header_data)

        # data chunk
        fid.write(b'data')
        fid.write(struct.pack('<I', data.nbytes))
        if data.dtype.byteorder == '>' or (data.dtype.byteorder == '=' and
                                           sys.byteorder == 'big'):
            data = data.byteswap()
        _array_tofile(fid, data)

        # Determine file size and place it in correct
        #  position at start of the file.
        size = fid.tell()
        fid.seek(4)
        fid.write(struct.pack('<I', size-8))

    finally:
        if not hasattr(filename, 'write'):
            fid.close()
        else:
            fid.seek(0)

if sys.version_info[0] >= 3:
    def _array_tofile(fid, data):
        # ravel gives a c-contiguous buffer
        fid.write(data.ravel().view('b').data)
else:
    def _array_tofile(fid, data):
        fid.write(data.tostring())
<|endoftext|>


<commit_before>"""Compressed Sparse Column matrix format"""
from __future__ import division, print_function, absolute_import

__docformat__ = "restructuredtext en"

__all__ = ['csc_matrix', 'isspmatrix_csc']

from warnings import warn

import numpy as np
from scipy.lib.six import xrange

from .base import isspmatrix
from ._sparsetools import csc_tocsr
from . import _sparsetools
from .sputils import upcast, isintlike, IndexMixin, get_index_dtype

from .compressed import _cs_matrix

class csc_matrix(_cs_matrix, IndexMixin):
    """
    Compressed Sparse Column matrix

    This can be instantiated in several ways:

        csc_matrix(D)
            with a dense matrix or rank-2 ndarray D

        csc_matrix(S)
            with another sparse matrix S (equivalent to S.tocsc())

        csc_matrix((M, N), [dtype])
            to construct an empty matrix with shape (M, N)
            dtype is optional, defaulting to dtype='d'.

        csc_matrix((data, ij), [shape=(M, N)])
            where ``data`` and ``ij`` satisfy the relationship
            ``a[ij[0, k], ij[1, k]] = data[k]``

        csc_matrix((data, indices, indptr), [shape=(M, N)])
            is the standard CSC representation where the row indices for
            column i are stored in ``indices[indptr[i]:indptr[i+1]]``
            and their corresponding values are stored in
            ``data[indptr[i]:indptr[i+1]]``.  If the shape parameter is
            not supplied, the matrix dimensions are inferred from
            the index arrays.

    Attributes
    ----------
    dtype : dtype
        Data type of the matrix
    shape : 2-tuple
        Shape of the matrix
    ndim : int
        Number of dimensions (this is always 2)
    nnz
        Number of nonzero elements
    data
        Data array of the matrix
    indices
        CSC format index array
    indptr
        CSC format index pointer array
    has_sorted_indices
        Whether indices are sorted

    Notes
    -----

    Sparse matrices can be used in arithmetic operations: they support
    addition, subtraction, multiplication, division, and matrix power.

    Advantages of the CSC format
        - efficient arithmetic operations CSC + CSC, CSC * CSC, etc.
        - efficient column slicing
        - fast matrix vector products (CSR, BSR may be faster)

    Disadvantages of the CSC format
      - slow row slicing operations (consider CSR)
      - changes to the sparsity structure are expensive (consider LIL or DOK)

    Examples
    --------

    >>> from scipy.sparse import *
    >>> from scipy import *
    >>> csc_matrix((3, 4), dtype=int8).toarray()
    array([[0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 0, 0]], dtype=int8)

    >>> row = array([0, 2, 2, 0, 1, 2])
    >>> col = array([0, 0, 1, 2, 2, 2])
    >>> data = array([1, 2, 3, 4, 5, 6])
    >>> csc_matrix((data, (row, col)), shape=(3, 3)).toarray()
    array([[1, 0, 4],
           [0, 0, 5],
           [2, 3, 6]])

    >>> indptr = array([0, 2, 3, 6])
    >>> indices = array([0, 2, 2, 0, 1, 2])
    >>> data = array([1, 2, 3, 4, 5, 6])
    >>> csc_matrix((data, indices, indptr), shape=(3, 3)).toarray()
    array([[1, 0, 4],
           [0, 0, 5],
           [2, 3, 6]])

    """

    def transpose(self, copy=False):
        from .csr import csr_matrix
        M,N = self.shape
        return csr_matrix((self.data,self.indices,self.indptr),(N,M),copy=copy)

    def __iter__(self):
        csr = self.tocsr()
        for r in xrange(self.shape[0]):
            yield csr[r,:]

    def tocsc(self, copy=False):
        if copy:
            return self.copy()
        else:
            return self

    def tocsr(self):
        M,N = self.shape
        idx_dtype = get_index_dtype((self.indptr, self.indices),
                                    maxval=max(self.nnz, N))
        indptr = np.empty(M + 1, dtype=idx_dtype)
        indices = np.empty(self.nnz, dtype=idx_dtype)
        data = np.empty(self.nnz, dtype=upcast(self.dtype))

        csc_tocsr(M, N,
                  self.indptr.astype(idx_dtype),
                  self.indices.astype(idx_dtype),
                  self.data,
                  indptr,
                  indices,
                  data)

        from .csr import csr_matrix
        A = csr_matrix((data, indices, indptr), shape=self.shape)
        A.has_sorted_indices = True
        return A

    def __getitem__(self, key):
        # Use CSR to implement fancy indexing.

        row, col = self._unpack_index(key)
        # Things that return submatrices. row or col is a int or slice.
        if (isinstance(row, slice) or isinstance(col, slice) or
                isintlike(row) or isintlike(col)):
            return self.T[col, row].T
        # Things that return a sequence of values.
        else:
            return self.T[col, row]

    def nonzero(self):
        # CSC can't use _cs_matrix's .nonzero method because it
        # returns the indices sorted for self transposed.

        # Get row and col indices, from _cs_matrix.tocoo
        major_dim, minor_dim = self._swap(self.shape)
        minor_indices = self.indices
        major_indices = np.empty(len(minor_indices), dtype=self.indptr.dtype)
        _sparsetools.expandptr(major_dim, self.indptr, major_indices)
        row, col = self._swap((major_indices, minor_indices))

        # Sort them to be in C-style order
        ind = np.lexsort((col, row))
        row = row[ind]
        col = col[ind]

        return row, col

    nonzero.__doc__ = _cs_matrix.nonzero.__doc__

    def getrow(self, i):
        """Returns a copy of row i of the matrix, as a (1 x n)
        CSR matrix (row vector).
        """
        # transpose to use CSR code
        # we convert to CSR to maintain compatibility with old impl.
        # in spmatrix.getrow()
        return self.T.getcol(i).T.tocsr()

    def getcol(self, i):
        """Returns a copy of column i of the matrix, as a (m x 1)
        CSC matrix (column vector).
        """
        # transpose to use CSR code
        return self.T.getrow(i).T

    # these functions are used by the parent class (_cs_matrix)
    # to remove redudancy between csc_matrix and csr_matrix
    def _swap(self,x):
        """swap the members of x if this is a column-oriented matrix
        """
        return (x[1],x[0])

def isspmatrix_csc(x):
    return isinstance(x, csc_matrix)
<commit_msg>STY avoid unnecessary transposes in csr_matrix.getcol/row<commit_after>"""Compressed Sparse Column matrix format"""
from __future__ import division, print_function, absolute_import

__docformat__ = "restructuredtext en"

__all__ = ['csc_matrix', 'isspmatrix_csc']

from warnings import warn

import numpy as np
from scipy.lib.six import xrange

from .base import isspmatrix
from ._sparsetools import csc_tocsr
from . import _sparsetools
from .sputils import upcast, isintlike, IndexMixin, get_index_dtype

from .compressed import _cs_matrix

class csc_matrix(_cs_matrix, IndexMixin):
    """
    Compressed Sparse Column matrix

    This can be instantiated in several ways:

        csc_matrix(D)
            with a dense matrix or rank-2 ndarray D

        csc_matrix(S)
            with another sparse matrix S (equivalent to S.tocsc())

        csc_matrix((M, N), [dtype])
            to construct an empty matrix with shape (M, N)
            dtype is optional, defaulting to dtype='d'.

        csc_matrix((data, ij), [shape=(M, N)])
            where ``data`` and ``ij`` satisfy the relationship
            ``a[ij[0, k], ij[1, k]] = data[k]``

        csc_matrix((data, indices, indptr), [shape=(M, N)])
            is the standard CSC representation where the row indices for
            column i are stored in ``indices[indptr[i]:indptr[i+1]]``
            and their corresponding values are stored in
            ``data[indptr[i]:indptr[i+1]]``.  If the shape parameter is
            not supplied, the matrix dimensions are inferred from
            the index arrays.

    Attributes
    ----------
    dtype : dtype
        Data type of the matrix
    shape : 2-tuple
        Shape of the matrix
    ndim : int
        Number of dimensions (this is always 2)
    nnz
        Number of nonzero elements
    data
        Data array of the matrix
    indices
        CSC format index array
    indptr
        CSC format index pointer array
    has_sorted_indices
        Whether indices are sorted

    Notes
    -----

    Sparse matrices can be used in arithmetic operations: they support
    addition, subtraction, multiplication, division, and matrix power.

    Advantages of the CSC format
        - efficient arithmetic operations CSC + CSC, CSC * CSC, etc.
        - efficient column slicing
        - fast matrix vector products (CSR, BSR may be faster)

    Disadvantages of the CSC format
      - slow row slicing operations (consider CSR)
      - changes to the sparsity structure are expensive (consider LIL or DOK)

    Examples
    --------

    >>> from scipy.sparse import *
    >>> from scipy import *
    >>> csc_matrix((3, 4), dtype=int8).toarray()
    array([[0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 0, 0]], dtype=int8)

    >>> row = array([0, 2, 2, 0, 1, 2])
    >>> col = array([0, 0, 1, 2, 2, 2])
    >>> data = array([1, 2, 3, 4, 5, 6])
    >>> csc_matrix((data, (row, col)), shape=(3, 3)).toarray()
    array([[1, 0, 4],
           [0, 0, 5],
           [2, 3, 6]])

    >>> indptr = array([0, 2, 3, 6])
    >>> indices = array([0, 2, 2, 0, 1, 2])
    >>> data = array([1, 2, 3, 4, 5, 6])
    >>> csc_matrix((data, indices, indptr), shape=(3, 3)).toarray()
    array([[1, 0, 4],
           [0, 0, 5],
           [2, 3, 6]])

    """

    def transpose(self, copy=False):
        from .csr import csr_matrix
        M,N = self.shape
        return csr_matrix((self.data,self.indices,self.indptr),(N,M),copy=copy)

    def __iter__(self):
        csr = self.tocsr()
        for r in xrange(self.shape[0]):
            yield csr[r,:]

    def tocsc(self, copy=False):
        if copy:
            return self.copy()
        else:
            return self

    def tocsr(self):
        M,N = self.shape
        idx_dtype = get_index_dtype((self.indptr, self.indices),
                                    maxval=max(self.nnz, N))
        indptr = np.empty(M + 1, dtype=idx_dtype)
        indices = np.empty(self.nnz, dtype=idx_dtype)
        data = np.empty(self.nnz, dtype=upcast(self.dtype))

        csc_tocsr(M, N,
                  self.indptr.astype(idx_dtype),
                  self.indices.astype(idx_dtype),
                  self.data,
                  indptr,
                  indices,
                  data)

        from .csr import csr_matrix
        A = csr_matrix((data, indices, indptr), shape=self.shape)
        A.has_sorted_indices = True
        return A

    def __getitem__(self, key):
        # Use CSR to implement fancy indexing.

        row, col = self._unpack_index(key)
        # Things that return submatrices. row or col is a int or slice.
        if (isinstance(row, slice) or isinstance(col, slice) or
                isintlike(row) or isintlike(col)):
            return self.T[col, row].T
        # Things that return a sequence of values.
        else:
            return self.T[col, row]

    def nonzero(self):
        # CSC can't use _cs_matrix's .nonzero method because it
        # returns the indices sorted for self transposed.

        # Get row and col indices, from _cs_matrix.tocoo
        major_dim, minor_dim = self._swap(self.shape)
        minor_indices = self.indices
        major_indices = np.empty(len(minor_indices), dtype=self.indptr.dtype)
        _sparsetools.expandptr(major_dim, self.indptr, major_indices)
        row, col = self._swap((major_indices, minor_indices))

        # Sort them to be in C-style order
        ind = np.lexsort((col, row))
        row = row[ind]
        col = col[ind]

        return row, col

    nonzero.__doc__ = _cs_matrix.nonzero.__doc__

    def getrow(self, i):
        """Returns a copy of row i of the matrix, as a (1 x n)
        CSR matrix (row vector).
        """
        # we convert to CSR to maintain compatibility with old impl.
        # in spmatrix.getrow()
        return self._get_submatrix(i, slice(None)).tocsr()

    def getcol(self, i):
        """Returns a copy of column i of the matrix, as a (m x 1)
        CSC matrix (column vector).
        """
        return self._get_submatrix(slice(None), i)

    # these functions are used by the parent class (_cs_matrix)
    # to remove redudancy between csc_matrix and csr_matrix
    def _swap(self,x):
        """swap the members of x if this is a column-oriented matrix
        """
        return (x[1],x[0])

def isspmatrix_csc(x):
    return isinstance(x, csc_matrix)
<|endoftext|>


<commit_before>from __future__ import with_statement

import os, sys, shutil, pkg_resources
from subprocess import Popen, PIPE
from tempfile import NamedTemporaryFile, mkdtemp

def get_spider_list_from_eggfile(eggfile, project):
    # FIXME: we use a temporary directory here to avoid permissions problems
    # when running as system service, as "scrapy list" command tries to write
    # the scrapy.db sqlite database in current directory
    tmpdir = mkdtemp()
    try:
        with NamedTemporaryFile(suffix='.egg', dir=tmpdir) as f:
            shutil.copyfileobj(eggfile, f)
            f.flush()
            eggfile.seek(0)
            pargs = [sys.executable, '-m', 'scrapyd.eggrunner', 'list']
            env = os.environ.copy()
            env['SCRAPY_PROJECT'] = project
            env['SCRAPY_EGGFILE'] = f.name
            proc = Popen(pargs, stdout=PIPE, cwd=tmpdir, env=env)
            out = proc.communicate()[0]
            return out.splitlines()
    finally:
        shutil.rmtree(tmpdir)

def activate_egg(eggpath):
    """Activate a Scrapy egg file. This is meant to be used from egg runners
    to activate a Scrapy egg file. Don't use it from other code as it may
    leave unwanted side effects.
    """
    d = pkg_resources.find_distributions(eggpath).next()
    d.activate()
    settings_module = d.get_entry_info('scrapy', 'settings').module_name
    os.environ['SCRAPY_SETTINGS_MODULE'] = settings_module
<commit_msg>get_spider_list_from_eggfile(): fixed bug when SCRAPY_SETTINGS_DISABLED is set<commit_after>from __future__ import with_statement

import os, sys, shutil, pkg_resources
from subprocess import Popen, PIPE
from tempfile import NamedTemporaryFile, mkdtemp

def get_spider_list_from_eggfile(eggfile, project):
    # FIXME: we use a temporary directory here to avoid permissions problems
    # when running as system service, as "scrapy list" command tries to write
    # the scrapy.db sqlite database in current directory
    tmpdir = mkdtemp()
    try:
        with NamedTemporaryFile(suffix='.egg', dir=tmpdir) as f:
            shutil.copyfileobj(eggfile, f)
            f.flush()
            eggfile.seek(0)
            pargs = [sys.executable, '-m', 'scrapyd.eggrunner', 'list']
            env = os.environ.copy()
            env['SCRAPY_PROJECT'] = project
            env['SCRAPY_EGGFILE'] = f.name
            env.pop('SCRAPY_SETTINGS_DISABLED', None)
            proc = Popen(pargs, stdout=PIPE, cwd=tmpdir, env=env)
            out = proc.communicate()[0]
            return out.splitlines()
    finally:
        shutil.rmtree(tmpdir)

def activate_egg(eggpath):
    """Activate a Scrapy egg file. This is meant to be used from egg runners
    to activate a Scrapy egg file. Don't use it from other code as it may
    leave unwanted side effects.
    """
    d = pkg_resources.find_distributions(eggpath).next()
    d.activate()
    settings_module = d.get_entry_info('scrapy', 'settings').module_name
    os.environ['SCRAPY_SETTINGS_MODULE'] = settings_module
<|endoftext|>


<commit_before>#!/usr/bin/env python
import os
import json
import time
import argparse
import subprocess
import terminal

snapshot_id = '7067f369f7b76f0a3276beb561820a21c9b5204ab60fbd90524560db96d7cb38'
key_name = 'tempkey'

def get_credentials(utoken, atoken, credsfile='creds.json'):
    if utoken is None and atoken is None:
        try:
            creds = json.load(open(credsfile, 'r'))
            utoken = creds['user_token']
            atoken = creds['access_token']
        except:
            print "Can't open credentials file. \n ", \
                "You must provide a user token and a access token at least one time, or a valid credentials file"
            exit(127)
    elif (utoken is not None and atoken is None) or (utoken is None and atoken is not None):
        print "--utoken AND --atoken parameters must be passed together"
        exit(1)
    else:
        with open(credsfile, 'w') as cfile:
            json.dump({'user_token': utoken, 'access_token': atoken}, cfile)
    return utoken, atoken

def generate_ssh_key(key_file):
    subprocess.call(['ssh-keygen','-f', key_file,'-P',''])

def get_public_key(key_file):
    with open(key_file) as f:
        content = f.readlines()[0].rstrip('\n')
    return str(content)

def start_snap(name, snapshot_id, size=None, script=None):
    output = terminal.start_snapshot(snapshot_id, size, None, name, None, script)
    request_id = output['request_id']
    output = terminal.request_progress(request_id)
    state = 'None'
    time.sleep(1)
    while output['status'] != 'success':
        output = terminal.request_progress(request_id)
        if output['state'] != state:
            state = output['state']
            print('(%s)' % state)
        time.sleep(1)
    container_key = output['result']['container_key']
    subdomain = output['result']['subdomain']
    container_ip = output['result']['container_ip']
    return container_key, container_ip, subdomain

def run_on_terminal(cip, user, script):
    try:
        p = subprocess.Popen(
        ['ssh', '-q' ,'-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-i', 'pemfile',
         user + '@' + cip, script], 0, None,  None, None, None)
        p.wait()
        return p.returncode
    except Exception, e:
        return 'Error: [%s]' % e

def get_script(filename):
    try:
        with open(filename) as f:
            data = f.read()
            data.replace('\n',';')
        return data
    except Exception, e:
        print '(%s)' % e
        return None

def send_script(cip, user, script, path=''):
    try:
        p = subprocess.Popen(
        ['scp', script ,'-q' ,'-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-i', 'pemfile',
         '%s@%s:%s' % (user,cip,path)], 0, None,  None, None, None)
        p.wait()
        return p.returncode
    except Exception, e:
        return 'Error: [%s]' % e

def args_sanitizer(args):
    if args.size not in terminal.instance_types()['instance_types']:
        print "Error. Wrong instance type"
        exit(1)
    if (int(args.quantity) < 1 or int(args.quantity) > 100):
        print "Error. Terminal amount out of bounds ( should be in between 1 and 100 )"
        exit(1)
    if args.method not in ('ssh, startup'):
        print "Error. Not a valid method. (use ssh or startup)"
        exit(1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--size", type=str, default='medium', help="micro, mini, small, medium, xlarge.. etc")
    parser.add_argument("-q", "--quantity", type=int, default=3, help="How many nodes will have your deploy")
    parser.add_argument("-u", "--utoken", type=str, default=None ,help="Your user token")
    parser.add_argument("-a", "--atoken", type=str, default=None, help="Your access token")
    parser.add_argument("-c", "--creds", type=str, default='/etc/creds.json', help="A credentials json file")
    parser.add_argument("-x", "--script", type=str, default=None, help="A script file to be executed in the new Terminals - \
    With ssh method you can also use a binary executable")
    parser.add_argument('-m', "--method", type=str, default='ssh', help="\"ssh\" or \"startup\" script methods")
    parser.add_argument('-n', "--name", type=str, default='Scripted Terminal', help="The name of your Terminal")
    parser.add_argument('-k', "--ssh_key_file", type=str, default=None, help="Use your own ssh key instead of create a new one - \
    Use your private key name")
    parser.description="Utility to start and setup Terminals"
    args = parser.parse_args()
    terminal.setup_credentials(get_credentials(args.utoken, args.atoken, args.creds),None)
    args_sanitizer(args)

    # Preparing
    if args.method == 'ssh':
        if args.ssh_key_file is None:
            generate_ssh_key(key_name)
        else:
            key_name=args.ssh_key_file
        publicKey=get_public_key('%s.pub' % key_name)
        script=None
    else:
        script=get_script(args.script)

    # Creating terminals
    terminals=[]
    for i in range(args.quantity):
        name = '%s-%s' % (args.name,i)
        container_key, container_ip, subdomain = start_snap(name, snapshot_id, args.size, script)
        terminals.append({'container_key':container_key, 'container_ip':container_ip, 'subdomain':subdomain, 'name':name})
    time.sleep(2) # Prevent race-condition issues

    # Installing stuff by ssh method
    if args.method == 'ssh':
        for i in range(len(terminals)):
            terminal.add_authorized_key_to_terminal(terminals[i]['container_key'],publicKey)
            time.sleep(1)
            send_script(terminals[i]['container_ip'], 'root', script)
            run_on_terminal(terminals[i]['container_ip'], 'user', '/bin/bash /root/%s' % os.path.basename(script))

    # Print results in json format
    print json.loads(terminals)<commit_msg>bugfixes<commit_after>#!/usr/bin/env python
import os
import json
import time
import argparse
import subprocess
import terminal

snapshot_id = '7067f369f7b76f0a3276beb561820a21c9b5204ab60fbd90524560db96d7cb38'
key_name = 'tempkey'

def get_credentials(utoken, atoken, credsfile='creds.json'):
    if utoken is None and atoken is None:
        try:
            creds = json.load(open(credsfile, 'r'))
            utoken = creds['user_token']
            atoken = creds['access_token']
        except:
            print "Can't open credentials file. \n ", \
                "You must provide a user token and a access token at least one time, or a valid credentials file"
            exit(127)
    elif (utoken is not None and atoken is None) or (utoken is None and atoken is not None):
        print "--utoken AND --atoken parameters must be passed together"
        exit(1)
    else:
        with open(credsfile, 'w') as cfile:
            json.dump({'user_token': utoken, 'access_token': atoken}, cfile)
    return utoken, atoken

def generate_ssh_key(key_file):
    subprocess.call(['ssh-keygen','-f', key_file,'-P',''])

def get_public_key(key_file):
    with open(key_file) as f:
        content = f.readlines()[0].rstrip('\n')
    return str(content)

def start_snap(name, snapshot_id, size=None, script=None):
    output = terminal.start_snapshot(snapshot_id, size, None, name, None, script)
    request_id = output['request_id']
    output = terminal.request_progress(request_id)
    state = 'None'
    time.sleep(1)
    while output['status'] != 'success':
        output = terminal.request_progress(request_id)
        if output['state'] != state:
            state = output['state']
            print('(%s)' % state)
        time.sleep(1)
    container_key = output['result']['container_key']
    subdomain = output['result']['subdomain']
    container_ip = output['result']['container_ip']
    return container_key, container_ip, subdomain

def run_on_terminal(cip, user, script):
    try:
        p = subprocess.Popen(
        ['ssh', '-q' ,'-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-i', 'pemfile',
         user + '@' + cip, script], 0, None,  None, None, None)
        p.wait()
        return p.returncode
    except Exception, e:
        return 'Error: [%s]' % e

def get_script(filename):
    try:
        with open(filename) as f:
            data = f.read()
            data.replace('\n',';')
        return data
    except Exception, e:
        print '(%s)' % e
        return None

def send_script(cip, user, script, path=''):
    try:
        p = subprocess.Popen(
        ['scp', script ,'-q' ,'-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-i', 'pemfile',
         '%s@%s:%s' % (user,cip,path)], 0, None,  None, None, None)
        p.wait()
        return p.returncode
    except Exception, e:
        return 'Error: [%s]' % e

def args_sanitizer(args):
    if args.size not in terminal.instance_types()['instance_types']:
        print "Error. Wrong instance type"
        exit(1)
    if (int(args.quantity) < 1 or int(args.quantity) > 100):
        print "Error. Terminal amount out of bounds ( should be in between 1 and 100 )"
        exit(1)
    if args.method not in ('ssh, startup'):
        print "Error. Not a valid method. (use ssh or startup)"
        exit(1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--size", type=str, default='medium', help="micro, mini, small, medium, xlarge.. etc")
    parser.add_argument("-q", "--quantity", type=int, default=3, help="How many nodes will have your deploy")
    parser.add_argument("-u", "--utoken", type=str, default=None ,help="Your user token")
    parser.add_argument("-a", "--atoken", type=str, default=None, help="Your access token")
    parser.add_argument("-c", "--creds", type=str, default=None, help="A credentials json file")
    parser.add_argument("-x", "--script", type=str, default=None, help="A script file to be executed in the new Terminals - \
    With ssh method you can also use a binary executable")
    parser.add_argument('-m', "--method", type=str, default='ssh', help="\"ssh\" or \"startup\" script methods")
    parser.add_argument('-n', "--name", type=str, default='Scripted Terminal', help="The name of your Terminal")
    parser.add_argument('-k', "--ssh_key_file", type=str, default=None, help="Use your own ssh key instead of create a new one - \
    Use your private key name")
    parser.description="Utility to start and setup Terminals"
    args = parser.parse_args()

    terminal.setup_credentials(args.utoken, args.atoken, args.creds)
    args_sanitizer(args)

    # Preparing
    if args.method == 'ssh':
        if args.ssh_key_file is None:
            generate_ssh_key(key_name)
        else:
            key_name=args.ssh_key_file
        publicKey=get_public_key('%s.pub' % key_name)
        script=None
    else:
        script=get_script(args.script)

    # Creating terminals
    terminals=[]
    for i in range(args.quantity):
        name = '%s-%s' % (args.name,i)
        container_key, container_ip, subdomain = start_snap(name, snapshot_id, args.size, script)
        terminals.append({'container_key':container_key, 'container_ip':container_ip, 'subdomain':subdomain, 'name':name})
    time.sleep(2) # Prevent race-condition issues

    # Installing stuff by ssh method
    if args.method == 'ssh':
        for i in range(len(terminals)):
            terminal.add_authorized_key_to_terminal(terminals[i]['container_key'],publicKey)
            time.sleep(1)
            send_script(terminals[i]['container_ip'], 'root', script)
            run_on_terminal(terminals[i]['container_ip'], 'user', '/bin/bash /root/%s' % os.path.basename(script))

    # Print results in json format
    print json.loads(terminals)<|endoftext|>