<commit_before>import psycopg2
import numpy as np
import pyiem.nws.vtec as vtec
from pyiem.network import Table as NetworkTable
from pandas.io.sql import read_sql
import calendar
from pyiem.util import get_autoplot_context

def get_description():
    """ Return a dict describing how to call this plotter """
    d = dict()
    d['data'] = True
    d['cache'] = 86400
    d['description'] = """This chart displays the monthly number of distinct
    NWS Office issued watch / warning / advisory product. For example, a
    single watch for a dozen counties would only count 1 in this chart. These
    values are based on unofficial archives maintained by the IEM."""
    d['arguments'] = [
        dict(type='networkselect', name='station', network='WFO',
             default='DMX', label='Select WFO:', all=True),
        dict(type='phenomena', name='phenomena',
             default='FF', label='Select Watch/Warning Phenomena Type:'),
        dict(type='significance', name='significance',
             default='W', label='Select Watch/Warning Significance Level:'),
    ]
    return d

def plotter(fdict):
    """ Go """
    import matplotlib
    matplotlib.use('agg')
    import matplotlib.pyplot as plt
    import matplotlib.colors as mpcolors
    import matplotlib.patheffects as PathEffects
    pgconn = psycopg2.connect(database='postgis', host='iemdb', user='nobody')
    ctx = get_autoplot_context(fdict, get_description())
    station = ctx['station']
    phenomena = ctx['phenomena']
    significance = ctx['significance']

    nt = NetworkTable('WFO')
    nt.sts['_ALL'] = {'name': 'All Offices'}

    wfo_limiter = (" and wfo = '%s' "
                   ) % (station if len(station) == 3 else station[1:],)
    if station == '_ALL':
        wfo_limiter = ''

    df = read_sql("""
        with data as (
            SELECT distinct extract(year from issue) as yr2,
            min(issue) as i, wfo, eventid
            from warnings where phenomena = %s and significance = %s
            """ + wfo_limiter + """ and issue is not null
            GROUP by yr2, wfo, eventid)

        SELECT extract(year from i) as yr, extract(month from i) as mo,
        count(*) from data GROUP by yr, mo ORDER by yr, mo ASC
      """, pgconn, params=(phenomena, significance), index_col=None)

    if len(df.index) == 0:
        return("Sorry, no data found!")
    (fig, ax) = plt.subplots(1, 1, figsize=(8, 8))

    minyear = df['yr'].min()
    maxyear = df['yr'].max()
    data = np.zeros((maxyear - minyear + 1, 12))
    for _, row in df.iterrows():
        data[row['yr'] - minyear, row['mo'] - 1] = row['count']
        txt = ax.text(row['mo'], row['yr'], "%.0f" % (row['count'],),
                      va='center', ha='center', color='white')
        txt.set_path_effects([PathEffects.withStroke(linewidth=2,
                                                     foreground="k")])

    cmap = plt.get_cmap('jet')
    cmap.set_under('white')
    bounds = np.linspace(1, df['count'].max(), 10, dtype='i')
    norm = mpcolors.BoundaryNorm(bounds, cmap.N)
    res = ax.imshow(data, extent=[0.5, 12.5, maxyear + 0.5, minyear - 0.5],
                    interpolation='nearest', aspect='auto', norm=norm)
    fig.colorbar(res, label='count')
    ax.grid(True)
    ax.set_xticks(range(1, 13))
    ax.set_xticklabels(calendar.month_abbr[1:])
    ax.set_title(("NWS %s\n%s %s (%s.%s) Issued by Year,Month"
                  ) % (nt.sts[station]['name'], vtec._phenDict[phenomena],
                       vtec._sigDict[significance], phenomena, significance))

    return fig, df
<commit_msg>correct ramp bug for values less than 10<commit_after>import psycopg2
import numpy as np
import pyiem.nws.vtec as vtec
from pyiem.network import Table as NetworkTable
from pandas.io.sql import read_sql
import calendar
from pyiem.util import get_autoplot_context

def get_description():
    """ Return a dict describing how to call this plotter """
    d = dict()
    d['data'] = True
    d['cache'] = 86400
    d['description'] = """This chart displays the monthly number of distinct
    NWS Office issued watch / warning / advisory product. For example, a
    single watch for a dozen counties would only count 1 in this chart. These
    values are based on unofficial archives maintained by the IEM."""
    d['arguments'] = [
        dict(type='networkselect', name='station', network='WFO',
             default='DMX', label='Select WFO:', all=True),
        dict(type='phenomena', name='phenomena',
             default='FF', label='Select Watch/Warning Phenomena Type:'),
        dict(type='significance', name='significance',
             default='W', label='Select Watch/Warning Significance Level:'),
    ]
    return d

def plotter(fdict):
    """ Go """
    import matplotlib
    matplotlib.use('agg')
    import matplotlib.pyplot as plt
    import matplotlib.colors as mpcolors
    import matplotlib.patheffects as PathEffects
    pgconn = psycopg2.connect(database='postgis', host='iemdb', user='nobody')
    ctx = get_autoplot_context(fdict, get_description())
    station = ctx['station']
    phenomena = ctx['phenomena']
    significance = ctx['significance']

    nt = NetworkTable('WFO')
    nt.sts['_ALL'] = {'name': 'All Offices'}

    wfo_limiter = (" and wfo = '%s' "
                   ) % (station if len(station) == 3 else station[1:],)
    if station == '_ALL':
        wfo_limiter = ''

    df = read_sql("""
        with data as (
            SELECT distinct extract(year from issue) as yr2,
            min(issue) as i, wfo, eventid
            from warnings where phenomena = %s and significance = %s
            """ + wfo_limiter + """ and issue is not null
            GROUP by yr2, wfo, eventid)

        SELECT extract(year from i) as yr, extract(month from i) as mo,
        count(*) from data GROUP by yr, mo ORDER by yr, mo ASC
      """, pgconn, params=(phenomena, significance), index_col=None)
    if len(df.index) == 0:
        return("Sorry, no data found!")
    (fig, ax) = plt.subplots(1, 1, figsize=(8, 8))

    minyear = df['yr'].min()
    maxyear = df['yr'].max()
    data = np.zeros((maxyear - minyear + 1, 12))
    for _, row in df.iterrows():
        data[row['yr'] - minyear, row['mo'] - 1] = row['count']
        txt = ax.text(row['mo'], row['yr'], "%.0f" % (row['count'],),
                      va='center', ha='center', color='white')
        txt.set_path_effects([PathEffects.withStroke(linewidth=2,
                                                     foreground="k")])
    cmap = plt.get_cmap('jet')
    cmap.set_under('white')
    maxval = max([df['count'].max(), 11])
    bounds = np.linspace(1, maxval, 10, dtype='i')
    norm = mpcolors.BoundaryNorm(bounds, cmap.N)
    res = ax.imshow(data, extent=[0.5, 12.5, maxyear + 0.5, minyear - 0.5],
                    interpolation='nearest', aspect='auto', norm=norm)
    fig.colorbar(res, label='count')
    ax.grid(True)
    ax.set_xticks(range(1, 13))
    ax.set_xticklabels(calendar.month_abbr[1:])
    ax.set_title(("NWS %s\n%s %s (%s.%s) Issued by Year,Month"
                  ) % (nt.sts[station]['name'], vtec._phenDict[phenomena],
                       vtec._sigDict[significance], phenomena, significance))

    return fig, df

if __name__ == '__main__':
    plotter(dict(wfo='DMX', network='WFO', phenomena='FG', significance='Y'))
<|endoftext|>


<commit_before>from copy import deepcopy

from .messages import TextMessage, ImageMessage, LocationMessage, VoiceMessage
from .parser import parse_user_msg

__all__ = ['WeTest']
message_temple = {
    'msgid': 1234567890123456,
    'ToUserName': 'test',
    'FromUserName': 'test',
    'CreateTime': 123456789,
}

class WeTest(object):
    def __init__(self, app):
        self._app = app

    def send(self, message):
        return self._app.get_reply(message)

    def send_xml(self, xml):
        message = parse_user_msg(xml)
        return self.send(message)

def make_text_message(content):
    message = deepcopy(message_temple)
    message["content"] = content
    message = TextMessage(message)
    return message

def make_image_message(img):
    message = deepcopy(message_temple)
    message["img"] = img
    message = ImageMessage(message)
    return message

def make_location_message(x, y, scale, label):
    message = deepcopy(message_temple)
    message["x"] = x
    message["y"] = y
    message["scale"] = scale
    message["label"] = label
    message = LocationMessage(message)
    return message

def make_voice_message(media_id, format, recognition):
    message = deepcopy(message_temple)
    message["media_id"] = media_id
    message["format"] = format
    message["recognition"] = recognition
    message = VoiceMessage(message)
    return message
<commit_msg>Drop a lot of funcs in ``werobot.testing``<commit_after>from .parser import parse_user_msg

__all__ = ['WeTest']

class WeTest(object):
    def __init__(self, app):
        self._app = app

    def send_xml(self, xml):
        message = parse_user_msg(xml)
        return self._app.get_reply(message)
<|endoftext|>


<commit_before>import sublime, sublime_plugin
import re
from ROBLOXLua.apiparser import parse_api_dump

completion_items = parse_api_dump()

FUNCTION_CALL_REGEX = re.compile(r"\b([\w\.]+)\s*\(?[\"\']")

service_detections = [ "GetService", "FindService", "getService", "service" ]
creatable_detections = [ "Instance.new" ]
services = set([ e["entry_completion"] for e in completion_items if e["entry_type"] == "Class" and "service" in e["entry_tags"] ])
creatables = set([ e["entry_completion"] for e in completion_items if e["entry_type"] == "Class" and ("notCreatable" not in e["entry_tags"] and "abstract" not in e["entry_tags"] and "service" not in e["entry_tags"]) ])

print(creatables)

class AutoCompleteProvider(sublime_plugin.EventListener):
	"""
	The automatic completions provider.
	"""
	def on_query_completions(self, view, prefix, locations):
		selected_completions = set()

		for location in locations:
			if "source.rbxlua" not in view.scope_name(location):
				return

			line = view.substr(view.line(location))
			function_match = FUNCTION_CALL_REGEX.search(line)

			if function_match is not None:
				function_name = function_match.group(1)

				if function_name in service_detections:
					selected_completions.update(services)
				elif function_name in creatable_detections:
					selected_completions.update(creatables)

		if len(selected_completions) > 0:
			return ([ [e] for e in selected_completions ], sublime.INHIBIT_EXPLICIT_COMPLETIONS | sublime.INHIBIT_WORD_COMPLETIONS)<commit_msg>Made creatables filter stricter<commit_after>import sublime, sublime_plugin
import re
from ROBLOXLua.apiparser import parse_api_dump

completion_items = parse_api_dump()

FUNCTION_CALL_REGEX = re.compile(r"\b([\w\.]+)\s*\(?[\"\']")

service_detections = [ "GetService", "FindService", "getService", "service" ]
creatable_detections = [ "Instance.new" ]
services = set([ e["entry_completion"] for e in completion_items if e["entry_type"] == "Class" and "service" in e["entry_tags"] ])
creatables = set([ e["entry_completion"] for e in completion_items if e["entry_type"] == "Class" and ("notCreatable" not in e["entry_tags"] and "abstract" not in e["entry_tags"] and "service" not in e["entry_tags"] and "deprecated" not in e["entry_tags"]) ])

class AutoCompleteProvider(sublime_plugin.EventListener):
	"""
	The automatic completions provider.
	"""
	def on_query_completions(self, view, prefix, locations):
		selected_completions = set()

		for location in locations:
			if "source.rbxlua" not in view.scope_name(location):
				return

			line = view.substr(view.line(location))
			function_match = FUNCTION_CALL_REGEX.search(line)

			if function_match is not None:
				function_name = function_match.group(1)

				if function_name in service_detections:
					selected_completions.update(services)
				elif function_name in creatable_detections:
					selected_completions.update(creatables)

		if len(selected_completions) > 0:
			return ([ [e] for e in selected_completions ], sublime.INHIBIT_EXPLICIT_COMPLETIONS | sublime.INHIBIT_WORD_COMPLETIONS)<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Handle loading six package from system or from the bundled copy

"""
import sys

# Trying to load alternate six packages
sys.modules['astropy.extern.six'] = None

# We have removed everything we already imported
# Importing again

import sys

def _load_six_moves(base, dest):
    _cur_sys_modules = list(sys.modules.items())
    for i,mod in _cur_sys_modules:
        if i.startswith(base):
            pre, full, trail = i.partition(base)
            if not pre:
                modname = dest + trail
                sys.modules[modname] = mod

six_system_package = False

_dest_moves = 'astropy.extern.six.moves'
_dest_root = 'astropy.extern.six'

try:
    import six
    _system_package = True
except ImportError:
    _system_package = False

if _system_package:
    # Check six version
    _valid_version = True
    if _valid_version:
        # handle 'moves'
        _base_moves = 'six.moves'
        _load_six_moves(_base_moves, _dest_moves)
        six.system_package = True
        six.bundled_package = False
        sys.modules[_dest_root] = six

if not _system_package:
    import astropy.extern.bundled.six as six
    # handle 'moves'
    _base_moves = 'astropy.extern.bundled.six.moves'
    _load_six_moves(_base_moves, _dest_moves)
    six.system_package = False
    six.bundled_package = True
    sys.modules[_dest_root] = six

<commit_msg>Handle the case where system six exists but __version__ is not valid<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Handle loading six package from system or from the bundled copy

"""
import sys

# Trying to load alternate six packages
sys.modules['astropy.extern.six'] = None

# We have removed everything we already imported
# Importing again

import sys

def _load_six_moves(base, dest):
    _cur_sys_modules = list(sys.modules.items())
    for i,mod in _cur_sys_modules:
        if i.startswith(base):
            pre, full, trail = i.partition(base)
            if not pre:
                modname = dest + trail
                sys.modules[modname] = mod

_dest_moves = 'astropy.extern.six.moves'
_dest_root = 'astropy.extern.six'

_system_package = False
try:
    import six
    _system_package = True
except ImportError:
    _system_package = False

if _system_package:
    # Check six version
    _valid_version = True
    if _valid_version:
        # handle 'moves'
        _base_moves = 'six.moves'
        _load_six_moves(_base_moves, _dest_moves)
        six.system_package = True
        six.bundled_package = False
        sys.modules[_dest_root] = six
    else:
        _system_package = False

if not _system_package:
    import astropy.extern.bundled.six as six
    # handle 'moves'
    _base_moves = 'astropy.extern.bundled.six.moves'
    _load_six_moves(_base_moves, _dest_moves)
    six.system_package = False
    six.bundled_package = True
    sys.modules[_dest_root] = six

<|endoftext|>


<commit_before>import charts
import datetime
import json
import os
import whisky
from flask import render_template, redirect, Response, request, abort
from flask import make_response
from sqlalchemy import desc
from whiskyton import app, models

@app.route('/')
def index():
    random_one = whisky.random_whisky()
    return render_template(
        'home.html',
        main_title=app.config['MAIN_TITLE'],
        headline=app.config['HEADLINE'],
        remote_scripts=app.config['GOOGLE_ANALYTICS'],
        random_one=random_one)

@app.route('/search', methods=['GET', 'POST'])
def search():
    slug = whisky.slugfy(request.args['s'])
    w = models.Whisky.query.filter_by(slug=slug).first()
    if w is None:
        random_one = whisky.random_whisky()
        return render_template(
            '404.html',
            main_title=app.config['MAIN_TITLE'],
            headline=app.config['HEADLINE'],
            remote_scripts=app.config['GOOGLE_ANALYTICS'],
            slug=slug,
            random_one=random_one)
    else:
        return redirect('/' + str(w.slug))

@app.route('/<whisky_slug>')
def whisky_page(whisky_slug):

    slugfied = whisky.slugfy(whisky_slug)
    if whisky_slug != slugfied:
        return redirect('/' + slugfied)

    reference = models.Whisky.query.filter_by(slug=whisky_slug).first()

    # error page if whisky doesn't exist
    if reference is None:
        return abort(404)

    # load correlations
    else:

        # query
        whiskies = models.Correlation.query\
            .add_entity(models.Whisky)\
            .filter(models.Correlation.reference == reference.id)\
            .filter(models.Correlation.r > 0.5)\
            .join(models.Whisky, models.Correlation.whisky == models.Whisky.id)\
            .order_by(desc(models.Correlation.r))\
            .limit(9)

        # if query succeeded
        if whiskies is not None:

            # build result
            title = 'Whiskies for %s lovers | %s' % (reference.distillery,
                                                     app.config['MAIN_TITLE'])
            return render_template(
                'whiskies.html',
                main_title=title,
                headline=app.config['HEADLINE'],
                remote_scripts=app.config['GOOGLE_ANALYTICS'],
                whiskies=whiskies,
                reference=reference,
                count=whiskies.count(),
                result_page=True)

        # if queries fail, return 404
        else:
            return abort(404)

@app.route('/w/<whiskyID>')
def searchID(whiskyID):
    reference = models.Whisky.query.filter_by(id=whiskyID).first()
    if reference is None:
        return abort(404)
    else:
        return redirect('/' + reference.slug)

@app.route('/charts/<reference_slug>-<whisky_slug>.svg')
def create_chart(reference_slug, whisky_slug):

    # URL check
    slug1 = whisky.slugfy(whisky_slug)
    slug2 = whisky.slugfy(reference_slug)
    if slug1 != whisky_slug or slug2 != reference_slug:
        return redirect('/charts/%s-%s.svg' % (slug1, slug2))

    # get whisky objects form db
    reference_obj = models.Whisky.query.filter_by(slug=reference_slug).first()
    whisky_obj = models.Whisky.query.filter_by(slug=whisky_slug).first()

    # error page if whisky doesn't exist
    if reference_obj is None or whisky_obj is None:
        return abort(404)

    # if file does not exists, create it
    reference = charts.tastes2list(reference_obj)
    comparison = charts.tastes2list(whisky_obj)
    filename = charts.cache_name(reference, comparison, True)
    if not filename.exists():
        charts.create(reference, comparison)

    # return the chart to the user
    return Response(filename.read_file(), mimetype='image/svg+xml')

@app.route('/whiskyton.json')
def whisky_json():
    whiskies = models.Whisky.query.all()
    distilleries = json.dumps([whisky.distillery for whisky in whiskies])
    resp = Response(
        response=distilleries,
        status=200,
        mimetype='application/json')
    return resp

@app.route('/robots.txt', methods=['GET'])
def robots():
    response = make_response(open('robots.txt').read())
    response.headers["Content-type"] = "text/plain"
    return response

@app.route('/sitemap.xml')
def sitemap():
    whiskies = models.Whisky.query.all()
    ref_file = 'whiskyton/views.py'
    dt_unix = int(os.path.getmtime(ref_file))
    last_change = datetime.datetime.fromtimestamp(dt_unix).strftime('%Y-%m-%d')
    return render_template(
        'sitemap.xml',
        whiskies=whiskies,
        last_change=last_change,
        url_root=request.url_root)

@app.errorhandler(404)
def page_not_found(e):
    random_one = whisky.random_whisky()
    return render_template(
        '404.html',
        main_title=app.config['MAIN_TITLE'],
        headline=app.config['HEADLINE'],
        remote_scripts=app.config['GOOGLE_ANALYTICS'],
        random_one=random_one), 404
<commit_msg>Use unipath for robots.txt<commit_after>import charts
import datetime
import json
import os
import whisky
from flask import render_template, redirect, Response, request, abort
from flask import make_response
from sqlalchemy import desc
from whiskyton import app, models

@app.route('/')
def index():
    random_one = whisky.random_whisky()
    return render_template(
        'home.html',
        main_title=app.config['MAIN_TITLE'],
        headline=app.config['HEADLINE'],
        remote_scripts=app.config['GOOGLE_ANALYTICS'],
        random_one=random_one)

@app.route('/search', methods=['GET', 'POST'])
def search():
    slug = whisky.slugfy(request.args['s'])
    w = models.Whisky.query.filter_by(slug=slug).first()
    if w is None:
        random_one = whisky.random_whisky()
        return render_template(
            '404.html',
            main_title=app.config['MAIN_TITLE'],
            headline=app.config['HEADLINE'],
            remote_scripts=app.config['GOOGLE_ANALYTICS'],
            slug=slug,
            random_one=random_one)
    else:
        return redirect('/' + str(w.slug))

@app.route('/<whisky_slug>')
def whisky_page(whisky_slug):

    slugfied = whisky.slugfy(whisky_slug)
    if whisky_slug != slugfied:
        return redirect('/' + slugfied)

    reference = models.Whisky.query.filter_by(slug=whisky_slug).first()

    # error page if whisky doesn't exist
    if reference is None:
        return abort(404)

    # load correlations
    else:

        # query
        whiskies = models.Correlation.query\
            .add_entity(models.Whisky)\
            .filter(models.Correlation.reference == reference.id)\
            .filter(models.Correlation.r > 0.5)\
            .join(models.Whisky, models.Correlation.whisky == models.Whisky.id)\
            .order_by(desc(models.Correlation.r))\
            .limit(9)

        # if query succeeded
        if whiskies is not None:

            # build result
            title = 'Whiskies for %s lovers | %s' % (reference.distillery,
                                                     app.config['MAIN_TITLE'])
            return render_template(
                'whiskies.html',
                main_title=title,
                headline=app.config['HEADLINE'],
                remote_scripts=app.config['GOOGLE_ANALYTICS'],
                whiskies=whiskies,
                reference=reference,
                count=whiskies.count(),
                result_page=True)

        # if queries fail, return 404
        else:
            return abort(404)

@app.route('/w/<whiskyID>')
def searchID(whiskyID):
    reference = models.Whisky.query.filter_by(id=whiskyID).first()
    if reference is None:
        return abort(404)
    else:
        return redirect('/' + reference.slug)

@app.route('/charts/<reference_slug>-<whisky_slug>.svg')
def create_chart(reference_slug, whisky_slug):

    # URL check
    slug1 = whisky.slugfy(whisky_slug)
    slug2 = whisky.slugfy(reference_slug)
    if slug1 != whisky_slug or slug2 != reference_slug:
        return redirect('/charts/%s-%s.svg' % (slug1, slug2))

    # get whisky objects form db
    reference_obj = models.Whisky.query.filter_by(slug=reference_slug).first()
    whisky_obj = models.Whisky.query.filter_by(slug=whisky_slug).first()

    # error page if whisky doesn't exist
    if reference_obj is None or whisky_obj is None:
        return abort(404)

    # if file does not exists, create it
    reference = charts.tastes2list(reference_obj)
    comparison = charts.tastes2list(whisky_obj)
    filename = charts.cache_name(reference, comparison, True)
    if not filename.exists():
        charts.create(reference, comparison)

    # return the chart to the user
    return Response(filename.read_file(), mimetype='image/svg+xml')

@app.route('/whiskyton.json')
def whisky_json():
    whiskies = models.Whisky.query.all()
    distilleries = json.dumps([whisky.distillery for whisky in whiskies])
    resp = Response(
        response=distilleries,
        status=200,
        mimetype='application/json')
    return resp

@app.route('/robots.txt', methods=['GET'])
def robots():
    robots = app.config['BASEDIR'].child('robots.txt').read_file()
    response = make_response(robots)
    response.headers["Content-type"] = "text/plain"
    return response

@app.route('/sitemap.xml')
def sitemap():
    whiskies = models.Whisky.query.all()
    ref_file = 'whiskyton/views.py'
    dt_unix = int(os.path.getmtime(ref_file))
    last_change = datetime.datetime.fromtimestamp(dt_unix).strftime('%Y-%m-%d')
    return render_template(
        'sitemap.xml',
        whiskies=whiskies,
        last_change=last_change,
        url_root=request.url_root)

@app.errorhandler(404)
def page_not_found(e):
    random_one = whisky.random_whisky()
    return render_template(
        '404.html',
        main_title=app.config['MAIN_TITLE'],
        headline=app.config['HEADLINE'],
        remote_scripts=app.config['GOOGLE_ANALYTICS'],
        random_one=random_one), 404
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""WAL-E is a program to assist in performing PostgreSQL continuous
archiving on S3: it handles the major four operations of
arching/receiving WAL segments and archiving/receiving base hot
backups of the PostgreSQL file cluster.

"""

def gevent_monkey(*args, **kwargs):
    import gevent.monkey
    gevent.monkey.patch_socket(dns=True, aggressive=True)
    gevent.monkey.patch_ssl()
    gevent.monkey.patch_time()

# Monkey-patch procedures early.  If it doesn't work with gevent,
# sadly it cannot be used (easily) in WAL-E.
gevent_monkey()

import argparse
import logging
import os
import re
import subprocess
import sys
import textwrap

import wal_e.log_help as log_help

from wal_e.exception import UserException
from wal_e.operator import s3_operator
from wal_e.piper import popen_sp
from wal_e.worker.psql_worker import PSQL_BIN, psql_csv_run
from wal_e.worker.s3_worker import LZOP_BIN, MBUFFER_BIN

# TODO: Make controllable from userland
log_help.configure(
    format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s')

logger = log_help.WalELogger('wal_e.main', level=logging.INFO)

def external_program_check(
    to_check=frozenset([PSQL_BIN, LZOP_BIN, MBUFFER_BIN])):
    """
    Validates the existence and basic working-ness of other programs

    Implemented because it is easy to get confusing error output when
    one does not install a dependency because of the fork-worker model
    that is both necessary for throughput and makes more obscure the
    cause of failures.  This is intended to be a time and frustration
    saving measure.  This problem has confused The Author in practice
    when switching rapidly between machines.

    """

    could_not_run = []
    error_msgs = []

    def psql_err_handler(popen):
        assert popen.returncode != 0
        error_msgs.append(textwrap.fill(
                'Could not get a connection to the database: '
                'note that superuser access is required'))

        # Bogus error message that is re-caught and re-raised
        raise Exception('INTERNAL: Had problems running psql '
                        'from external_program_check')

    with open(os.devnull, 'w') as nullf:
        for program in to_check:
            try:
                if program is PSQL_BIN:
                    psql_csv_run('SELECT 1', error_handler=psql_err_handler)
                else:
                    if program is MBUFFER_BIN:
                        # Prevent noise on the TTY, as mbuffer writes
                        # text there: suppressing stdout/stderr
                        # doesn't seem to work.
                        extra_args = ['-q']
                    else:
                        extra_args = []

                    proc = popen_sp([program] + extra_args,
                                    stdout=nullf, stderr=nullf,
                                    stdin=subprocess.PIPE)

                    # Close stdin for processes that default to
                    # reading from the pipe; the programs WAL-E uses
                    # of this kind will terminate in this case.
                    proc.stdin.close()
                    proc.wait()
            except EnvironmentError:
                could_not_run.append(program)

    if could_not_run:
        error_msgs.append(
            'Could not run the following programs, are they installed? ' +
            ', '.join(could_not_run))

    if error_msgs:
        raise UserException(
            'could not run one or more external programs WAL-E depends upon',
            '\n'.join(error_msgs))

    return None

def extract_segment(text_with_extractable_segment):
    from wal_e.storage.s3_storage import BASE_BACKUP_REGEXP
    from wal_e.storage.s3_storage import SegmentNumber

    match = re.match(BASE_BACKUP_REGEXP, text_with_extractable_segment)
    if match is None:
        return None
    else:
        groupdict = match.groupdict()
        return SegmentNumber(log=groupdict['log'], seg=groupdict['seg'])

def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__)

    parser.add_argument('-k', '--aws-access-key-id',
                        help='public AWS access key. Can also be defined in '
                        'an environment variable. If both are defined, '
                        'the one defined in the programs arguments takes '
                        'precedence.')

    parser.add_argument('--s3-prefix',
                        help='S3 prefix to run all commands against.  '
                        'Can also be defined via environment variable '
                        'WALE_S3_PREFIX')

    subparsers = parser.add_subparsers(title='subcommands',
                                       dest='subcommand')

    # Common arguments for backup-fetch and backup-push
    backup_fetchpush_parent = argparse.ArgumentParser(add_help=False)
    backup_fetchpush_parent.add_argument('PG_CLUSTER_DIRECTORY',
                                         help="Postgres cluster path, "
                                         "such as '/var/lib/database'")
    backup_fetchpush_parent.add_argument('--pool-size', '-p',
                                         type=int, default=4,
                                         help='Download pooling size')

    # Common arguments for backup-list and backup-fetch
    #
    # NB: This does not include the --detail options because some
    # other commands use backup listing functionality in a way where
    # --detail is never required.
    backup_list_nodetail_parent = argparse.ArgumentParser(add_help=False)

    # Common arguments between wal-push and wal-fetch
    wal_fetchpush_parent = argparse.ArgumentParser(add_help=False)
    wal_fetchpush_parent.add_argument('WAL_SEGMENT',
                                      help='Path to a WAL segment to upload')

    backup_destroy_parser = subparsers.add_parser(
        'backup-destroy', help='destroy old backups in S3')
    backup_fetch_parser = subparsers.add_parser(
        'backup-fetch', help='fetch a hot backup from S3',
        parents=[backup_fetchpush_parent, backup_list_nodetail_parent])
    backup_list_parser = subparsers.add_parser(
        'backup-list', parents=[backup_list_nodetail_parent],
        help='list backups in S3')
    backup_push_parser = subparsers.add_parser(
        'backup-push', help='pushing a fresh hot backup to S3',
        parents=[backup_fetchpush_parent])
    backup_push_parser.add_argument(
        '--cluster-read-rate-limit',
        help='Rate limit reading the PostgreSQL cluster directory to a '
        'tunable number of bytes per second', dest='rate_limit',
        metavar='BYTES_PER_SECOND',
        type=int, default=None)

    wal_fetch_parser = subparsers.add_parser(
        'wal-fetch', help='fetch a WAL file from S3',
        parents=[wal_fetchpush_parent])
    subparsers.add_parser('wal-push', help='push a WAL file to S3',
                          parents=[wal_fetchpush_parent])

    # backup-fetch operator section
    backup_fetch_parser.add_argument('BACKUP_NAME',
                                     help='the name of the backup to fetch')

    # backup-list operator section
    backup_list_parser.add_argument(
        'QUERY', nargs='?', default=None,
        help='a string qualifying backups to list')
    backup_list_parser.add_argument(
        '--detail', default=False, action='store_true',
        help='show more detailed information about every backup')

    # wal-push operator section
    wal_fetch_parser.add_argument('WAL_DESTINATION',
                                  help='Path to download the WAL segment to')

    # delete subparser section
    delete_parser = subparsers.add_parser(
        'delete', help=('operators to destroy specified data in S3'))
    delete_parser.add_argument('--dry-run', '-n', action='store_true',
                               help=('Only print what would be deleted, '
                                     'do not actually delete anything'))
    delete_parser.add_argument('--force', '-f', action='store_true',
                               help=('Actually delete data.  '
                                     'By default, a dry run is performed.  '
                                     'Overridden by --dry-run.'))
    delete_subparsers = delete_parser.add_subparsers(
        title='delete subcommands',
        description=('All operators that may delete data are contained '
                     'in this subcommand.'),
        dest='delete_subcommand')

    # delete 'before' operator
    delete_before_parser = delete_subparsers.add_parser(
        'before', help=('Delete all backups and WAL segments strictly before '
                        'the given base backup name or WAL segment number.  '
                        'The passed backup is *not* deleted.'))
    delete_before_parser.add_argument(
        'BEFORE_SEGMENT_EXCLUSIVE',
        help='A WAL segment number or base backup name')

    # delete old versions operator
    delete_old_versions_parser = delete_subparsers.add_parser('old-versions')

    # delete *everything* operator
    delete_old_versions_parser = delete_subparsers.add_parser('everything')

    # Okay, parse some arguments, finally
    args = parser.parse_args()

    # Attempt to read a few key parameters from environment variables
    # *or* the command line, enforcing a precedence order and
    # complaining should the required parameter not be defined in
    # either location.
    secret_key = os.getenv('AWS_SECRET_ACCESS_KEY')
    if secret_key is None:
        logger.error(
            msg='no AWS_SECRET_ACCESS_KEY defined',
            hint='Define the environment variable AWS_SECRET_ACCESS_KEY.')
        sys.exit(1)

    s3_prefix = args.s3_prefix or os.getenv('WALE_S3_PREFIX')

    if s3_prefix is None:
        logger.error(
            msg='no storage prefix defined',
            hint=('Either set the --s3-prefix option or define the '
                  'environment variable WALE_S3_PREFIX.'))
        sys.exit(1)

    if args.aws_access_key_id is None:
        aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
        if aws_access_key_id is None:
            logger.error(
                msg='no storage prefix defined',
                hint=('Either set the --aws-access-key-id option or define '
                      'the environment variable AWS_ACCESS_KEY_ID.'))
            sys.exit(1)
    else:
        aws_access_key_id = args.aws_access_key_id

    backup_cxt = s3_operator.S3Backup(aws_access_key_id, secret_key, s3_prefix)

    subcommand = args.subcommand

    try:
        if subcommand == 'backup-fetch':
            external_program_check([LZOP_BIN])
            backup_cxt.database_s3_fetch(
                args.PG_CLUSTER_DIRECTORY,
                args.BACKUP_NAME,
                pool_size=args.pool_size)
        elif subcommand == 'backup-list':
            backup_cxt.backup_list(query=args.QUERY, detail=args.detail)
        elif subcommand == 'backup-push':
            external_program_check([LZOP_BIN, PSQL_BIN, MBUFFER_BIN])
            rate_limit = args.rate_limit
            if rate_limit is not None and rate_limit < 8192:
                logger.error(
                    msg='bad rate limit passed',
                    detail='The passed rate limit was {0}'.format(rate_limit),
                    hint=('Pass a rate limit that is positive and equal or '
                          'greater than 8192'))
                sys.exit(1)

            backup_cxt.database_s3_backup(
                args.PG_CLUSTER_DIRECTORY, rate_limit=rate_limit,
                pool_size=args.pool_size)
        elif subcommand == 'wal-fetch':
            external_program_check([LZOP_BIN])
            backup_cxt.wal_s3_restore(args.WAL_SEGMENT, args.WAL_DESTINATION)
        elif subcommand == 'wal-push':
            external_program_check([LZOP_BIN])
            backup_cxt.wal_s3_archive(args.WAL_SEGMENT)
        elif subcommand == 'delete':
            # Set up pruning precedence, optimizing for *not* deleting data
            print args

            # Canonicalize the passed arguments into the value
            # "is_dry_run_really"
            if args.dry_run is False and args.force is True:
                # Actually delete data *only* if there are *no* --dry-runs
                # present and --force is present.
                logger.info(msg='deleting data in S3')
                is_dry_run_really = False
            else:
                logger.info(msg='performing dry run of S3 data deletion')
                is_dry_run_really = True

                import boto.s3.key

                # This is not necessary, but "just in case" to find bugs.
                def just_error(*args, **kwargs):
                    assert False, ('About to delete something in '
                                   'dry-run mode.  Please report a bug.')

                boto.s3.key.Key.delete = just_error

            # Handle the subcommands and route them to the right
            # implementations.
            if args.delete_subcommand == 'old-versions':
                backup_cxt.delete_old_versions(is_dry_run_really)
            elif args.delete_subcommand == 'everything':
                backup_cxt.delete_all(is_dry_run_really)
            elif args.delete_subcommand == 'before':
                segment_info = extract_segment(args.BEFORE_SEGMENT_EXCLUSIVE)
                backup_cxt.delete_before(is_dry_run_really, segment_info)
            else:
                assert False, 'Should be rejected by argument parsing.'
        else:
            logger.error(msg='subcommand not implemented',
                         detail=('The submitted subcommand was {0}.'
                                 .format(subcommand)),
                         hint='Check for typos or consult wal-e --help.')
            sys.exit(127)

        # Report on all encountered exceptions, and raise the last one
        # to take advantage of the final catch-all reporting and exit
        # code management.
        if backup_cxt.exceptions:
            for exc in backup_cxt.exceptions[:-1]:
                logger.log(level=exc.severity,
                           msg=log_help.WalELogger
                           .fmt_logline(exc.msg, exc.detail, exc.hint))
            raise backup_cxt.exceptions[-1]

    except UserException, e:
        logger.log(level=e.severity,
                   msg=log_help.WalELogger
                   .fmt_logline(e.msg, e.detail, e.hint))
        sys.exit(1)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Remove a stray debugging print<commit_after>#!/usr/bin/env python
"""WAL-E is a program to assist in performing PostgreSQL continuous
archiving on S3: it handles the major four operations of
arching/receiving WAL segments and archiving/receiving base hot
backups of the PostgreSQL file cluster.

"""

def gevent_monkey(*args, **kwargs):
    import gevent.monkey
    gevent.monkey.patch_socket(dns=True, aggressive=True)
    gevent.monkey.patch_ssl()
    gevent.monkey.patch_time()

# Monkey-patch procedures early.  If it doesn't work with gevent,
# sadly it cannot be used (easily) in WAL-E.
gevent_monkey()

import argparse
import logging
import os
import re
import subprocess
import sys
import textwrap

import wal_e.log_help as log_help

from wal_e.exception import UserException
from wal_e.operator import s3_operator
from wal_e.piper import popen_sp
from wal_e.worker.psql_worker import PSQL_BIN, psql_csv_run
from wal_e.worker.s3_worker import LZOP_BIN, MBUFFER_BIN

# TODO: Make controllable from userland
log_help.configure(
    format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s')

logger = log_help.WalELogger('wal_e.main', level=logging.INFO)

def external_program_check(
    to_check=frozenset([PSQL_BIN, LZOP_BIN, MBUFFER_BIN])):
    """
    Validates the existence and basic working-ness of other programs

    Implemented because it is easy to get confusing error output when
    one does not install a dependency because of the fork-worker model
    that is both necessary for throughput and makes more obscure the
    cause of failures.  This is intended to be a time and frustration
    saving measure.  This problem has confused The Author in practice
    when switching rapidly between machines.

    """

    could_not_run = []
    error_msgs = []

    def psql_err_handler(popen):
        assert popen.returncode != 0
        error_msgs.append(textwrap.fill(
                'Could not get a connection to the database: '
                'note that superuser access is required'))

        # Bogus error message that is re-caught and re-raised
        raise Exception('INTERNAL: Had problems running psql '
                        'from external_program_check')

    with open(os.devnull, 'w') as nullf:
        for program in to_check:
            try:
                if program is PSQL_BIN:
                    psql_csv_run('SELECT 1', error_handler=psql_err_handler)
                else:
                    if program is MBUFFER_BIN:
                        # Prevent noise on the TTY, as mbuffer writes
                        # text there: suppressing stdout/stderr
                        # doesn't seem to work.
                        extra_args = ['-q']
                    else:
                        extra_args = []

                    proc = popen_sp([program] + extra_args,
                                    stdout=nullf, stderr=nullf,
                                    stdin=subprocess.PIPE)

                    # Close stdin for processes that default to
                    # reading from the pipe; the programs WAL-E uses
                    # of this kind will terminate in this case.
                    proc.stdin.close()
                    proc.wait()
            except EnvironmentError:
                could_not_run.append(program)

    if could_not_run:
        error_msgs.append(
            'Could not run the following programs, are they installed? ' +
            ', '.join(could_not_run))

    if error_msgs:
        raise UserException(
            'could not run one or more external programs WAL-E depends upon',
            '\n'.join(error_msgs))

    return None

def extract_segment(text_with_extractable_segment):
    from wal_e.storage.s3_storage import BASE_BACKUP_REGEXP
    from wal_e.storage.s3_storage import SegmentNumber

    match = re.match(BASE_BACKUP_REGEXP, text_with_extractable_segment)
    if match is None:
        return None
    else:
        groupdict = match.groupdict()
        return SegmentNumber(log=groupdict['log'], seg=groupdict['seg'])

def main(argv=None):
    if argv is None:
        argv = sys.argv

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__)

    parser.add_argument('-k', '--aws-access-key-id',
                        help='public AWS access key. Can also be defined in '
                        'an environment variable. If both are defined, '
                        'the one defined in the programs arguments takes '
                        'precedence.')

    parser.add_argument('--s3-prefix',
                        help='S3 prefix to run all commands against.  '
                        'Can also be defined via environment variable '
                        'WALE_S3_PREFIX')

    subparsers = parser.add_subparsers(title='subcommands',
                                       dest='subcommand')

    # Common arguments for backup-fetch and backup-push
    backup_fetchpush_parent = argparse.ArgumentParser(add_help=False)
    backup_fetchpush_parent.add_argument('PG_CLUSTER_DIRECTORY',
                                         help="Postgres cluster path, "
                                         "such as '/var/lib/database'")
    backup_fetchpush_parent.add_argument('--pool-size', '-p',
                                         type=int, default=4,
                                         help='Download pooling size')

    # Common arguments for backup-list and backup-fetch
    #
    # NB: This does not include the --detail options because some
    # other commands use backup listing functionality in a way where
    # --detail is never required.
    backup_list_nodetail_parent = argparse.ArgumentParser(add_help=False)

    # Common arguments between wal-push and wal-fetch
    wal_fetchpush_parent = argparse.ArgumentParser(add_help=False)
    wal_fetchpush_parent.add_argument('WAL_SEGMENT',
                                      help='Path to a WAL segment to upload')

    backup_destroy_parser = subparsers.add_parser(
        'backup-destroy', help='destroy old backups in S3')
    backup_fetch_parser = subparsers.add_parser(
        'backup-fetch', help='fetch a hot backup from S3',
        parents=[backup_fetchpush_parent, backup_list_nodetail_parent])
    backup_list_parser = subparsers.add_parser(
        'backup-list', parents=[backup_list_nodetail_parent],
        help='list backups in S3')
    backup_push_parser = subparsers.add_parser(
        'backup-push', help='pushing a fresh hot backup to S3',
        parents=[backup_fetchpush_parent])
    backup_push_parser.add_argument(
        '--cluster-read-rate-limit',
        help='Rate limit reading the PostgreSQL cluster directory to a '
        'tunable number of bytes per second', dest='rate_limit',
        metavar='BYTES_PER_SECOND',
        type=int, default=None)

    wal_fetch_parser = subparsers.add_parser(
        'wal-fetch', help='fetch a WAL file from S3',
        parents=[wal_fetchpush_parent])
    subparsers.add_parser('wal-push', help='push a WAL file to S3',
                          parents=[wal_fetchpush_parent])

    # backup-fetch operator section
    backup_fetch_parser.add_argument('BACKUP_NAME',
                                     help='the name of the backup to fetch')

    # backup-list operator section
    backup_list_parser.add_argument(
        'QUERY', nargs='?', default=None,
        help='a string qualifying backups to list')
    backup_list_parser.add_argument(
        '--detail', default=False, action='store_true',
        help='show more detailed information about every backup')

    # wal-push operator section
    wal_fetch_parser.add_argument('WAL_DESTINATION',
                                  help='Path to download the WAL segment to')

    # delete subparser section
    delete_parser = subparsers.add_parser(
        'delete', help=('operators to destroy specified data in S3'))
    delete_parser.add_argument('--dry-run', '-n', action='store_true',
                               help=('Only print what would be deleted, '
                                     'do not actually delete anything'))
    delete_parser.add_argument('--force', '-f', action='store_true',
                               help=('Actually delete data.  '
                                     'By default, a dry run is performed.  '
                                     'Overridden by --dry-run.'))
    delete_subparsers = delete_parser.add_subparsers(
        title='delete subcommands',
        description=('All operators that may delete data are contained '
                     'in this subcommand.'),
        dest='delete_subcommand')

    # delete 'before' operator
    delete_before_parser = delete_subparsers.add_parser(
        'before', help=('Delete all backups and WAL segments strictly before '
                        'the given base backup name or WAL segment number.  '
                        'The passed backup is *not* deleted.'))
    delete_before_parser.add_argument(
        'BEFORE_SEGMENT_EXCLUSIVE',
        help='A WAL segment number or base backup name')

    # delete old versions operator
    delete_old_versions_parser = delete_subparsers.add_parser('old-versions')

    # delete *everything* operator
    delete_old_versions_parser = delete_subparsers.add_parser('everything')

    # Okay, parse some arguments, finally
    args = parser.parse_args()

    # Attempt to read a few key parameters from environment variables
    # *or* the command line, enforcing a precedence order and
    # complaining should the required parameter not be defined in
    # either location.
    secret_key = os.getenv('AWS_SECRET_ACCESS_KEY')
    if secret_key is None:
        logger.error(
            msg='no AWS_SECRET_ACCESS_KEY defined',
            hint='Define the environment variable AWS_SECRET_ACCESS_KEY.')
        sys.exit(1)

    s3_prefix = args.s3_prefix or os.getenv('WALE_S3_PREFIX')

    if s3_prefix is None:
        logger.error(
            msg='no storage prefix defined',
            hint=('Either set the --s3-prefix option or define the '
                  'environment variable WALE_S3_PREFIX.'))
        sys.exit(1)

    if args.aws_access_key_id is None:
        aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
        if aws_access_key_id is None:
            logger.error(
                msg='no storage prefix defined',
                hint=('Either set the --aws-access-key-id option or define '
                      'the environment variable AWS_ACCESS_KEY_ID.'))
            sys.exit(1)
    else:
        aws_access_key_id = args.aws_access_key_id

    backup_cxt = s3_operator.S3Backup(aws_access_key_id, secret_key, s3_prefix)

    subcommand = args.subcommand

    try:
        if subcommand == 'backup-fetch':
            external_program_check([LZOP_BIN])
            backup_cxt.database_s3_fetch(
                args.PG_CLUSTER_DIRECTORY,
                args.BACKUP_NAME,
                pool_size=args.pool_size)
        elif subcommand == 'backup-list':
            backup_cxt.backup_list(query=args.QUERY, detail=args.detail)
        elif subcommand == 'backup-push':
            external_program_check([LZOP_BIN, PSQL_BIN, MBUFFER_BIN])
            rate_limit = args.rate_limit
            if rate_limit is not None and rate_limit < 8192:
                logger.error(
                    msg='bad rate limit passed',
                    detail='The passed rate limit was {0}'.format(rate_limit),
                    hint=('Pass a rate limit that is positive and equal or '
                          'greater than 8192'))
                sys.exit(1)

            backup_cxt.database_s3_backup(
                args.PG_CLUSTER_DIRECTORY, rate_limit=rate_limit,
                pool_size=args.pool_size)
        elif subcommand == 'wal-fetch':
            external_program_check([LZOP_BIN])
            backup_cxt.wal_s3_restore(args.WAL_SEGMENT, args.WAL_DESTINATION)
        elif subcommand == 'wal-push':
            external_program_check([LZOP_BIN])
            backup_cxt.wal_s3_archive(args.WAL_SEGMENT)
        elif subcommand == 'delete':
            # Set up pruning precedence, optimizing for *not* deleting data
            #
            # Canonicalize the passed arguments into the value
            # "is_dry_run_really"
            if args.dry_run is False and args.force is True:
                # Actually delete data *only* if there are *no* --dry-runs
                # present and --force is present.
                logger.info(msg='deleting data in S3')
                is_dry_run_really = False
            else:
                logger.info(msg='performing dry run of S3 data deletion')
                is_dry_run_really = True

                import boto.s3.key

                # This is not necessary, but "just in case" to find bugs.
                def just_error(*args, **kwargs):
                    assert False, ('About to delete something in '
                                   'dry-run mode.  Please report a bug.')

                boto.s3.key.Key.delete = just_error

            # Handle the subcommands and route them to the right
            # implementations.
            if args.delete_subcommand == 'old-versions':
                backup_cxt.delete_old_versions(is_dry_run_really)
            elif args.delete_subcommand == 'everything':
                backup_cxt.delete_all(is_dry_run_really)
            elif args.delete_subcommand == 'before':
                segment_info = extract_segment(args.BEFORE_SEGMENT_EXCLUSIVE)
                backup_cxt.delete_before(is_dry_run_really, segment_info)
            else:
                assert False, 'Should be rejected by argument parsing.'
        else:
            logger.error(msg='subcommand not implemented',
                         detail=('The submitted subcommand was {0}.'
                                 .format(subcommand)),
                         hint='Check for typos or consult wal-e --help.')
            sys.exit(127)

        # Report on all encountered exceptions, and raise the last one
        # to take advantage of the final catch-all reporting and exit
        # code management.
        if backup_cxt.exceptions:
            for exc in backup_cxt.exceptions[:-1]:
                logger.log(level=exc.severity,
                           msg=log_help.WalELogger
                           .fmt_logline(exc.msg, exc.detail, exc.hint))
            raise backup_cxt.exceptions[-1]

    except UserException, e:
        logger.log(level=e.severity,
                   msg=log_help.WalELogger
                   .fmt_logline(e.msg, e.detail, e.hint))
        sys.exit(1)

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


<commit_before>"""Plot SPC Outlook Infographics."""
import datetime

# third party
import pytz
import pandas as pd
from geopandas import read_postgis
import cartopy.crs as ccrs
from matplotlib.patches import Rectangle
from pyiem.plot import MapPlot
from pyiem.util import get_autoplot_context, get_dbconn, utc
from pyiem.exceptions import NoDataFound
from pyiem.reference import Z_OVERLAY2_LABEL, Z_POLITICAL

CENTRALTZ = pytz.timezone("America/Chicago")
PDICT = {
    "cwa": "Plot by NWS Forecast Office",
    "state": "Plot by State/Sector",
}
PDICT2 = {
    "1C": "Day 1 Convective",
    "2C": "Day 2 Convective",
    "3C": "Day 3 Convective",
    "4C": "Day 4 Convective",
    "5C": "Day 5 Convective",
    "6C": "Day 6 Convective",
    "7C": "Day 7 Convective",
    "8C": "Day 8 Convective",
    "0C": "Day 4-8 Convective",
    "1F": "Day 1 Fire Weather",
    "2F": "Day 2 Fire Weather",
    "0F": "Day 3-8 Fire Weather",
}
PDICT3 = {
    "categorical": "Categorical (D1-3), Any Severe (D4-8)",
    "hail": "Hail (D1-2)",
    "tornado": "Tornado (D1-2)",
    "wind": "Wind (D1-2)",
}
THRESHOLD_LEVELS = {
    "TSTM": "Thunderstorms",
    "MRGL": "1. Marginal",
    "SLGT": "2. Slight",
    "ENH": "3. Enhanced",
    "MDT": "4. Moderate",
    "HIGH": "5. High",
}
COLORS = {
    "TSTM": "#c0e8c0",
    "MRGL": "#66c57d",
    "SLGT": "#f6f67b",
    "ENH": "#edbf7c",
    "MDT": "#f67a7d",
    "HIGH": "#ff78ff",
    "ELEV": "#ffbb7c",
    "CRIT": "#ff787d",
    "EXTM": "#ff78ff",
    "IDRT": "#c5a392",
    "SDRT": "#ff7f7f",
    "0.02": "#008b00",
    "0.05": "#8b4726",
    "0.10": "#ffc800",
    "0.15": "#ff0000",
    "0.30": "#ff00ff",
    "0.45": "#912cee",
    "0.60": "#104e8b",
}
DAY_COLORS = {
    3: "#ff00ff",
    4: "#ff0000",
    5: "#ffff00",
    6: "#edbf7c",
    7: "#f67a7d",
    8: "#ff78ff",
}
OUTLINE_COLORS = {
    "TSTM": "#566453",
    "MRGL": "#437a43",
    "SLGT": "#d8a31f",
    "ENH": "#d9921c",
    "MDT": "#a6160b",
    "HIGH": "#d21fc3",
}
ISO = "%Y-%m-%d %H:%M"

def get_description():
    """Return a dict describing how to call this plotter"""
    desc = dict()
    desc["data"] = True
    desc["cache"] = 600
    desc[
        "description"
    ] = """
    This app generates infographics for the Storm Prediction Center
    outlooks.  The trick here is how the valid time works.  The app will first
    attempt to match an issuance to that timestamp, if it fails, it then
    looks backwards in time for the most recent issuance to that timestamp.

Another bit of ambiguity is which outlook you get between about
    midnight and the 13z issuance of the Day 1 Outlook.  In this case and
    as the code stands now, you get the next day's outlook.

A GIS Shapefile download
    option exists for downloading these outlooks in-bulk.

    """
    desc["arguments"] = [
        dict(
            type="select",
            default="1C",
            name="which",
            options=PDICT2,
            label="Select SPC Outlook Day and Type",
        ),
        dict(
            type="select",
            default="categorical",
            name="cat",
            options=PDICT3,
            label="Select SPC Outlook Category",
        ),
        dict(
            type="select",
            name="t",
            default="state",
            options=PDICT,
            label="Select plot extent type:",
        ),
        dict(
            type="networkselect",
            name="wfo",
            network="WFO",
            default="DMX",
            label="Select WFO: (ignored if plotting state)",
        ),
        dict(
            type="csector",
            name="csector",
            default="IA",
            label="Select state/sector to plot",
        ),
        dict(
            type="datetime",
            name="valid",
            default=utc().strftime("%Y/%m/%d %H%M"),
            label="Outlook Issuance/Valid Timestamp (UTC Timezone):",
            min="1987/01/01 0000",
        ),
    ]
    return desc

def outlook_search(pgconn, valid, days, outlook_type):
    """Go look for a nearest in time TAF."""
    cursor = pgconn.cursor()
    cursor.execute(
        "SELECT product_issue at time zone 'UTC' from spc_outlook "
        "WHERE day in %s and outlook_type = %s and product_issue > %s and "
        "product_issue < %s ORDER by product_issue DESC",
        (days, outlook_type, valid - datetime.timedelta(hours=24), valid),
    )
    if cursor.rowcount == 0:
        return None
    return cursor.fetchone()[0].replace(tzinfo=datetime.timezone.utc)

def compute_datelabel(df):
    """Figure out something pretty."""
    date1 = df["issue"].min().tz_convert(CENTRALTZ).to_pydatetime()
    date2 = df["issue"].max().tz_convert(CENTRALTZ).to_pydatetime()
    if date1 == date2:
        return date1.strftime("%B %-d, %Y")
    if date1.month == date2.month:
        return (
            date1.strftime("%B %-d-")
            + date2.strftime("%-d ")
            + str(date2.year)
        )
    return (
        date1.strftime("%b %-d-") + date2.strftime("%b %-d ") + str(date2.year)
    )

def get_threshold_label(threshold):
    """Make it pretty."""
    if threshold in THRESHOLD_LEVELS:
        return THRESHOLD_LEVELS.get(threshold)
    if threshold.startswith("0."):
        return "%.0f%%" % (float(threshold) * 100.0,)
    return threshold

def plotter(fdict):
    """Go"""
    ctx = get_autoplot_context(fdict, get_description())
    valid = ctx["valid"].replace(tzinfo=datetime.timezone.utc)
    pgconn = get_dbconn("postgis")
    day, outlook_type = int(ctx["which"][0]), ctx["which"][1]
    category = ctx["cat"].upper()
    if outlook_type == "F":
        category = "FIRE WEATHER CATEGORICAL"
    if (day == 0 or day > 2) and outlook_type == "F":
        category = "CRITICAL FIRE WEATHER AREA"
    if (day == 0 or day > 3) and outlook_type == "C":
        category = "ANY SEVERE"
    days = (day,)
    if day == 0:
        days = tuple(range(4, 9))
        if outlook_type == "F":
            days = tuple(range(3, 9))

    def fetch(ts):
        """Getme data."""
        return read_postgis(
            "WITH data as ("
            "SELECT * from spc_outlooks WHERE category = %s and "
            "product_issue = %s and day in %s and outlook_type = %s) "
            "SELECT d.*, t.priority from data d LEFT JOIN "
            "spc_outlook_thresholds t on (d.threshold = t.threshold) "
            "ORDER by day ASC, priority ASC",
            pgconn,
            params=(category, ts, days, outlook_type),
            index_col=None,
            geom_col="geom",
        )

    df = fetch(valid)
    if df.empty:
        valid = outlook_search(pgconn, valid, days, outlook_type)
        if valid is None:
            raise NoDataFound("SPC Outlook data was not found!")
        df = fetch(valid)
    if ctx["t"] == "cwa":
        sector = "cwa"
    else:
        sector = "state" if len(ctx["csector"]) == 2 else ctx["csector"]
    daylabel = day if day > 0 else "%s-%s" % (days[0], days[-1])
    datelabel = compute_datelabel(df)

    catlabel = " ".join([x.capitalize() for x in category.split()])
    mp = MapPlot(
        title=(
            f"{datelabel} Storm Prediction Center Day {daylabel} "
            f"{catlabel} Outlook"
        ),
        subtitle=(
            f"Issued: {df.iloc[0]['product_issue'].strftime(ISO)} UTC "
            f"| Valid: {df['issue'].min().strftime(ISO)} UTC "
            f"| Expire: {df['expire'].max().strftime(ISO)} UTC"
        ),
        sector=sector,
        twitter=True,
        state=ctx["csector"],
        cwa=(ctx["wfo"] if len(ctx["wfo"]) == 3 else ctx["wfo"][1:]),
        nocaption=True,
    )
    rects = []
    rectlabels = []
    for _idx, row in df[~pd.isna(df["threshold"])].iterrows():
        if row["threshold"] == "SIGN":
            mp.ax.add_geometries(
                [row["geom"]],
                ccrs.PlateCarree(),
                facecolor="None",
                edgecolor="k",
                linewidth=2,
                hatch="/",
                zorder=Z_POLITICAL - 1,
            )
            rect = Rectangle((0, 0), 1, 1, fc="None", hatch="/")
        else:
            fc = COLORS.get(row["threshold"], "red")
            ec = OUTLINE_COLORS.get(row["threshold"], "k")
            if day == 0:
                fc, ec = "None", DAY_COLORS[row["day"]]
                if row["threshold"] == "0.30":
                    fc, ec = ec, ec
            mp.ax.add_geometries(
                [row["geom"]],
                ccrs.PlateCarree(),
                facecolor=fc,
                edgecolor=ec,
                linewidth=2,
                zorder=Z_POLITICAL - 1,
            )
            rect = Rectangle((0, 0), 1, 1, fc=fc, ec=ec)
        rects.append(rect)
        label = get_threshold_label(row["threshold"])
        if day == 0:
            label = f"D{row['day']} {label}"
        rectlabels.append(label)
    if rects:
        mp.ax.legend(
            rects,
            rectlabels,
            ncol=1,
            loc="lower right",
            fontsize=12,
            bbox_to_anchor=(1.08, 0.01),
            fancybox=True,
            framealpha=1,
        ).set_zorder(Z_OVERLAY2_LABEL + 100)
    else:
        emptytext = "Potential Too Low"
        if outlook_type == "C" and day in [1, 2, 3]:
            emptytext = "No Thunderstorms Forecast"
        mp.ax.text(
            0.5,
            0.5,
            emptytext,
            transform=mp.ax.transAxes,
            fontsize=30,
            ha="center",
            color="yellow",
            bbox=dict(color="black", alpha=0.5),
            zorder=Z_OVERLAY2_LABEL,
        )

    if sector == "cwa":
        mp.draw_cwas(color="k", linewidth=2.5)
        mp.drawcounties()
        mp.drawcities()

    for col in ["product_issue", "issue", "expire", "updated"]:
        # some rows could be NaN
        df[col] = df[~pd.isna(df[col])][col].apply(lambda x: x.strftime(ISO))
    return mp.fig, df.drop("geom", axis=1)

if __name__ == "__main__":
    # has all 5 days with something included
    # plotter(dict(cat="categorical", which="0C", valid="2019-05-14 2022"))
    # has three days of F
    # plotter(dict(cat="categorical", which="0F", valid="2018-05-07 2322"))
    plotter(dict(cat="tornado", which="1C", valid="2011-04-07 2058"))
<commit_msg>revert use of view<commit_after>"""Plot SPC Outlook Infographics."""
import datetime

# third party
import pytz
import pandas as pd
from geopandas import read_postgis
import cartopy.crs as ccrs
from matplotlib.patches import Rectangle
from pyiem.plot import MapPlot
from pyiem.util import get_autoplot_context, get_dbconn, utc
from pyiem.exceptions import NoDataFound
from pyiem.reference import Z_OVERLAY2_LABEL, Z_POLITICAL

CENTRALTZ = pytz.timezone("America/Chicago")
PDICT = {
    "cwa": "Plot by NWS Forecast Office",
    "state": "Plot by State/Sector",
}
PDICT2 = {
    "1C": "Day 1 Convective",
    "2C": "Day 2 Convective",
    "3C": "Day 3 Convective",
    "4C": "Day 4 Convective",
    "5C": "Day 5 Convective",
    "6C": "Day 6 Convective",
    "7C": "Day 7 Convective",
    "8C": "Day 8 Convective",
    "0C": "Day 4-8 Convective",
    "1F": "Day 1 Fire Weather",
    "2F": "Day 2 Fire Weather",
    "0F": "Day 3-8 Fire Weather",
}
PDICT3 = {
    "categorical": "Categorical (D1-3), Any Severe (D4-8)",
    "hail": "Hail (D1-2)",
    "tornado": "Tornado (D1-2)",
    "wind": "Wind (D1-2)",
}
THRESHOLD_LEVELS = {
    "TSTM": "Thunderstorms",
    "MRGL": "1. Marginal",
    "SLGT": "2. Slight",
    "ENH": "3. Enhanced",
    "MDT": "4. Moderate",
    "HIGH": "5. High",
}
COLORS = {
    "TSTM": "#c0e8c0",
    "MRGL": "#66c57d",
    "SLGT": "#f6f67b",
    "ENH": "#edbf7c",
    "MDT": "#f67a7d",
    "HIGH": "#ff78ff",
    "ELEV": "#ffbb7c",
    "CRIT": "#ff787d",
    "EXTM": "#ff78ff",
    "IDRT": "#c5a392",
    "SDRT": "#ff7f7f",
    "0.02": "#008b00",
    "0.05": "#8b4726",
    "0.10": "#ffc800",
    "0.15": "#ff0000",
    "0.30": "#ff00ff",
    "0.45": "#912cee",
    "0.60": "#104e8b",
}
DAY_COLORS = {
    3: "#ff00ff",
    4: "#ff0000",
    5: "#ffff00",
    6: "#edbf7c",
    7: "#f67a7d",
    8: "#ff78ff",
}
OUTLINE_COLORS = {
    "TSTM": "#566453",
    "MRGL": "#437a43",
    "SLGT": "#d8a31f",
    "ENH": "#d9921c",
    "MDT": "#a6160b",
    "HIGH": "#d21fc3",
}
ISO = "%Y-%m-%d %H:%M"

def get_description():
    """Return a dict describing how to call this plotter"""
    desc = dict()
    desc["data"] = True
    desc["cache"] = 600
    desc[
        "description"
    ] = """
    This app generates infographics for the Storm Prediction Center
    outlooks.  The trick here is how the valid time works.  The app will first
    attempt to match an issuance to that timestamp, if it fails, it then
    looks backwards in time for the most recent issuance to that timestamp.

Another bit of ambiguity is which outlook you get between about
    midnight and the 13z issuance of the Day 1 Outlook.  In this case and
    as the code stands now, you get the next day's outlook.

A GIS Shapefile download
    option exists for downloading these outlooks in-bulk.

    """
    desc["arguments"] = [
        dict(
            type="select",
            default="1C",
            name="which",
            options=PDICT2,
            label="Select SPC Outlook Day and Type",
        ),
        dict(
            type="select",
            default="categorical",
            name="cat",
            options=PDICT3,
            label="Select SPC Outlook Category",
        ),
        dict(
            type="select",
            name="t",
            default="state",
            options=PDICT,
            label="Select plot extent type:",
        ),
        dict(
            type="networkselect",
            name="wfo",
            network="WFO",
            default="DMX",
            label="Select WFO: (ignored if plotting state)",
        ),
        dict(
            type="csector",
            name="csector",
            default="IA",
            label="Select state/sector to plot",
        ),
        dict(
            type="datetime",
            name="valid",
            default=utc().strftime("%Y/%m/%d %H%M"),
            label="Outlook Issuance/Valid Timestamp (UTC Timezone):",
            min="1987/01/01 0000",
        ),
    ]
    return desc

def outlook_search(pgconn, valid, days, outlook_type):
    """Go look for a nearest in time TAF."""
    cursor = pgconn.cursor()
    cursor.execute(
        "SELECT product_issue at time zone 'UTC' from spc_outlook "
        "WHERE day in %s and outlook_type = %s and product_issue > %s and "
        "product_issue < %s ORDER by product_issue DESC",
        (days, outlook_type, valid - datetime.timedelta(hours=24), valid),
    )
    if cursor.rowcount == 0:
        return None
    return cursor.fetchone()[0].replace(tzinfo=datetime.timezone.utc)

def compute_datelabel(df):
    """Figure out something pretty."""
    date1 = df["issue"].min().tz_convert(CENTRALTZ).to_pydatetime()
    date2 = df["issue"].max().tz_convert(CENTRALTZ).to_pydatetime()
    if date1 == date2:
        return date1.strftime("%B %-d, %Y")
    if date1.month == date2.month:
        return (
            date1.strftime("%B %-d-")
            + date2.strftime("%-d ")
            + str(date2.year)
        )
    return (
        date1.strftime("%b %-d-") + date2.strftime("%b %-d ") + str(date2.year)
    )

def get_threshold_label(threshold):
    """Make it pretty."""
    if threshold in THRESHOLD_LEVELS:
        return THRESHOLD_LEVELS.get(threshold)
    if threshold.startswith("0."):
        return "%.0f%%" % (float(threshold) * 100.0,)
    return threshold

def plotter(fdict):
    """Go"""
    ctx = get_autoplot_context(fdict, get_description())
    valid = ctx["valid"].replace(tzinfo=datetime.timezone.utc)
    pgconn = get_dbconn("postgis")
    day, outlook_type = int(ctx["which"][0]), ctx["which"][1]
    category = ctx["cat"].upper()
    if outlook_type == "F":
        category = "FIRE WEATHER CATEGORICAL"
    if (day == 0 or day > 2) and outlook_type == "F":
        category = "CRITICAL FIRE WEATHER AREA"
    if (day == 0 or day > 3) and outlook_type == "C":
        category = "ANY SEVERE"
    days = (day,)
    if day == 0:
        days = tuple(range(4, 9))
        if outlook_type == "F":
            days = tuple(range(3, 9))

    def fetch(ts):
        """Getme data."""
        # NB careful here with the joins and not to use the view!
        return read_postgis(
            "WITH data as ("
            "SELECT o.*, g.* from spc_outlook o LEFT JOIN "
            "spc_outlook_geometries g on (o.id = g.spc_outlook_id and "
            "g.category = %s) WHERE product_issue = %s and day in %s and "
            "outlook_type = %s) "
            "SELECT d.*, t.priority from data d LEFT JOIN "
            "spc_outlook_thresholds t on (d.threshold = t.threshold) "
            "ORDER by day ASC, priority ASC",
            pgconn,
            params=(category, ts, days, outlook_type),
            index_col=None,
            geom_col="geom",
        )

    df = fetch(valid)
    if df.empty:
        valid = outlook_search(pgconn, valid, days, outlook_type)
        if valid is None:
            raise NoDataFound("SPC Outlook data was not found!")
        df = fetch(valid)
    if ctx["t"] == "cwa":
        sector = "cwa"
    else:
        sector = "state" if len(ctx["csector"]) == 2 else ctx["csector"]
    daylabel = day if day > 0 else "%s-%s" % (days[0], days[-1])
    datelabel = compute_datelabel(df)

    catlabel = " ".join([x.capitalize() for x in category.split()])
    mp = MapPlot(
        title=(
            f"{datelabel} Storm Prediction Center Day {daylabel} "
            f"{catlabel} Outlook"
        ),
        subtitle=(
            f"Issued: {df.iloc[0]['product_issue'].strftime(ISO)} UTC "
            f"| Valid: {df['issue'].min().strftime(ISO)} UTC "
            f"| Expire: {df['expire'].max().strftime(ISO)} UTC"
        ),
        sector=sector,
        twitter=True,
        state=ctx["csector"],
        cwa=(ctx["wfo"] if len(ctx["wfo"]) == 3 else ctx["wfo"][1:]),
        nocaption=True,
    )
    rects = []
    rectlabels = []
    for _idx, row in df[~pd.isna(df["threshold"])].iterrows():
        if row["threshold"] == "SIGN":
            mp.ax.add_geometries(
                [row["geom"]],
                ccrs.PlateCarree(),
                facecolor="None",
                edgecolor="k",
                linewidth=2,
                hatch="/",
                zorder=Z_POLITICAL - 1,
            )
            rect = Rectangle((0, 0), 1, 1, fc="None", hatch="/")
        else:
            fc = COLORS.get(row["threshold"], "red")
            ec = OUTLINE_COLORS.get(row["threshold"], "k")
            if day == 0:
                fc, ec = "None", DAY_COLORS[row["day"]]
                if row["threshold"] == "0.30":
                    fc, ec = ec, ec
            mp.ax.add_geometries(
                [row["geom"]],
                ccrs.PlateCarree(),
                facecolor=fc,
                edgecolor=ec,
                linewidth=2,
                zorder=Z_POLITICAL - 1,
            )
            rect = Rectangle((0, 0), 1, 1, fc=fc, ec=ec)
        rects.append(rect)
        label = get_threshold_label(row["threshold"])
        if day == 0:
            label = f"D{row['day']} {label}"
        rectlabels.append(label)
    if rects:
        mp.ax.legend(
            rects,
            rectlabels,
            ncol=1,
            loc="lower right",
            fontsize=12,
            bbox_to_anchor=(1.08, 0.01),
            fancybox=True,
            framealpha=1,
        ).set_zorder(Z_OVERLAY2_LABEL + 100)
    else:
        emptytext = "Potential Too Low"
        if outlook_type == "C" and day in [1, 2, 3]:
            emptytext = "No Thunderstorms Forecast"
        mp.ax.text(
            0.5,
            0.5,
            emptytext,
            transform=mp.ax.transAxes,
            fontsize=30,
            ha="center",
            color="yellow",
            bbox=dict(color="black", alpha=0.5),
            zorder=Z_OVERLAY2_LABEL,
        )

    if sector == "cwa":
        mp.draw_cwas(color="k", linewidth=2.5)
        mp.drawcounties()
        mp.drawcities()

    for col in ["product_issue", "issue", "expire", "updated"]:
        # some rows could be NaN
        df[col] = df[~pd.isna(df[col])][col].apply(lambda x: x.strftime(ISO))
    return mp.fig, df.drop("geom", axis=1)

if __name__ == "__main__":
    # has all 5 days with something included
    # plotter(dict(cat="categorical", which="0C", valid="2019-05-14 2022"))
    # has three days of F
    # plotter(dict(cat="categorical", which="0F", valid="2018-05-07 2322"))
    plotter(dict(cat="categorical", which="0C", valid="2021-07-09 0854"))
<|endoftext|>