<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import gzip
import logging
import shapely
import numpy
from openquake.baselib import sap
from openquake.hazardlib.geo.utils import get_bounding_box
from openquake.calculators.extract import Extractor, WebExtractor

def basemap(projection, sitecol):
    from mpl_toolkits.basemap import Basemap  # costly import
    minlon, minlat, maxlon, maxlat = get_bounding_box(sitecol, maxdist=10)
    bmap = Basemap(projection=projection,
                   llcrnrlon=minlon, llcrnrlat=minlat,
                   urcrnrlon=maxlon, urcrnrlat=maxlat,
                   lat_0=sitecol['lat'].mean(), lon_0=sitecol['lon'].mean())
    bmap.drawcoastlines()
    return bmap

def make_figure_hcurves(extractors, what):
    """
    $ oq plot 'hcurves?kind=mean&imt=PGA&site_id=0'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    got = {}  # (calc_id, kind) -> curves
    for i, ex in enumerate(extractors):
        hcurves = ex.get(what)
        for kind in hcurves.kind:
            got[ex.calc_id, kind] = hcurves[kind]
    oq = ex.oqparam
    n_imts = len(hcurves.imt)
    [site] = hcurves.site_id
    for j, imt in enumerate(hcurves.imt):
        imls = oq.imtls[imt]
        imt_slice = oq.imtls(imt)
        ax = fig.add_subplot(n_imts, 1, j + 1)
        ax.set_xlabel('%s, site %s, inv_time=%dy' %
                      (imt, site, oq.investigation_time))
        ax.set_ylabel('PoE')
        for ck, arr in got.items():
            if (arr == 0).all():
                logging.warning('There is a zero curve %s_%s', *ck)
            ax.loglog(imls, arr[0, imt_slice], '-', label='%s_%s' % ck)
            ax.loglog(imls, arr[0, imt_slice], '.')
        ax.grid(True)
        ax.legend()
    return plt

def make_figure_hmaps(extractors, what):
    """
    $ oq plot 'hmaps?kind=mean&imt=PGA'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ncalcs = len(extractors)
    if ncalcs > 2:
        raise RuntimeError('Could not plot more than two calculations at once')
    elif ncalcs == 2:  # plot the differences
        ex1, ex2 = extractors
        oq1 = ex1.oqparam
        oq2 = ex2.oqparam
        n_poes = len(oq1.poes)
        assert n_poes == len(oq2.poes)
        itime = oq1.investigation_time
        assert oq2.investigation_time == itime
        sitecol = ex1.get('sitecol')
        assert ex1.get('sitecol') == sitecol
        hmaps1 = ex1.get(what)
        hmaps2 = ex2.get(what)
        [imt] = hmaps1.imt
        assert [imt] == hmaps2.imt
        [kind] = hmaps1.kind
        assert hmaps1.kind == [kind]
        for j, poe in enumerate(oq1.poes):
            diff = hmaps1[kind][:, 0, j] - hmaps2[kind][:, 0, j]
            maxdiff = numpy.abs(diff).max()
            ax = fig.add_subplot(1, n_poes, j + 1)
            ax.grid(True)
            ax.set_xlabel('IMT=%s, kind=%s, poe=%s\ncalcs %d-%d, '
                          'inv_time=%dy\nmaxdiff=%s' %
                          (imt, kind, poe, ex1.calc_id, ex2.calc_id,
                           itime, maxdiff))
            bmap = basemap('cyl', sitecol)
            bmap.scatter(sitecol['lon'], sitecol['lat'],
                         c=diff, cmap='jet')
    elif ncalcs == 1:  # plot the hmap
        [ex] = extractors
        oq = ex.oqparam
        n_poes = len(oq.poes)
        sitecol = ex.get('sitecol')
        hmaps = ex.get(what)
        [imt] = hmaps.imt
        [kind] = hmaps.kind
        for j, poe in enumerate(oq.poes):
            ax = fig.add_subplot(1, n_poes, j + 1)
            ax.grid(True)
            ax.set_xlabel('hmap for IMT=%s, kind=%s, poe=%s\ncalculation %d, '
                          'inv_time=%dy' %
                          (imt, kind, poe, ex.calc_id, oq.investigation_time))
            bmap = basemap('cyl', sitecol)
            bmap.scatter(sitecol['lon'], sitecol['lat'],
                         c=hmaps[kind][:, 0, j], cmap='jet')
    return plt

def make_figure_uhs(extractors, what):
    """
    $ oq plot 'uhs?kind=mean&site_id=0'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    got = {}  # (calc_id, kind) -> curves
    for i, ex in enumerate(extractors):
        uhs = ex.get(what)
        for kind in uhs.kind:
            got[ex.calc_id, kind] = uhs[kind]
    oq = ex.oqparam
    n_poes = len(oq.poes)
    periods = [imt.period for imt in oq.imt_periods()]
    [site] = uhs.site_id
    for j, poe in enumerate(oq.poes):
        ax = fig.add_subplot(n_poes, 1, j + 1)
        ax.set_xlabel('UHS on site %s, poe=%s, inv_time=%dy' %
                      (site, poe, oq.investigation_time))
        ax.set_ylabel('SA')
        for ck, arr in got.items():
            ax.plot(periods, arr[0, :, j], '-', label='%s_%s' % ck)
            ax.plot(periods, arr[0, :, j], '.')
        ax.grid(True)
        ax.legend()
    return plt

def make_figure_disagg(extractors, what):
    """
    $ oq plot 'disagg?kind=Mag&imt=PGA'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    oq = extractors[0].oqparam
    disagg = extractors[0].get(what)
    [sid] = disagg.site_id
    [imt] = disagg.imt
    [poe_id] = disagg.poe_id
    ax.set_xlabel('Disagg%s on site %s, imt=%s, poe_id=%d, inv_time=%dy' %
                  (disagg.kind, sid, imt, poe_id, oq.investigation_time))
    for name, values in zip(disagg.names, disagg.array):
        x, y = values.T
        print(y)
        ax.plot(x, y, label=name.split('-')[1])
    ax.legend()
    return plt

def make_figure_task_info(extractors, what):
    """
    $ oq plot 'task_info?kind=classical'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    [ex] = extractors
    [(task_name, task_info)] = ex.get(what).to_dict().items()
    x = task_info['duration']
    ax = fig.add_subplot(1, 1, 1)
    mean, std = x.mean(), x.std(ddof=1)
    ax.hist(x, bins=50, rwidth=0.9)
    ax.set_xlabel("mean=%d+-%d seconds" % (mean, std))
    ax.set_ylabel("tasks=%d" % len(x))
    #from scipy.stats import linregress
    #ax = fig.add_subplot(2, 1, 2)
    #arr = numpy.sort(task_info, order='duration')
    #x, y = arr['duration'], arr['weight']
    #reg = linregress(x, y)
    #ax.plot(x, reg.intercept + reg.slope * x)
    #ax.plot(x, y)
    #ax.set_ylabel("weight")
    #ax.set_xlabel("duration")
    return plt

def make_figure_memory(extractors, what):
    """
    $ oq plot 'memory?'
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt

    [ex] = extractors
    task_info = ex.get('task_info').to_dict()
    fig, ax = plt.subplots()
    ax.grid(True)
    ax.set_xlabel('tasks')
    ax.set_ylabel('GB')
    start = 0
    for task_name in task_info:
        mem = task_info[task_name]['mem_gb']
        ax.plot(range(start, start + len(mem)), mem, label=task_name)
        start += len(mem)
    ax.legend()
    return plt

class PolygonPlotter():
    """
    Add polygons to a given axis object
    """
    def __init__(self, ax):
        self.ax = ax
        self.minxs = []
        self.maxxs = []
        self.minys = []
        self.maxys = []

    def add(self, poly, **kw):
        from openquake.hmtk.plotting.patch import PolygonPatch
        minx, miny, maxx, maxy = poly.bounds
        self.minxs.append(minx)
        self.maxxs.append(maxx)
        self.minys.append(miny)
        self.maxys.append(maxy)
        self.ax.add_patch(PolygonPatch(poly, **kw))

    def set_lim(self):
        if self.minxs and self.maxxs:
            self.ax.set_xlim(min(self.minxs), max(self.maxxs))
        if self.minys and self.maxys:
            self.ax.set_ylim(min(self.minys), max(self.maxys))

@sap.script
def make_figure_sources(extractors, what):
    """
    $ oq plot sources?sm_id=0
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt
    [ex] = extractors
    info = ex.get(what)
    fig, ax = plt.subplots()
    ax.grid(True)
    # sitecol = ex.get('sitecol')
    # bmap = basemap('cyl', sitecol)
    # bmap.plot(sitecol['lon'], sitecol['lat'], '+')
    pp = PolygonPlotter(ax)
    n = 0
    tot = 0
    for rec in info:
        if not rec['wkt'].startswith('POINT'):
            if rec['eff_ruptures']:  # not filtered out
                alpha = .3
                n += 1
            else:
                alpha = .1
            pp.add(shapely.wkt.loads(rec['wkt']), alpha=alpha)
            tot += 1
    pp.set_lim()
    ax.set_title('%d/%d sources for source model #%d' % (n, tot, info.sm_id))
    return plt

@sap.script
def make_figure_rupture_info(extractors, what):
    """
    $ oq plot rupture_info?min_mag=6
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt
    [ex] = extractors
    info = ex.get(what)
    fig, ax = plt.subplots()
    ax.grid(True)
    # sitecol = ex.get('sitecol')
    # bmap = basemap('cyl', sitecol)
    # bmap.plot(sitecol['lon'], sitecol['lat'], '+')
    n = 0
    tot = 0
    pp = PolygonPlotter(ax)
    geoms = gzip.decompress(info['boundaries']).decode('utf8').split('\n')
    for rec, wkt in zip(info, geoms):
        poly = shapely.wkt.loads(wkt)
        if poly.is_valid:
            pp.add(poly)
            n += 1
        else:
            print('Invalid %s' % rec['boundary'].decode('utf8'))
        tot += 1
    pp.set_lim()
    ax.set_title('%d/%d valid ruptures' % (n, tot))
    if tot == 1:
        # print the full geometry
        print(ex.get('rupture/%d' % rec['rupid']).toml())
    return plt

@sap.script
def plot(what='examples', calc_id=-1, other_id=None, webapi=False):
    """
    Generic plotter for local and remote calculations.
    """
    if what == 'examples':
        help_msg = ['Examples of possible plots:']
        for k, v in globals().items():
            if k.startswith('make_figure_'):
                help_msg.append(v.__doc__)
        raise SystemExit(''.join(help_msg))
    if '?' not in what:
        raise SystemExit('Missing ? in %r' % what)
    prefix, rest = what.split('?', 1)
    if prefix in 'hcurves hmaps' and 'imt=' not in rest:
        raise SystemExit('Missing imt= in %r' % what)
    elif prefix == 'uhs' and 'imt=' in rest:
        raise SystemExit('Invalid IMT in %r' % what)
    elif prefix in 'hcurves uhs disagg' and 'site_id=' not in rest:
        what += '&site_id=0'
    if prefix == 'disagg' and 'poe=' not in rest:
        what += '&poe_id=0'
    if webapi:
        xs = [WebExtractor(calc_id)]
        if other_id:
            xs.append(WebExtractor(other_id))
    else:
        xs = [Extractor(calc_id)]
        if other_id:
            xs.append(Extractor(other_id))
    make_figure = globals()['make_figure_' + prefix]
    plt = make_figure(xs, what)
    plt.show()

plot.arg('what', 'what to extract')
plot.arg('calc_id', 'computation ID', type=int)
plot.arg('other_id', 'ID of another computation', type=int)
plot.flg('webapi', 'if given, pass through the WebAPI')
<commit_msg>[skip hazardlib][demos]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2019 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import gzip
import logging
import shapely
import numpy
from openquake.baselib import sap
from openquake.hazardlib.geo.utils import get_bounding_box
from openquake.calculators.extract import Extractor, WebExtractor

def basemap(projection, sitecol):
    from mpl_toolkits.basemap import Basemap  # costly import
    minlon, minlat, maxlon, maxlat = get_bounding_box(sitecol, maxdist=10)
    bmap = Basemap(projection=projection,
                   llcrnrlon=minlon, llcrnrlat=minlat,
                   urcrnrlon=maxlon, urcrnrlat=maxlat,
                   lat_0=sitecol['lat'].mean(), lon_0=sitecol['lon'].mean())
    bmap.drawcoastlines()
    return bmap

def make_figure_hcurves(extractors, what):
    """
    $ oq plot 'hcurves?kind=mean&imt=PGA&site_id=0'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    got = {}  # (calc_id, kind) -> curves
    for i, ex in enumerate(extractors):
        hcurves = ex.get(what)
        for kind in hcurves.kind:
            got[ex.calc_id, kind] = hcurves[kind]
    oq = ex.oqparam
    n_imts = len(hcurves.imt)
    [site] = hcurves.site_id
    for j, imt in enumerate(hcurves.imt):
        imls = oq.imtls[imt]
        imt_slice = oq.imtls(imt)
        ax = fig.add_subplot(n_imts, 1, j + 1)
        ax.set_xlabel('%s, site %s, inv_time=%dy' %
                      (imt, site, oq.investigation_time))
        ax.set_ylabel('PoE')
        for ck, arr in got.items():
            if (arr == 0).all():
                logging.warning('There is a zero curve %s_%s', *ck)
            ax.loglog(imls, arr[0, imt_slice], '-', label='%s_%s' % ck)
            ax.loglog(imls, arr[0, imt_slice], '.')
        ax.grid(True)
        ax.legend()
    return plt

def make_figure_hmaps(extractors, what):
    """
    $ oq plot 'hmaps?kind=mean&imt=PGA'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ncalcs = len(extractors)
    if ncalcs > 2:
        raise RuntimeError('Could not plot more than two calculations at once')
    elif ncalcs == 2:  # plot the differences
        ex1, ex2 = extractors
        oq1 = ex1.oqparam
        oq2 = ex2.oqparam
        n_poes = len(oq1.poes)
        assert n_poes == len(oq2.poes)
        itime = oq1.investigation_time
        assert oq2.investigation_time == itime
        sitecol = ex1.get('sitecol')
        assert ex1.get('sitecol') == sitecol
        hmaps1 = ex1.get(what)
        hmaps2 = ex2.get(what)
        [imt] = hmaps1.imt
        assert [imt] == hmaps2.imt
        [kind] = hmaps1.kind
        assert hmaps1.kind == [kind]
        for j, poe in enumerate(oq1.poes):
            diff = hmaps1[kind][:, 0, j] - hmaps2[kind][:, 0, j]
            maxdiff = numpy.abs(diff).max()
            ax = fig.add_subplot(1, n_poes, j + 1)
            ax.grid(True)
            ax.set_xlabel('IMT=%s, kind=%s, poe=%s\ncalcs %d-%d, '
                          'inv_time=%dy\nmaxdiff=%s' %
                          (imt, kind, poe, ex1.calc_id, ex2.calc_id,
                           itime, maxdiff))
            bmap = basemap('cyl', sitecol)
            bmap.scatter(sitecol['lon'], sitecol['lat'],
                         c=diff, cmap='jet')
    elif ncalcs == 1:  # plot the hmap
        [ex] = extractors
        oq = ex.oqparam
        n_poes = len(oq.poes)
        sitecol = ex.get('sitecol')
        hmaps = ex.get(what)
        [imt] = hmaps.imt
        [kind] = hmaps.kind
        for j, poe in enumerate(oq.poes):
            ax = fig.add_subplot(1, n_poes, j + 1)
            ax.grid(True)
            ax.set_xlabel('hmap for IMT=%s, kind=%s, poe=%s\ncalculation %d, '
                          'inv_time=%dy' %
                          (imt, kind, poe, ex.calc_id, oq.investigation_time))
            bmap = basemap('cyl', sitecol)
            bmap.scatter(sitecol['lon'], sitecol['lat'],
                         c=hmaps[kind][:, 0, j], cmap='jet')
    return plt

def make_figure_uhs(extractors, what):
    """
    $ oq plot 'uhs?kind=mean&site_id=0'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    got = {}  # (calc_id, kind) -> curves
    for i, ex in enumerate(extractors):
        uhs = ex.get(what)
        for kind in uhs.kind:
            got[ex.calc_id, kind] = uhs[kind]
    oq = ex.oqparam
    n_poes = len(oq.poes)
    periods = [imt.period for imt in oq.imt_periods()]
    [site] = uhs.site_id
    for j, poe in enumerate(oq.poes):
        ax = fig.add_subplot(n_poes, 1, j + 1)
        ax.set_xlabel('UHS on site %s, poe=%s, inv_time=%dy' %
                      (site, poe, oq.investigation_time))
        ax.set_ylabel('SA')
        for ck, arr in got.items():
            ax.plot(periods, arr[0, :, j], '-', label='%s_%s' % ck)
            ax.plot(periods, arr[0, :, j], '.')
        ax.grid(True)
        ax.legend()
    return plt

def make_figure_disagg(extractors, what):
    """
    $ oq plot 'disagg?kind=Mag&imt=PGA'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    oq = extractors[0].oqparam
    disagg = extractors[0].get(what)
    [sid] = disagg.site_id
    [imt] = disagg.imt
    [poe_id] = disagg.poe_id
    ax.set_xlabel('Disagg%s on site %s, imt=%s, poe_id=%d, inv_time=%dy' %
                  (disagg.kind, sid, imt, poe_id, oq.investigation_time))
    for name, values in zip(disagg.names, disagg.array):
        x, y = values.T
        print(y)
        ax.plot(x, y, label=name.split('-')[1])
    ax.legend()
    return plt

def make_figure_task_info(extractors, what):
    """
    $ oq plot 'task_info?kind=classical'
    """
    import matplotlib.pyplot as plt
    fig = plt.figure()
    [ex] = extractors
    [(task_name, task_info)] = ex.get(what).to_dict().items()
    x = task_info['duration']
    ax = fig.add_subplot(1, 1, 1)
    mean, std = x.mean(), x.std(ddof=1)
    ax.hist(x, bins=50, rwidth=0.9)
    ax.set_xlabel("mean=%d+-%d seconds" % (mean, std))
    ax.set_ylabel("tasks=%d" % len(x))
    #from scipy.stats import linregress
    #ax = fig.add_subplot(2, 1, 2)
    #arr = numpy.sort(task_info, order='duration')
    #x, y = arr['duration'], arr['weight']
    #reg = linregress(x, y)
    #ax.plot(x, reg.intercept + reg.slope * x)
    #ax.plot(x, y)
    #ax.set_ylabel("weight")
    #ax.set_xlabel("duration")
    return plt

def make_figure_memory(extractors, what):
    """
    $ oq plot 'memory?'
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt

    [ex] = extractors
    task_info = ex.get('task_info').to_dict()
    fig, ax = plt.subplots()
    ax.grid(True)
    ax.set_xlabel('tasks')
    ax.set_ylabel('GB')
    start = 0
    for task_name in task_info:
        mem = task_info[task_name]['mem_gb']
        ax.plot(range(start, start + len(mem)), mem, label=task_name)
        start += len(mem)
    ax.legend()
    return plt

class PolygonPlotter():
    """
    Add polygons to a given axis object
    """
    def __init__(self, ax):
        self.ax = ax
        self.minxs = []
        self.maxxs = []
        self.minys = []
        self.maxys = []

    def add(self, poly, **kw):
        from openquake.hmtk.plotting.patch import PolygonPatch
        minx, miny, maxx, maxy = poly.bounds
        self.minxs.append(minx)
        self.maxxs.append(maxx)
        self.minys.append(miny)
        self.maxys.append(maxy)
        self.ax.add_patch(PolygonPatch(poly, **kw))

    def set_lim(self):
        if self.minxs and self.maxxs:
            self.ax.set_xlim(min(self.minxs), max(self.maxxs))
        if self.minys and self.maxys:
            self.ax.set_ylim(min(self.minys), max(self.maxys))

@sap.script
def make_figure_sources(extractors, what):
    """
    $ oq plot sources?sm_id=0
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt
    [ex] = extractors
    info = ex.get(what)
    fig, ax = plt.subplots()
    ax.grid(True)
    # sitecol = ex.get('sitecol')
    # bmap = basemap('cyl', sitecol)
    # bmap.plot(sitecol['lon'], sitecol['lat'], '+')
    pp = PolygonPlotter(ax)
    n = 0
    tot = 0
    for rec in info:
        if not rec['wkt'].startswith('POINT'):
            if rec['eff_ruptures']:  # not filtered out
                alpha = .3
                n += 1
            else:
                alpha = .1
            pp.add(shapely.wkt.loads(rec['wkt']), alpha=alpha)
            tot += 1
    pp.set_lim()
    ax.set_title('%d/%d sources for source model #%d' % (n, tot, info.sm_id))
    return plt

@sap.script
def make_figure_rupture_info(extractors, what):
    """
    $ oq plot rupture_info?min_mag=6
    """
    # NB: matplotlib is imported inside since it is a costly import
    import matplotlib.pyplot as plt
    [ex] = extractors
    info = ex.get(what)
    fig, ax = plt.subplots()
    ax.grid(True)
    # sitecol = ex.get('sitecol')
    # bmap = basemap('cyl', sitecol)
    # bmap.plot(sitecol['lon'], sitecol['lat'], '+')
    n = 0
    tot = 0
    pp = PolygonPlotter(ax)
    geoms = gzip.decompress(info['boundaries']).decode('utf8').split('\n')
    for rec, wkt in zip(info, geoms):
        poly = shapely.wkt.loads(wkt)
        if poly.is_valid:
            pp.add(poly)
            n += 1
        else:
            print('Invalid %s' % wkt)
        tot += 1
    pp.set_lim()
    ax.set_title('%d/%d valid ruptures' % (n, tot))
    if tot == 1:
        # print the full geometry
        print(ex.get('rupture/%d' % rec['rupid']).toml())
    return plt

@sap.script
def plot(what='examples', calc_id=-1, other_id=None, webapi=False):
    """
    Generic plotter for local and remote calculations.
    """
    if what == 'examples':
        help_msg = ['Examples of possible plots:']
        for k, v in globals().items():
            if k.startswith('make_figure_'):
                help_msg.append(v.__doc__)
        raise SystemExit(''.join(help_msg))
    if '?' not in what:
        raise SystemExit('Missing ? in %r' % what)
    prefix, rest = what.split('?', 1)
    if prefix in 'hcurves hmaps' and 'imt=' not in rest:
        raise SystemExit('Missing imt= in %r' % what)
    elif prefix == 'uhs' and 'imt=' in rest:
        raise SystemExit('Invalid IMT in %r' % what)
    elif prefix in 'hcurves uhs disagg' and 'site_id=' not in rest:
        what += '&site_id=0'
    if prefix == 'disagg' and 'poe=' not in rest:
        what += '&poe_id=0'
    if webapi:
        xs = [WebExtractor(calc_id)]
        if other_id:
            xs.append(WebExtractor(other_id))
    else:
        xs = [Extractor(calc_id)]
        if other_id:
            xs.append(Extractor(other_id))
    make_figure = globals()['make_figure_' + prefix]
    plt = make_figure(xs, what)
    plt.show()

plot.arg('what', 'what to extract')
plot.arg('calc_id', 'computation ID', type=int)
plot.arg('other_id', 'ID of another computation', type=int)
plot.flg('webapi', 'if given, pass through the WebAPI')
<|endoftext|>


<commit_before>"""
    Google Drive (Sheets and Docs) preprocessors allow you to store content in
    Google Drive and bring it into Grow. Grow will authenticate to the Google
    Drive API using OAuth2 and then download content as specified in
    `podspec.yaml`.

    Grow supports various ways to transform the content, e.g. Sheets can be
    downloaded and converted to yaml, and Docs can be downloaded and converted
    to markdown.
"""

from . import base
from googleapiclient import discovery
from googleapiclient import errors
from grow.common import oauth
from grow.common import utils
from grow.pods import formats
from protorpc import messages
import bs4
import cStringIO
import csv
import html2text
import httplib2
import json
import logging
import os
import re
import urllib
import yaml

OAUTH_SCOPE = 'https://www.googleapis.com/auth/drive'

# Silence extra logging from googleapiclient.
discovery.logger.setLevel(logging.WARNING)

class BaseGooglePreprocessor(base.BasePreprocessor):
    scheduleable = True

    @staticmethod
    def create_service():
        credentials = oauth.get_or_create_credentials(
            scope=OAUTH_SCOPE, storage_key='Grow SDK')
        http = httplib2.Http(ca_certs=utils.get_cacerts_path())
        http = credentials.authorize(http)
        return discovery.build('drive', 'v2', http=http)

    def run(self, build=True):
        try:
            self.execute(self.config)
        except errors.HttpError as e:
            self.logger.error(str(e))

    def execute(self, config):
        raise NotImplementedError

class GoogleDocsPreprocessor(BaseGooglePreprocessor):
    KIND = 'google_docs'

    class Config(messages.Message):
        path = messages.StringField(1)
        id = messages.StringField(2)
        convert = messages.BooleanField(3)

    @classmethod
    def _process_google_hrefs(cls, soup):
        for tag in soup.find_all('a'):
            if tag.attrs.get('href'):
                tag['href'] = cls._clean_google_href(tag['href'])

    @classmethod
    def _clean_google_href(cls, href):
        regex = ('^'
                 + re.escape('https://www.google.com/url?q=')
                 + '(.*?)'
                 + re.escape('&'))
        match = re.match(regex, href)
        if match:
            encoded_url = match.group(1)
            return urllib.unquote(encoded_url)
        return href

    @classmethod
    def download(cls, path, doc_id, logger=None, raise_errors=False):
        logger = logger or logging
        service = BaseGooglePreprocessor.create_service()
        resp = service.files().get(fileId=doc_id).execute()
        if 'exportLinks' not in resp:
            text = 'Unable to export Google Doc: {}'
            logger.error(text.format(path))
            logger.error('Received: {}'.format(resp))
            return
        for mimetype, url in resp['exportLinks'].iteritems():
            if not mimetype.endswith('html'):
                continue
            resp, content = service._http.request(url)
            if resp.status != 200:
                text = 'Error {} downloading Google Doc: {}'
                text = text.format(resp.status, path)
                if raise_errors:
                    raise base.PreprocessorError(text)
                logger.error(text)
            return content
        if raise_errors:
            text = 'No file to export from Google Docs: {}'.format(path)
            raise base.PreprocessorError(text)

    @classmethod
    def format_content(cls, path, content, convert=True, existing_data=None):
        ext = os.path.splitext(path)[1]
        convert_to_markdown = ext == '.md' and convert is not False
        soup = bs4.BeautifulSoup(content, 'html.parser')
        GoogleDocsPreprocessor._process_google_hrefs(soup)
        content = unicode(soup.body)
        if convert_to_markdown:
            h2t = html2text.HTML2Text()
            content = h2t.handle(content)
        content = content.encode('utf-8')
        # Preserve any existing frontmatter.
        if existing_data:
            if formats.Format.has_front_matter(existing_data):
                content = formats.Format.update(
                    existing_data, body=content)
        return content

    def execute(self, config):
        doc_id = config.id
        path = config.path
        convert = config.convert is not False
        content = GoogleDocsPreprocessor.download(
            path, doc_id=doc_id, logger=self.pod.logger)
        existing_data = None
        if self.pod.file_exists(path):
            existing_data = self.pod.read_file(path)
        content = GoogleDocsPreprocessor.format_content(
            path, content, convert=convert, existing_data=existing_data)
        self.pod.write_file(path, content)
        self.logger.info('Downloaded Google Doc -> {}'.format(path))

class GoogleSheetsPreprocessor(BaseGooglePreprocessor):
    KIND = 'google_sheets'

    class Config(messages.Message):
        path = messages.StringField(1)
        id = messages.StringField(2)
        gid = messages.IntegerField(3)
        output_style = messages.StringField(4, default='compressed')
        format = messages.StringField(5, default='list')
        preserve = messages.StringField(6, default='builtins')

    @staticmethod
    def _convert_rows_to_mapping(reader):
        results = {}
        def _update_node(root, part):
            if isinstance(root, dict) and part not in root:
                root[part] = {}
        for row in reader:
            key = row[0]
            value = row[1]
            if key.startswith('#'):
                continue
            if '.' in key:
                parts = key.split('.')
                parent = results
                for i, part in enumerate(parts):
                    _update_node(parent, part)
                    if i + 1 < len(parts):
                        parent = parent[part]
                if isinstance(parent, dict):
                    parent[part] = value
            else:
                results[key] = value
        return results

    @staticmethod
    def format_as_map(fp):
        reader = csv.reader(fp)
        results = GoogleSheetsPreprocessor._convert_rows_to_mapping(reader)
        return results

    @classmethod
    def download(cls, path, sheet_id, gid=None, logger=None,
                 raise_errors=False):
        logger = logger or logging
        ext = os.path.splitext(path)[1]
        service = BaseGooglePreprocessor.create_service()
        resp = service.files().get(fileId=sheet_id).execute()
        if 'exportLinks' not in resp:
            text = 'Unable to export Google Sheet: {} / Received: {}'
            logger.error(text.format(path, resp))
            if raise_errors:
                raise base.PreprocessorError(text)
            return
        for mimetype, url in resp['exportLinks'].iteritems():
            if not mimetype.endswith('csv'):
                continue
            if gid is not None:
                url += '&gid={}'.format(gid)
            resp, content = service._http.request(url)
            if resp.status != 200:
                text = 'Error {} downloading Google Sheet: {}'
                text = text.format(resp.status, path)
                logger.error(text)
                if raise_errors:
                    raise base.PreprocessorError(text)
            return content
        if raise_errors:
            text = 'No file to export from Google Sheets: {}'.format(path)
            raise base.PreprocessorError(text)

    def execute(self, config):
        path = config.path
        sheet_id = config.id
        gid = config.gid
        content = GoogleSheetsPreprocessor.download(
            path=path, sheet_id=sheet_id, gid=gid, logger=self.pod.logger)
        existing_data = None
        if (path.endswith(('.yaml', '.yml'))
                and self.config.preserve
                and self.pod.file_exists(path)):
            existing_data = self.pod.read_yaml(path)
        content = GoogleSheetsPreprocessor.format_content(
            content=content, path=path, format_as=self.config.format,
            preserve=self.config.preserve,
            existing_data=existing_data)
        content = GoogleSheetsPreprocessor.serialize_content(
            formatted_data=content, path=path,
            output_style=self.config.output_style)
        self.pod.write_file(path, content)
        self.logger.info('Downloaded Google Sheet -> {}'.format(path))

    @classmethod
    def get_convert_to(cls, path):
        ext = os.path.splitext(path)[1]
        convert_to = None
        if ext == '.json':
            return ext
        elif ext in ['.yaml', '.yml']:
            return ext
        return convert_to

    @classmethod
    def format_content(cls, content, path, format_as=None, preserve=None,
                       existing_data=None):
        """Formats content into either a CSV (text), list, or dictionary."""
        convert_to = cls.get_convert_to(path)
        if convert_to in ['.json', '.yaml', '.yml']:
            fp = cStringIO.StringIO()
            fp.write(content)
            fp.seek(0)
            if format_as == 'map':
                formatted_data = GoogleSheetsPreprocessor.format_as_map(fp)
            else:
                reader = csv.DictReader(fp)
                formatted_data = list(reader)
            if existing_data:
                if preserve == 'builtins':
                    for key in existing_data.keys():
                        if not key.startswith('$'):
                            del existing_data[key]
                existing_data.update(formatted_data)
                formatted_data = existing_data
            return formatted_data
        return content

    @classmethod
    def serialize_content(cls, formatted_data, path, output_style=None):
        """Serializes an object into a string as JSON, YAML, or a CSV
        (default)."""
        kwargs = {}
        convert_to = cls.get_convert_to(path)
        if convert_to == '.json':
            if output_style == 'pretty':
                kwargs['indent'] = 2
                kwargs['separators'] = (',', ': ')
                kwargs['sort_keys'] = True
            return json.dumps(formatted_data, **kwargs)
        elif convert_to in ('.yaml', '.yml'):
            return yaml.safe_dump(formatted_data, default_flow_style=False)
        return formatted_data

    def inject(self, doc):
        path = doc.pod_path
        try:
            content = GoogleSheetsPreprocessor.download(
                path=path, sheet_id=self.config.id, gid=self.config.gid,
                raise_errors=True)
        except (errors.HttpError, base.PreprocessorError):
            doc.pod.logger.error('Error downloading sheet -> %s', path)
            raise
        if not content:
            return
        existing_data = doc.pod.read_yaml(doc.pod_path)
        fields = GoogleSheetsPreprocessor.format_content(
            content, path=path, format_as=self.config.format,
            preserve=self.config.preserve, existing_data=existing_data)
        doc.inject(fields=fields)
<commit_msg>Support injection for Google Sheets preprocessor.<commit_after>"""
    Google Drive (Sheets and Docs) preprocessors allow you to store content in
    Google Drive and bring it into Grow. Grow will authenticate to the Google
    Drive API using OAuth2 and then download content as specified in
    `podspec.yaml`.

    Grow supports various ways to transform the content, e.g. Sheets can be
    downloaded and converted to yaml, and Docs can be downloaded and converted
    to markdown.
"""

from . import base
from googleapiclient import discovery
from googleapiclient import errors
from grow.common import oauth
from grow.common import utils
from grow.pods import formats
from protorpc import messages
import bs4
import cStringIO
import csv
import html2text
import httplib2
import json
import logging
import os
import re
import urllib
import yaml

OAUTH_SCOPE = 'https://www.googleapis.com/auth/drive'

# Silence extra logging from googleapiclient.
discovery.logger.setLevel(logging.WARNING)

class BaseGooglePreprocessor(base.BasePreprocessor):
    scheduleable = True

    @staticmethod
    def create_service():
        credentials = oauth.get_or_create_credentials(
            scope=OAUTH_SCOPE, storage_key='Grow SDK')
        http = httplib2.Http(ca_certs=utils.get_cacerts_path())
        http = credentials.authorize(http)
        return discovery.build('drive', 'v2', http=http)

    def run(self, build=True):
        try:
            self.execute(self.config)
        except errors.HttpError as e:
            self.logger.error(str(e))

    def execute(self, config):
        raise NotImplementedError

class GoogleDocsPreprocessor(BaseGooglePreprocessor):
    KIND = 'google_docs'

    class Config(messages.Message):
        path = messages.StringField(1)
        id = messages.StringField(2)
        convert = messages.BooleanField(3)

    @classmethod
    def _process_google_hrefs(cls, soup):
        for tag in soup.find_all('a'):
            if tag.attrs.get('href'):
                tag['href'] = cls._clean_google_href(tag['href'])

    @classmethod
    def _clean_google_href(cls, href):
        regex = ('^'
                 + re.escape('https://www.google.com/url?q=')
                 + '(.*?)'
                 + re.escape('&'))
        match = re.match(regex, href)
        if match:
            encoded_url = match.group(1)
            return urllib.unquote(encoded_url)
        return href

    @classmethod
    def download(cls, path, doc_id, logger=None, raise_errors=False):
        logger = logger or logging
        service = BaseGooglePreprocessor.create_service()
        resp = service.files().get(fileId=doc_id).execute()
        if 'exportLinks' not in resp:
            text = 'Unable to export Google Doc: {}'
            logger.error(text.format(path))
            logger.error('Received: {}'.format(resp))
            return
        for mimetype, url in resp['exportLinks'].iteritems():
            if not mimetype.endswith('html'):
                continue
            resp, content = service._http.request(url)
            if resp.status != 200:
                text = 'Error {} downloading Google Doc: {}'
                text = text.format(resp.status, path)
                if raise_errors:
                    raise base.PreprocessorError(text)
                logger.error(text)
            return content
        if raise_errors:
            text = 'No file to export from Google Docs: {}'.format(path)
            raise base.PreprocessorError(text)

    @classmethod
    def format_content(cls, path, content, convert=True, existing_data=None):
        ext = os.path.splitext(path)[1]
        convert_to_markdown = ext == '.md' and convert is not False
        soup = bs4.BeautifulSoup(content, 'html.parser')
        GoogleDocsPreprocessor._process_google_hrefs(soup)
        content = unicode(soup.body)
        if convert_to_markdown:
            h2t = html2text.HTML2Text()
            content = h2t.handle(content)
        content = content.encode('utf-8')
        # Preserve any existing frontmatter.
        if existing_data:
            if formats.Format.has_front_matter(existing_data):
                content = formats.Format.update(
                    existing_data, body=content)
        return content

    def execute(self, config):
        doc_id = config.id
        path = config.path
        convert = config.convert is not False
        content = GoogleDocsPreprocessor.download(
            path, doc_id=doc_id, logger=self.pod.logger)
        existing_data = None
        if self.pod.file_exists(path):
            existing_data = self.pod.read_file(path)
        content = GoogleDocsPreprocessor.format_content(
            path, content, convert=convert, existing_data=existing_data)
        self.pod.write_file(path, content)
        self.logger.info('Downloaded Google Doc -> {}'.format(path))

class GoogleSheetsPreprocessor(BaseGooglePreprocessor):
    KIND = 'google_sheets'

    class Config(messages.Message):
        path = messages.StringField(1)
        id = messages.StringField(2)
        gid = messages.IntegerField(3)
        output_style = messages.StringField(4, default='compressed')
        format = messages.StringField(5, default='list')
        preserve = messages.StringField(6, default='builtins')

    @staticmethod
    def _convert_rows_to_mapping(reader):
        results = {}
        def _update_node(root, part):
            if isinstance(root, dict) and part not in root:
                root[part] = {}
        for row in reader:
            key = row[0]
            value = row[1]
            if key.startswith('#'):
                continue
            if '.' in key:
                parts = key.split('.')
                parent = results
                for i, part in enumerate(parts):
                    _update_node(parent, part)
                    if i + 1 < len(parts):
                        parent = parent[part]
                if isinstance(parent, dict):
                    parent[part] = value
            else:
                results[key] = value
        return results

    @staticmethod
    def format_as_map(fp):
        reader = csv.reader(fp)
        results = GoogleSheetsPreprocessor._convert_rows_to_mapping(reader)
        return results

    @classmethod
    def download(cls, path, sheet_id, gid=None, logger=None,
                 raise_errors=False):
        logger = logger or logging
        ext = os.path.splitext(path)[1]
        service = BaseGooglePreprocessor.create_service()
        resp = service.files().get(fileId=sheet_id).execute()
        if 'exportLinks' not in resp:
            text = 'Unable to export Google Sheet: {} / Received: {}'
            logger.error(text.format(path, resp))
            if raise_errors:
                raise base.PreprocessorError(text)
            return
        for mimetype, url in resp['exportLinks'].iteritems():
            if not mimetype.endswith('csv'):
                continue
            if gid is not None:
                url += '&gid={}'.format(gid)
            resp, content = service._http.request(url)
            if resp.status != 200:
                text = 'Error {} downloading Google Sheet: {}'
                text = text.format(resp.status, path)
                logger.error(text)
                if raise_errors:
                    raise base.PreprocessorError(text)
            return content
        if raise_errors:
            text = 'No file to export from Google Sheets: {}'.format(path)
            raise base.PreprocessorError(text)

    def execute(self, config):
        path = config.path
        sheet_id = config.id
        gid = config.gid
        content = GoogleSheetsPreprocessor.download(
            path=path, sheet_id=sheet_id, gid=gid, logger=self.pod.logger)
        existing_data = None
        if (path.endswith(('.yaml', '.yml'))
                and self.config.preserve
                and self.pod.file_exists(path)):
            existing_data = self.pod.read_yaml(path)
        content = GoogleSheetsPreprocessor.format_content(
            content=content, path=path, format_as=self.config.format,
            preserve=self.config.preserve,
            existing_data=existing_data)
        content = GoogleSheetsPreprocessor.serialize_content(
            formatted_data=content, path=path,
            output_style=self.config.output_style)
        self.pod.write_file(path, content)
        self.logger.info('Downloaded Google Sheet -> {}'.format(path))

    @classmethod
    def get_convert_to(cls, path):
        ext = os.path.splitext(path)[1]
        convert_to = None
        if ext == '.json':
            return ext
        elif ext in ['.yaml', '.yml']:
            return ext
        return convert_to

    @classmethod
    def format_content(cls, content, path, format_as=None, preserve=None,
                       existing_data=None):
        """Formats content into either a CSV (text), list, or dictionary."""
        convert_to = cls.get_convert_to(path)
        if convert_to in ['.json', '.yaml', '.yml']:
            fp = cStringIO.StringIO()
            fp.write(content)
            fp.seek(0)
            if format_as == 'map':
                formatted_data = GoogleSheetsPreprocessor.format_as_map(fp)
            else:
                reader = csv.DictReader(fp)
                formatted_data = list(reader)
            if existing_data:
                if preserve == 'builtins':
                    for key in existing_data.keys():
                        if not key.startswith('$'):
                            del existing_data[key]
                existing_data.update(formatted_data)
                formatted_data = existing_data
            return formatted_data
        return content

    @classmethod
    def serialize_content(cls, formatted_data, path, output_style=None):
        """Serializes an object into a string as JSON, YAML, or a CSV
        (default)."""
        kwargs = {}
        convert_to = cls.get_convert_to(path)
        if convert_to == '.json':
            if output_style == 'pretty':
                kwargs['indent'] = 2
                kwargs['separators'] = (',', ': ')
                kwargs['sort_keys'] = True
            return json.dumps(formatted_data, **kwargs)
        elif convert_to in ('.yaml', '.yml'):
            return yaml.safe_dump(formatted_data, default_flow_style=False)
        return formatted_data

    def can_inject(self, doc=None, collection=None):
        if not self.injected:
            return False
        if doc and doc.pod_path == self.config.path:
            return True
        return False

    def inject(self, doc):
        path = doc.pod_path
        try:
            content = GoogleSheetsPreprocessor.download(
                path=path, sheet_id=self.config.id, gid=self.config.gid,
                raise_errors=True)
        except (errors.HttpError, base.PreprocessorError):
            doc.pod.logger.error('Error downloading sheet -> %s', path)
            raise
        if not content:
            return
        existing_data = doc.pod.read_yaml(doc.pod_path)
        fields = GoogleSheetsPreprocessor.format_content(
            content, path=path, format_as=self.config.format,
            preserve=self.config.preserve, existing_data=existing_data)
        fields = utils.untag_fields(fields)
        doc.inject(fields=fields)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2017 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function
import io
import os
import getpass
import logging
import numpy

from openquake.baselib import sap, config
from openquake.hazardlib import stats
from openquake.baselib import datastore
from openquake.commonlib.writers import write_csv
from openquake.commonlib.util import rmsep
from openquake.commonlib import logs, calc
from openquake.calculators.views import view
from openquake.calculators.extract import extract

if config.dbserver.multi_user:
    # get the datastore of the user who ran the job
    def read(calc_id):
        job = logs.dbcmd('get_job', calc_id, getpass.getuser())
        datadir = os.path.dirname(job.ds_calc_dir)
        return datastore.read(job.id, datadir=datadir)
else:  # get the datastore of the current user
    read = datastore.read

def get_hcurves_and_means(dstore):
    """
    Extract hcurves from the datastore and compute their means.

    :returns: curves_by_rlz, mean_curves
    """
    getter = calc.PmapGetter(dstore)
    sitecol = dstore['sitecol']
    pmaps = getter.get_pmaps(sitecol.sids)
    return dict(zip(getter.rlzs, pmaps)), dstore['hcurves/mean']

@sap.Script
def show(what='contents', calc_id=-1, extra=()):
    """
    Show the content of a datastore (by default the last one).
    """
    datadir = datastore.get_datadir()
    if what == 'all':  # show all
        if not os.path.exists(datadir):
            return
        rows = []
        for calc_id in datastore.get_calc_ids(datadir):
            try:
                ds = datastore.read(calc_id)
                oq = ds['oqparam']
                cmode, descr = oq.calculation_mode, oq.description
            except:
                # invalid datastore file, or missing calculation_mode
                # and description attributes, perhaps due to a manual kill
                f = os.path.join(datadir, 'calc_%s.hdf5' % calc_id)
                logging.warn('Unreadable datastore %s', f)
                continue
            else:
                rows.append((calc_id, cmode, descr.encode('utf-8')))
        for row in sorted(rows, key=lambda row: row[0]):  # by calc_id
            print('#%d %s: %s' % row)
        return

    ds = read(calc_id)

    # this part is experimental
    if what == 'rlzs' and 'poes' in ds:
        min_value = 0.01  # used in rmsep
        getter = calc.PmapGetter(ds)
        sitecol = ds['sitecol']
        pmaps = getter.get_pmaps(sitecol.sids)
        weights = [rlz.weight for rlz in getter.rlzs]
        mean = stats.compute_pmap_stats(pmaps, [numpy.mean], weights)
        dists = []
        for rlz, pmap in zip(getter.rlzs, pmaps):
            dist = rmsep(mean.array, pmap.array, min_value)
            dists.append((dist, rlz))
        print('Realizations in order of distance from the mean curves')
        for dist, rlz in sorted(dists):
            print('%s: rmsep=%s' % (rlz, dist))
    elif view.keyfunc(what) in view:
        print(view(what, ds))
    elif what.split('/', 1)[0] in extract:
        print(extract(ds, what, *extra))
    elif what in ds:
        obj = ds[what]
        if hasattr(obj, 'value'):  # an array
            print(write_csv(io.StringIO(), obj.value))
        else:
            print(obj)
    else:
        print('%s not found' % what)

    ds.close()

show.arg('what', 'key or view of the datastore')
show.arg('calc_id', 'calculation ID', type=int)
show.arg('extra', 'extra arguments', nargs='*')
<commit_msg>Fixed oq show<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2017 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function
import io
import os
import getpass
import logging
import numpy

from openquake.baselib import sap, config
from openquake.hazardlib import stats
from openquake.baselib import datastore
from openquake.commonlib.writers import write_csv
from openquake.commonlib.util import rmsep
from openquake.commonlib import logs, calc
from openquake.calculators.views import view
from openquake.calculators.extract import extract

if config.dbserver.multi_user:
    # get the datastore of the user who ran the job
    def read(calc_id):
        job = logs.dbcmd('get_job', calc_id, getpass.getuser())
        datadir = os.path.dirname(job.ds_calc_dir)
        return datastore.read(job.id, datadir=datadir)
else:  # get the datastore of the current user
    read = datastore.read

def get_hcurves_and_means(dstore):
    """
    Extract hcurves from the datastore and compute their means.

    :returns: curves_by_rlz, mean_curves
    """
    getter = calc.PmapGetter(dstore)
    sitecol = dstore['sitecol']
    pmaps = getter.get_pmaps(sitecol.sids)
    return dict(zip(getter.rlzs, pmaps)), dstore['hcurves/mean']

@sap.Script
def show(what='contents', calc_id=-1, extra=()):
    """
    Show the content of a datastore (by default the last one).
    """
    datadir = datastore.get_datadir()
    if what == 'all':  # show all
        if not os.path.exists(datadir):
            return
        rows = []
        for calc_id in datastore.get_calc_ids(datadir):
            try:
                ds = datastore.read(calc_id)
                oq = ds['oqparam']
                cmode, descr = oq.calculation_mode, oq.description
            except:
                # invalid datastore file, or missing calculation_mode
                # and description attributes, perhaps due to a manual kill
                f = os.path.join(datadir, 'calc_%s.hdf5' % calc_id)
                logging.warn('Unreadable datastore %s', f)
                continue
            else:
                rows.append((calc_id, cmode, descr.encode('utf-8')))
        for row in sorted(rows, key=lambda row: row[0]):  # by calc_id
            print('#%d %s: %s' % row)
        return

    ds = read(calc_id)

    # this part is experimental
    if what == 'rlzs' and 'poes' in ds:
        min_value = 0.01  # used in rmsep
        getter = calc.PmapGetter(ds)
        sitecol = ds['sitecol']
        pmaps = getter.get_pmaps(sitecol.sids)
        weights = [rlz.weight for rlz in getter.rlzs]
        mean = stats.compute_pmap_stats(pmaps, [numpy.mean], weights)
        dists = []
        for rlz, pmap in zip(getter.rlzs, pmaps):
            dist = rmsep(mean.array, pmap.array, min_value)
            dists.append((dist, rlz))
        print('Realizations in order of distance from the mean curves')
        for dist, rlz in sorted(dists):
            print('%s: rmsep=%s' % (rlz, dist))
    elif view.keyfunc(what) in view:
        print(view(what, ds))
    elif what.split('/', 1)[0] in extract:
        print(extract(ds, what, *extra))
    elif what in ds:
        obj = ds[what]
        if hasattr(obj, 'value'):  # an array
            print(write_csv(io.BytesIO(), obj.value).decode('utf8'))
        else:
            print(obj)
    else:
        print('%s not found' % what)

    ds.close()

show.arg('what', 'key or view of the datastore')
show.arg('calc_id', 'calculation ID', type=int)
show.arg('extra', 'extra arguments', nargs='*')
<|endoftext|>


<commit_before>import numpy
from numpy.testing import assert_array_equal

import unittest2 as unittest

from tkp.sourcefinder.utils import maximum_pixel_method_variance, fudge_max_pix
from tkp.sourcefinder.utils import circular_mask

class TestCircularMask(unittest.TestCase):
    def test_known_results(self):
        known_results = [
            ((0,0,0), numpy.zeros([0,0])),
            ((0,0,100), numpy.zeros([0,0])),
            ((1,1,0), numpy.array([[True]])),
            ((1,1,1), numpy.array([[False]])),
            ((5,1,0), numpy.array([[True],[True],[True],[True],[True]])),
            ((5,1,2), numpy.array([[True],[False],[False],[False],[True]])),
            ((1,4,0), numpy.array([[True,True,True,True]])),
            ((1,4,1), numpy.array([[True,False,False,True]]))
        ]
        for parameters, result in known_results:
            assert_array_equal(circular_mask(*parameters), result)

class UtilsTest(unittest.TestCase):
    # format: semimajor, semiminor, theta, correction, variance
    def setUp(self):
        self.correct_data = (
            ( 1.46255645752, 1.30175183614, 1.51619779602, 1.06380431709, 0.0017438970306615786 ),
            ( 1.56219151815, 1.25084762573, 1.49080535398, 1.06328376441, 0.0017765054643135159 ),
            ( 1.55231285095, 1.24513448079, 1.44660127410, 1.06398549354, 0.0018208980607872238 )
        )
    def testMaximumPixelMethodVariance(self):
        for (semimajor, semiminor, theta, correction, variance) in self.correct_data:
            self.assertEqual(maximum_pixel_method_variance(semimajor, semiminor, theta), variance)

    def testFudgeMaxPix(self):
        for (semimajor, semiminor, theta, correction, variance) in self.correct_data:
            self.assertAlmostEqual(fudge_max_pix(semimajor, semiminor, theta), correction)

<commit_msg>Test case for #5236<commit_after>import numpy
from numpy.testing import assert_array_equal
from collections import namedtuple
from tkp.utility.uncertain import Uncertain

import unittest2 as unittest

from tkp.sourcefinder.utils import maximum_pixel_method_variance, fudge_max_pix
from tkp.sourcefinder.utils import generate_result_maps
from tkp.sourcefinder.utils import circular_mask

class TestCircularMask(unittest.TestCase):
    def test_known_results(self):
        known_results = [
            ((0,0,0), numpy.zeros([0,0])),
            ((0,0,100), numpy.zeros([0,0])),
            ((1,1,0), numpy.array([[True]])),
            ((1,1,1), numpy.array([[False]])),
            ((5,1,0), numpy.array([[True],[True],[True],[True],[True]])),
            ((5,1,2), numpy.array([[True],[False],[False],[False],[True]])),
            ((1,4,0), numpy.array([[True,True,True,True]])),
            ((1,4,1), numpy.array([[True,False,False,True]]))
        ]
        for parameters, result in known_results:
            assert_array_equal(circular_mask(*parameters), result)

class TestResultMaps(unittest.TestCase):
    def testPositions(self):
        # The pixel position x, y of the source should be the same as the
        # position of the maximum in the generated result map.
        x_pos = 40
        y_pos = 60

        empty_data = numpy.zeros((100, 100))
        SrcMeasurement = namedtuple("SrcMeasurement",
            ['peak', 'x', 'y', 'smaj', 'smin', 'theta']
        )
        src_measurement = SrcMeasurement(
            peak=Uncertain(10), x=Uncertain(x_pos), y=Uncertain(y_pos),
            smaj=Uncertain(10), smin=Uncertain(10), theta=Uncertain(0)
        )
        gaussian_map, residual_map = generate_result_maps(empty_data, [src_measurement])
        gaussian_max_x = numpy.where(gaussian_map==gaussian_map.max())[0][0]
        gaussian_max_y = numpy.where(gaussian_map==gaussian_map.max())[1][0]

        self.assertEqual(gaussian_max_x, x_pos)
        self.assertEqual(gaussian_max_y, y_pos)

class UtilsTest(unittest.TestCase):
    # format: semimajor, semiminor, theta, correction, variance
    def setUp(self):
        self.correct_data = (
            ( 1.46255645752, 1.30175183614, 1.51619779602, 1.06380431709, 0.0017438970306615786 ),
            ( 1.56219151815, 1.25084762573, 1.49080535398, 1.06328376441, 0.0017765054643135159 ),
            ( 1.55231285095, 1.24513448079, 1.44660127410, 1.06398549354, 0.0018208980607872238 )
        )
    def testMaximumPixelMethodVariance(self):
        for (semimajor, semiminor, theta, correction, variance) in self.correct_data:
            self.assertEqual(maximum_pixel_method_variance(semimajor, semiminor, theta), variance)

    def testFudgeMaxPix(self):
        for (semimajor, semiminor, theta, correction, variance) in self.correct_data:
            self.assertAlmostEqual(fudge_max_pix(semimajor, semiminor, theta), correction)

<|endoftext|>


<commit_before>from django.utils import simplejson as json

from open_municipio.data_import.lib import DataSource, BaseReader, BaseWriter, JSONWriter, XMLWriter, OMWriter
# import OM-XML language tags
from open_municipio.data_import.om_xml import *

from lxml import etree

class ActsDataSource(DataSource):
    """
    A data source containing acts-related data.
    """
    pass

class BaseActsReader(BaseReader):
    """
    An abstract reader class parsing acts-relate//d data.

    This class encapsulates generic -- but acts-specific -- parsing logic.
    It's intended to be subclassed (and its methods overriden) in order to adapt to
    concrete scenarios.
    """
    def read(self):
        # initialize the reader
        self.setup()
        # get the data source to read from
        data_source = self.get_data_source()
        # initialize the data source
        data_source.setup()

        self.acts = []

        # TODO extract the relevant information

        return self.acts;

class BaseActsWriter(BaseWriter):
    """
    An abstract writer class that serializes internal representations
    of acts-related data.

    This class encapsulates acts-specific serialization logic which is independent
    from a specific output format.
    """
    pass

class OMActsWriter(BaseActsWriter, OMWriter):
    """
    A writer class which outputs acts data as objects in the OpenMunicio model.
    """
    def write(self):
        raise NotImplementedError

# python object layer for imported data

# document section

class Act:
    id = ""
    content = ""
    file = None
    subscribers = [] # list of Charges

    def add_subscriber(self, charge):
        self.subscribers.append(charge)

class CouncilDeliberation(Act):
    final_id = ""
    execution_date = ""
    initiative = "" # (council_member, ...)

class CityGovernmentDeliberation(Act):
    pass

class Interrogation(Act):
    pass

class Interpellation(Act):
    pass

class Motion(Act):
    pass

class Agenda(Act):
    pass

class Emendation(Act):
    pass

class Attachment:
    description = ""
    file = None
    type = "" # type of file (define an array of possible file types)

# people section

class Person:
    id = ""
    name = ""
    surname = ""
    ssn = ""

    def __init__(self, id, name, surname, ssn):
        self.id = id
        self.name = name
        self.surname = surname
        self.ssn = ssn

    def __str__(self):
        return "%s %s (%s)" % (self.name, self.surname, self.id)

    def __unicode__(self):
        return u"%s %s (%s)" % (self.name, self.surname, self.id)

class Charge:
    id = ""
    start_date = ""
    person = None
    name = None
    description = ""

    def __init__(self, id, start_date, person, name, description):
        self.id = id
        self.start_date = start_date
        self.person = person
        self.name = name
        self.description = description

# institution section

class Institution:
    charges = []

    def add_charge(self, charge):
        self.charges.append(charge)

class Mayor(Institution):

    def __init__(self, mayor):
        self.append(mayor)

class CityGovernment(Institution):
    pass

class CityCouncil(Institution):
    pass

class Subscriber:
    charge = None
    type = "" # first subscriber or co-subscriber

<commit_msg>added __str__ and __unicode__ for Charge in data_import<commit_after>from django.utils import simplejson as json

from open_municipio.data_import.lib import DataSource, BaseReader, BaseWriter, JSONWriter, XMLWriter, OMWriter
# import OM-XML language tags
from open_municipio.data_import.om_xml import *

from lxml import etree

class ActsDataSource(DataSource):
    """
    A data source containing acts-related data.
    """
    pass

class BaseActsReader(BaseReader):
    """
    An abstract reader class parsing acts-relate//d data.

    This class encapsulates generic -- but acts-specific -- parsing logic.
    It's intended to be subclassed (and its methods overriden) in order to adapt to
    concrete scenarios.
    """
    def read(self):
        # initialize the reader
        self.setup()
        # get the data source to read from
        data_source = self.get_data_source()
        # initialize the data source
        data_source.setup()

        self.acts = []

        # TODO extract the relevant information

        return self.acts;

class BaseActsWriter(BaseWriter):
    """
    An abstract writer class that serializes internal representations
    of acts-related data.

    This class encapsulates acts-specific serialization logic which is independent
    from a specific output format.
    """
    pass

class OMActsWriter(BaseActsWriter, OMWriter):
    """
    A writer class which outputs acts data as objects in the OpenMunicio model.
    """
    def write(self):
        raise NotImplementedError

# python object layer for imported data

# document section

class Act:
    id = ""
    content = ""
    file = None
    subscribers = [] # list of Charges

    def add_subscriber(self, charge):
        self.subscribers.append(charge)

class CouncilDeliberation(Act):
    final_id = ""
    execution_date = ""
    initiative = "" # (council_member, ...)

class CityGovernmentDeliberation(Act):
    pass

class Interrogation(Act):
    pass

class Interpellation(Act):
    pass

class Motion(Act):
    pass

class Agenda(Act):
    pass

class Emendation(Act):
    pass

class Attachment:
    description = ""
    file = None
    type = "" # type of file (define an array of possible file types)

# people section

class Person:
    id = ""
    name = ""
    surname = ""
    ssn = ""

    def __init__(self, id, name, surname, ssn):
        self.id = id
        self.name = name
        self.surname = surname
        self.ssn = ssn

    def __str__(self):
        return "%s %s (%s)" % (self.name, self.surname, self.id)

    def __unicode__(self):
        return u"%s %s (%s)" % (self.name, self.surname, self.id)

class Charge:
    id = ""
    start_date = ""
    person = None
    name = None
    description = ""

    def __init__(self, id, start_date, person, name, description):
        self.id = id
        self.start_date = start_date
        self.person = person
        self.name = name
        self.description = description

    def __str__(self):
        return "%s as %s from %s (%d)" % (self.person, self.name, self.start_date, self.id)

    def __unicode__(self):
        return u"%s as %s from %s (%d)" % (self.person, self.name, self.start_date, self.id)

# institution section

class Institution:
    charges = []

    def add_charge(self, charge):
        self.charges.append(charge)

class Mayor(Institution):

    def __init__(self, mayor):
        self.append(mayor)

class CityGovernment(Institution):
    pass

class CityCouncil(Institution):
    pass

class Subscriber:
    charge = None
    type = "" # first subscriber or co-subscriber

<|endoftext|>


<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2014-2016, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Here is a minimal example of usage:

.. code-block:: python

    >>> from openquake.commonlib import sap
    >>> def fun(input, inplace, output=None, out='/tmp'):
    ...     'Example'
    ...     for argname, argvalue in sorted(locals().iteritems()):
    ...         print argname, '=', argvalue

    >>> p = sap.Parser(fun)
    >>> p.arg('input', 'input file or archive')
    >>> p.flg('inplace', 'convert inplace')
    >>> p.arg('output', 'output archive')
    >>> p.opt('out', 'optional output file')

    >>> p.callfunc(['a'])
    inplace = False
    input = a
    out = /tmp
    output = None

    >>> p.callfunc(['a', 'b', '-i', '-o', 'OUT'])
    inplace = True
    input = a
    out = OUT
    output = b

Parsers can be composed too.
"""

import sys
import inspect
import argparse
from collections import OrderedDict

NODEFAULT = object()

def get_parentparser(parser, description=None, help=True):
    """
    :param parser: :class:`argparse.ArgumentParser` instance or None
    :param description: string used to build a new parser if parser is None
    :param help: flag used to build a new parser if parser is None
    :returns: if parser is None the new parser; otherwise the `.parentparser`
              attribute (if set) or the parser itself (if not set)
    """
    if parser is None:
        return argparse.ArgumentParser(
            description=description, add_help=help)
    elif hasattr(parser, 'parentparser'):
        return parser.parentparser
    else:
        return parser

class Parser(object):
    """
    A simple way to define command processors based on argparse.
    Each parser is associated to a function and parsers can be
    composed together, by dispatching on a given name (if not given,
    the function name is used).
    """
    def __init__(self, func, name=None, parentparser=None, help=True):
        self.func = func
        self.name = name or func.__name__
        args, self.varargs, varkw, defaults = inspect.getargspec(func)
        assert self.varargs is None, self.varargs
        defaults = defaults or ()
        nodefaults = len(args) - len(defaults)
        alldefaults = (NODEFAULT,) * nodefaults + defaults
        self.argdict = OrderedDict(zip(args, alldefaults))
        self.parentparser = get_parentparser(
            parentparser, description=func.__doc__, help=help)
        self.names = set()
        self.all_arguments = []
        self._group = self.parentparser
        self._argno = 0  # used in the NameError check in the _add method
        self.checked = False  # used in the check_arguments method

    def group(self, descr):
        """Added a new group of arguments with the given description"""
        self._group = self.parentparser.add_argument_group(descr)

    def _add(self, name, *args, **kw):
        """
        Add an argument to the underlying parser and grow the list
        .all_arguments and the set .names
        """
        argname = list(self.argdict)[self._argno]
        if argname != name:
            raise NameError(
                'Setting argument %s, but it should be %s' % (name, argname))
        self._group.add_argument(*args, **kw)
        self.all_arguments.append((args, kw))
        self.names.add(name)
        self._argno += 1

    def arg(self, name, help, type=None, choices=None, metavar=None,
            nargs='?'):
        """Describe a positional argument"""
        kw = dict(help=help, type=type, choices=choices, metavar=metavar)
        default = self.argdict[name]
        if default is not NODEFAULT:
            kw['nargs'] = nargs
            kw['default'] = default
            kw['help'] = kw['help'] + ' [default: %s]' % str(default)
        self._add(name, name, **kw)

    def opt(self, name, help, abbrev=None,
            type=None, choices=None, metavar=None, nargs=None):
        """Describe an option"""
        kw = dict(help=help, type=type, choices=choices, metavar=metavar,
                  nargs=nargs)
        default = self.argdict[name]
        if default is not NODEFAULT:
            kw['default'] = default
            kw['metavar'] = metavar or str(default)
        abbrev = abbrev or '-' + name[0]
        abbrevs = set(args[0] for args, kw in self.all_arguments)
        longname = '--' + name.replace('_', '-')
        if abbrev == '-h' or abbrev in abbrevs:
            # avoid conflicts with predefined abbreviations
            self._add(name, longname, **kw)
        else:
            self._add(name, abbrev, longname, **kw)

    def flg(self, name, help, abbrev=None):
        """Describe a flag"""
        abbrev = abbrev or '-' + name[0]
        longname = '--' + name.replace('_', '-')
        self._add(name, abbrev, longname, action='store_true', help=help)

    def check_arguments(self):
        """Make sure all arguments have a specification"""
        for name, default in self.argdict.items():
            if name not in self.names and default is NODEFAULT:
                raise NameError('Missing argparse specification for %r' % name)

    def callfunc(self, argv=None):
        """
        Parse the argv list and extract a dictionary of arguments which
        is then passed to  the function underlying the Parser.
        """
        if not self.checked:
            self.check_arguments()
            self.checked = True
        namespace = self.parentparser.parse_args(argv or sys.argv[1:])
        return self.func(**vars(namespace))

    def help(self):
        """
        Return the help message as a string
        """
        return self.parentparser.format_help()

def compose(parsers, name='main', description=None, help=True, prog=None):
    """
    Collects together different arguments parsers and builds a single
    Parser dispatching on the subparsers depending on
    the first argument, i.e. the name of the subparser to invoke.

    :param parsers: a list of Parser instances
    :param name: the name of the composed parser
    :param description: description of the composed parser
    :param help: help flag
    :param prog: name of the script printed in the usage message
    """
    assert len(parsers) >= 1, parsers
    parentparser = argparse.ArgumentParser(
        description=description, add_help=help)
    subparsers = parentparser.add_subparsers(
        help='available subcommands (see sub help)', prog=prog)
    for p in parsers:
        subp = subparsers.add_parser(p.name)
        for args, kw in p.all_arguments:
            subp.add_argument(*args, **kw)
        subp.set_defaults(_func=p.func)

    def main(**kw):
        func = kw.pop('_func')
        return func(**kw)
    main.__name__ = name
    return Parser(main, name, parentparser)
<commit_msg>Fixed nargs in sap<commit_after>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2014-2016, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Here is a minimal example of usage:

.. code-block:: python

    >>> from openquake.commonlib import sap
    >>> def fun(input, inplace, output=None, out='/tmp'):
    ...     'Example'
    ...     for argname, argvalue in sorted(locals().iteritems()):
    ...         print argname, '=', argvalue

    >>> p = sap.Parser(fun)
    >>> p.arg('input', 'input file or archive')
    >>> p.flg('inplace', 'convert inplace')
    >>> p.arg('output', 'output archive')
    >>> p.opt('out', 'optional output file')

    >>> p.callfunc(['a'])
    inplace = False
    input = a
    out = /tmp
    output = None

    >>> p.callfunc(['a', 'b', '-i', '-o', 'OUT'])
    inplace = True
    input = a
    out = OUT
    output = b

Parsers can be composed too.
"""

import sys
import inspect
import argparse
from collections import OrderedDict

NODEFAULT = object()

def get_parentparser(parser, description=None, help=True):
    """
    :param parser: :class:`argparse.ArgumentParser` instance or None
    :param description: string used to build a new parser if parser is None
    :param help: flag used to build a new parser if parser is None
    :returns: if parser is None the new parser; otherwise the `.parentparser`
              attribute (if set) or the parser itself (if not set)
    """
    if parser is None:
        return argparse.ArgumentParser(
            description=description, add_help=help)
    elif hasattr(parser, 'parentparser'):
        return parser.parentparser
    else:
        return parser

class Parser(object):
    """
    A simple way to define command processors based on argparse.
    Each parser is associated to a function and parsers can be
    composed together, by dispatching on a given name (if not given,
    the function name is used).
    """
    def __init__(self, func, name=None, parentparser=None, help=True):
        self.func = func
        self.name = name or func.__name__
        args, self.varargs, varkw, defaults = inspect.getargspec(func)
        assert self.varargs is None, self.varargs
        defaults = defaults or ()
        nodefaults = len(args) - len(defaults)
        alldefaults = (NODEFAULT,) * nodefaults + defaults
        self.argdict = OrderedDict(zip(args, alldefaults))
        self.parentparser = get_parentparser(
            parentparser, description=func.__doc__, help=help)
        self.names = set()
        self.all_arguments = []
        self._group = self.parentparser
        self._argno = 0  # used in the NameError check in the _add method
        self.checked = False  # used in the check_arguments method

    def group(self, descr):
        """Added a new group of arguments with the given description"""
        self._group = self.parentparser.add_argument_group(descr)

    def _add(self, name, *args, **kw):
        """
        Add an argument to the underlying parser and grow the list
        .all_arguments and the set .names
        """
        argname = list(self.argdict)[self._argno]
        if argname != name:
            raise NameError(
                'Setting argument %s, but it should be %s' % (name, argname))
        self._group.add_argument(*args, **kw)
        self.all_arguments.append((args, kw))
        self.names.add(name)
        self._argno += 1

    def arg(self, name, help, type=None, choices=None, metavar=None,
            nargs=None):
        """Describe a positional argument"""
        kw = dict(help=help, type=type, choices=choices, metavar=metavar,
                  nargs=nargs)
        default = self.argdict[name]
        if default is not NODEFAULT:
            kw['nargs'] = nargs or '?'
            kw['default'] = default
            kw['help'] = kw['help'] + ' [default: %s]' % str(default)
        self._add(name, name, **kw)

    def opt(self, name, help, abbrev=None,
            type=None, choices=None, metavar=None, nargs=None):
        """Describe an option"""
        kw = dict(help=help, type=type, choices=choices, metavar=metavar,
                  nargs=nargs)
        default = self.argdict[name]
        if default is not NODEFAULT:
            kw['default'] = default
            kw['metavar'] = metavar or str(default)
        abbrev = abbrev or '-' + name[0]
        abbrevs = set(args[0] for args, kw in self.all_arguments)
        longname = '--' + name.replace('_', '-')
        if abbrev == '-h' or abbrev in abbrevs:
            # avoid conflicts with predefined abbreviations
            self._add(name, longname, **kw)
        else:
            self._add(name, abbrev, longname, **kw)

    def flg(self, name, help, abbrev=None):
        """Describe a flag"""
        abbrev = abbrev or '-' + name[0]
        longname = '--' + name.replace('_', '-')
        self._add(name, abbrev, longname, action='store_true', help=help)

    def check_arguments(self):
        """Make sure all arguments have a specification"""
        for name, default in self.argdict.items():
            if name not in self.names and default is NODEFAULT:
                raise NameError('Missing argparse specification for %r' % name)

    def callfunc(self, argv=None):
        """
        Parse the argv list and extract a dictionary of arguments which
        is then passed to  the function underlying the Parser.
        """
        if not self.checked:
            self.check_arguments()
            self.checked = True
        namespace = self.parentparser.parse_args(argv or sys.argv[1:])
        return self.func(**vars(namespace))

    def help(self):
        """
        Return the help message as a string
        """
        return self.parentparser.format_help()

def compose(parsers, name='main', description=None, help=True, prog=None):
    """
    Collects together different arguments parsers and builds a single
    Parser dispatching on the subparsers depending on
    the first argument, i.e. the name of the subparser to invoke.

    :param parsers: a list of Parser instances
    :param name: the name of the composed parser
    :param description: description of the composed parser
    :param help: help flag
    :param prog: name of the script printed in the usage message
    """
    assert len(parsers) >= 1, parsers
    parentparser = argparse.ArgumentParser(
        description=description, add_help=help)
    subparsers = parentparser.add_subparsers(
        help='available subcommands (see sub help)', prog=prog)
    for p in parsers:
        subp = subparsers.add_parser(p.name)
        for args, kw in p.all_arguments:
            subp.add_argument(*args, **kw)
        subp.set_defaults(_func=p.func)

    def main(**kw):
        func = kw.pop('_func')
        return func(**kw)
    main.__name__ = name
    return Parser(main, name, parentparser)
<|endoftext|>