<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import numpy as np
from astropy.convolution import Kernel2D
from ..extern.imageutils import sigmaclip_stats

__all__ = ['detect_threshold', 'detect_sources', 'find_peaks']

def detect_threshold(data, snr, background=None, error=None, mask=None,
                     mask_val=None, sigclip_sigma=3.0, sigclip_iters=None):
    """
    Calculate a pixel-wise threshold image to be used to detect sources.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    snr : float
        The signal-to-noise ratio per pixel above the ``background`` for
        which to consider a pixel as possibly being part of a source.

    background : float or array_like, optional
        The background value(s) of the input ``data``.  ``background``
        may either be a scalar value or a 2D image with the same shape
        as the input ``data``.  If the input ``data`` has been
        background-subtracted, then set ``background`` to ``0.0``.  If
        `None`, then a scalar background value will be estimated using
        sigma-clipped statistics.

    error : float or array_like, optional
        The Gaussian 1-sigma standard deviation of the background noise
        in ``data``.  ``error`` should include all sources of
        "background" error, but *exclude* the Poisson error of the
        sources.  If ``error`` is a 2D image, then it should represent
        the 1-sigma background error in each pixel of ``data``.  If
        `None`, then a scalar background rms value will be estimated
        using sigma-clipped statistics.

    mask : array_like, bool, optional
        A boolean mask with the same shape as ``data``, where a `True`
        value indicates the corresponding element of ``data`` is masked.
        Masked pixels are ignored when computing the image background
        statistics.

    mask_val : float, optional
        An image data value (e.g., ``0.0``) that is ignored when
        computing the image background statistics.  ``mask_val`` will be
        ignored if ``mask`` is input.

    sigclip_sigma : float, optional
        The number of standard deviations to use as the clipping limit
        when calculating the image background statistics.

    sigclip_iters : float, optional
       The number of iterations to perform sigma clipping, or `None` to
       clip until convergence is achieved (i.e., continue until the last
       iteration clips nothing) when calculating the image background
       statistics.

    Returns
    -------
    threshold : 2D `~numpy.ndarray`
        A 2D image with the same shape as ``data`` containing the
        pixel-wise threshold values.

    See Also
    --------
    detect_sources

    Notes
    -----
    The ``mask``, ``mask_val``, ``sigclip_sigma``, and ``sigclip_iters``
    inputs are used only if it is necessary to estimate ``background``
    or ``error`` using sigma-clipped background statistics.  If
    ``background`` and ``error`` are both input, then ``mask``,
    ``mask_val``, ``sigclip_sigma``, and ``sigclip_iters`` are ignored.
    """

    if background is None or error is None:
        data_mean, data_median, data_std = sigmaclip_stats(
            data, mask=mask, mask_val=mask_val, sigma=sigclip_sigma,
            iters=sigclip_iters)
        bkgrd_image = np.zeros_like(data) + data_mean
        bkgrdrms_image = np.zeros_like(data) + data_std

    if background is None:
        background = bkgrd_image
    else:
        if np.isscalar(background):
            background = np.zeros_like(data) + background
        else:
            if background.shape != data.shape:
                raise ValueError('If input background is 2D, then it '
                                 'must have the same shape as the input '
                                 'data.')

    if error is None:
        error = bkgrdrms_image
    else:
        if np.isscalar(error):
            error = np.zeros_like(data) + error
        else:
            if error.shape != data.shape:
                raise ValueError('If input error is 2D, then it '
                                 'must have the same shape as the input '
                                 'data.')

    return background + (error * snr)

def detect_sources(data, threshold, npixels, filter_kernel=None,
                   connectivity=8):
    """
    Detect sources above a specified threshold value in an image and
    return a segmentation image.

    Detected sources must have ``npixels`` connected pixels that are
    each greater than the ``threshold`` value.  If the filtering option
    is used, then the ``threshold`` is applied to the filtered image.

    This function does not deblend overlapping sources.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    threshold : float or array-like
        The data value or pixel-wise data values to be used for the
        detection threshold.  A 2D ``threshold`` must have the same
        shape as ``data``.  See `detect_threshold` for one way to create
        a ``threshold`` image.

    npixels : int
        The number of connected pixels, each greater than ``threshold``,
        that an object must have to be detected.  ``npixels`` must be a
        positive integer.

    filter_kernel : array-like (2D) or `~astropy.convolution.Kernel2D`, optional
        The 2D array of the kernel used to filter the image before
        thresholding.  Filtering the image will smooth the noise and
        maximize detectability of objects with a shape similar to the
        kernel.

    connectivity : {4, 8}, optional
        The type of pixel connectivity used in determining how pixels
        are grouped into a detected source.  The options are 4 or 8
        (default).  4-connected pixels touch along their edges.
        8-connected pixels touch along their edges or corners.  For
        reference, SExtractor uses 8-connected pixels.

    Returns
    -------
    segment_image : `~numpy.ndarray` (int)
        A 2D segmentation image, with the same shape as ``data``, where
        sources are marked by different positive integer values.  A
        value of zero is reserved for the background.

    See Also
    --------
    detect_threshold, :class:`photutils.segmentation.segment_properties`

    Examples
    --------

    .. plot::
        :include-source:

        # make a table of Gaussian sources
        from astropy.table import Table
        table = Table()
        table['amplitude'] = [50, 70, 150, 210]
        table['x_mean'] = [160, 25, 150, 90]
        table['y_mean'] = [70, 40, 25, 60]
        table['x_stddev'] = [15.2, 5.1, 3., 8.1]
        table['y_stddev'] = [2.6, 2.5, 3., 4.7]
        table['theta'] = np.array([145., 20., 0., 60.]) * np.pi / 180.

        # make an image of the sources with Gaussian noise
        from photutils.datasets import make_gaussian_sources
        from photutils.datasets import make_noise_image
        shape = (100, 200)
        sources = make_gaussian_sources(shape, table)
        noise = make_noise_image(shape, type='gaussian', mean=0.,
                                 stddev=5., random_state=12345)
        image = sources + noise

        # detect the sources
        from photutils import detect_threshold, detect_sources
        threshold = detect_threshold(image, snr=3)
        from astropy.convolution import Gaussian2DKernel
        sigma = 3.0 / (2.0 * np.sqrt(2.0 * np.log(2.0)))   # FWHM = 3
        filter_kernel = Gaussian2DKernel(sigma, x_size=3, y_size=3)
        segm_image = detect_sources(image, threshold, npixels=5,
                                    filter_kernel=filter_kernel)

        # plot the image and the segmentation image
        import matplotlib.pyplot as plt
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))
        ax1.imshow(image, origin='lower', interpolation='nearest')
        ax2.imshow(segm_image, origin='lower', interpolation='nearest')
    """

    from scipy import ndimage

    if (npixels <= 0) or (int(npixels) != npixels):
        raise ValueError('npixels must be a positive integer, got '
                         '"{0}"'.format(npixels))

    conv_mode = 'constant'    # SExtractor mode
    conv_val = 0.0
    if filter_kernel is not None:
        if isinstance(filter_kernel, Kernel2D):
            image = ndimage.convolve(data, filter_kernel.array, mode=conv_mode,
                                     cval=conv_val)
        else:
            image = ndimage.convolve(data, filter_kernel, mode=conv_mode,
                                     cval=conv_val)
    else:
        image = data

    image = (image > threshold)
    if connectivity == 4:
        selem = ndimage.generate_binary_structure(2, 1)
    elif connectivity == 8:    # e.g., SExtractor
        selem = ndimage.generate_binary_structure(2, 2)
    objlabels, nobj = ndimage.label(image, structure=selem)
    objslices = ndimage.find_objects(objlabels)

    # remove objects with less than npixels
    for objslice in objslices:
        objlabel = objlabels[objslice]
        obj_npix = len(np.where(objlabel.ravel() != 0)[0])
        if obj_npix < npixels:
            objlabels[objslice] = 0

    # relabel to make sequential label indices
    objlabels, nobj = ndimage.label(objlabels, structure=selem)
    return objlabels

def find_peaks(data, threshold, min_separation=2, exclude_border=True,
               segment_image=None, npeaks=np.inf, footprint=None):
    """
    Find local peaks in an image that are above above a specified
    threshold value.

    Peaks are the local maxima above the ``threshold`` in a region of
    ``(2 * min_separation) + 1`` (i.e., peaks are separated by at least
    ``min_separation`` pixels).

    If peaks are flat (i.e., multiple adjacent pixels have identical
    intensities), then the coordinates of all such pixels are returned,
    even if they are not separated by at least ``min_separation``.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    threshold : float or array-like
        The data value or pixel-wise data values to be used for the
        detection threshold.  A 2D ``threshold`` must have the same
        shape as ``data``.  See `detect_threshold` for one way to create
        a ``threshold`` image.

    min_separation : int, optional
        Minimum number of pixels separating peaks (i.e., peaks are
        separated by at least ``min_separation`` pixels).  To find the
        maximum number of peaks, use ``min_separation=1``.  If
        ``min_separation`` is not an integer, then it will be truncated.

    exclude_border : bool, optional
        If `True`, exclude peaks within ``min_separation`` from the
        border of the image as well as from each other.

    segment_image : `~numpy.ndarray` (int), optional
        If provided, then search for peaks located only within the
        labeled regions of a 2D segmentation image, where sources are
        marked by different positive integer values.  In the
        segmentation image a value of zero is reserved for the
        background.  ``segment_image`` must have the same shape as
        ``data``.

    npeaks : int, optional
        The maximum number of peaks to return.  When the number of
        detected peaks exceeds ``npeaks``, the peaks with the highest
        peak intensities will be returned.

    footprint : `~numpy.ndarray` of bools, optional
        A boolean array where `True` values describe the local footprint
        region within which to search for peaks at every point in
        ``data``.  Overrides ``min_separation``, except for border
        exclusion when ``exclude_border=True``.

    Returns
    -------
    output : `~numpy.ndarray`
        An ``Nx2`` array where the rows contain the ``(y, x)`` pixel
        coordinates of the local peaks.
    """

    if segment_image is not None:
        if segment_image.shape != data.shape:
            raise ValueError('segment_image and data must have the same '
                             'shape')

    from skimage.feature import peak_local_max
    coords = peak_local_max(data, min_distance=int(min_separation),
                            threshold_abs=threshold, threshold_rel=0.0,
                            exclude_border=exclude_border, indices=True,
                            num_peaks=npeaks, footprint=footprint,
                            labels=segment_image)
    if coords.shape[0] <= npeaks:
        return coords
    else:
        # NOTE: num_peaks is ignored by peak_local_max() if labels are input
        peak_values = data[coords[:, 0], coords[:, 1]]
        idx_maxsort = np.argsort(peak_values)[::-1]
        return coords[idx_maxsort][:npeaks]
<commit_msg>find_peaks threshold must be scalar (scikit-image) [skip ci]<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import numpy as np
from astropy.convolution import Kernel2D
from ..extern.imageutils import sigmaclip_stats

__all__ = ['detect_threshold', 'detect_sources', 'find_peaks']

def detect_threshold(data, snr, background=None, error=None, mask=None,
                     mask_val=None, sigclip_sigma=3.0, sigclip_iters=None):
    """
    Calculate a pixel-wise threshold image to be used to detect sources.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    snr : float
        The signal-to-noise ratio per pixel above the ``background`` for
        which to consider a pixel as possibly being part of a source.

    background : float or array_like, optional
        The background value(s) of the input ``data``.  ``background``
        may either be a scalar value or a 2D image with the same shape
        as the input ``data``.  If the input ``data`` has been
        background-subtracted, then set ``background`` to ``0.0``.  If
        `None`, then a scalar background value will be estimated using
        sigma-clipped statistics.

    error : float or array_like, optional
        The Gaussian 1-sigma standard deviation of the background noise
        in ``data``.  ``error`` should include all sources of
        "background" error, but *exclude* the Poisson error of the
        sources.  If ``error`` is a 2D image, then it should represent
        the 1-sigma background error in each pixel of ``data``.  If
        `None`, then a scalar background rms value will be estimated
        using sigma-clipped statistics.

    mask : array_like, bool, optional
        A boolean mask with the same shape as ``data``, where a `True`
        value indicates the corresponding element of ``data`` is masked.
        Masked pixels are ignored when computing the image background
        statistics.

    mask_val : float, optional
        An image data value (e.g., ``0.0``) that is ignored when
        computing the image background statistics.  ``mask_val`` will be
        ignored if ``mask`` is input.

    sigclip_sigma : float, optional
        The number of standard deviations to use as the clipping limit
        when calculating the image background statistics.

    sigclip_iters : float, optional
       The number of iterations to perform sigma clipping, or `None` to
       clip until convergence is achieved (i.e., continue until the last
       iteration clips nothing) when calculating the image background
       statistics.

    Returns
    -------
    threshold : 2D `~numpy.ndarray`
        A 2D image with the same shape as ``data`` containing the
        pixel-wise threshold values.

    See Also
    --------
    detect_sources

    Notes
    -----
    The ``mask``, ``mask_val``, ``sigclip_sigma``, and ``sigclip_iters``
    inputs are used only if it is necessary to estimate ``background``
    or ``error`` using sigma-clipped background statistics.  If
    ``background`` and ``error`` are both input, then ``mask``,
    ``mask_val``, ``sigclip_sigma``, and ``sigclip_iters`` are ignored.
    """

    if background is None or error is None:
        data_mean, data_median, data_std = sigmaclip_stats(
            data, mask=mask, mask_val=mask_val, sigma=sigclip_sigma,
            iters=sigclip_iters)
        bkgrd_image = np.zeros_like(data) + data_mean
        bkgrdrms_image = np.zeros_like(data) + data_std

    if background is None:
        background = bkgrd_image
    else:
        if np.isscalar(background):
            background = np.zeros_like(data) + background
        else:
            if background.shape != data.shape:
                raise ValueError('If input background is 2D, then it '
                                 'must have the same shape as the input '
                                 'data.')

    if error is None:
        error = bkgrdrms_image
    else:
        if np.isscalar(error):
            error = np.zeros_like(data) + error
        else:
            if error.shape != data.shape:
                raise ValueError('If input error is 2D, then it '
                                 'must have the same shape as the input '
                                 'data.')

    return background + (error * snr)

def detect_sources(data, threshold, npixels, filter_kernel=None,
                   connectivity=8):
    """
    Detect sources above a specified threshold value in an image and
    return a segmentation image.

    Detected sources must have ``npixels`` connected pixels that are
    each greater than the ``threshold`` value.  If the filtering option
    is used, then the ``threshold`` is applied to the filtered image.

    This function does not deblend overlapping sources.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    threshold : float or array-like
        The data value or pixel-wise data values to be used for the
        detection threshold.  A 2D ``threshold`` must have the same
        shape as ``data``.  See `detect_threshold` for one way to create
        a ``threshold`` image.

    npixels : int
        The number of connected pixels, each greater than ``threshold``,
        that an object must have to be detected.  ``npixels`` must be a
        positive integer.

    filter_kernel : array-like (2D) or `~astropy.convolution.Kernel2D`, optional
        The 2D array of the kernel used to filter the image before
        thresholding.  Filtering the image will smooth the noise and
        maximize detectability of objects with a shape similar to the
        kernel.

    connectivity : {4, 8}, optional
        The type of pixel connectivity used in determining how pixels
        are grouped into a detected source.  The options are 4 or 8
        (default).  4-connected pixels touch along their edges.
        8-connected pixels touch along their edges or corners.  For
        reference, SExtractor uses 8-connected pixels.

    Returns
    -------
    segment_image : `~numpy.ndarray` (int)
        A 2D segmentation image, with the same shape as ``data``, where
        sources are marked by different positive integer values.  A
        value of zero is reserved for the background.

    See Also
    --------
    detect_threshold, :class:`photutils.segmentation.segment_properties`

    Examples
    --------

    .. plot::
        :include-source:

        # make a table of Gaussian sources
        from astropy.table import Table
        table = Table()
        table['amplitude'] = [50, 70, 150, 210]
        table['x_mean'] = [160, 25, 150, 90]
        table['y_mean'] = [70, 40, 25, 60]
        table['x_stddev'] = [15.2, 5.1, 3., 8.1]
        table['y_stddev'] = [2.6, 2.5, 3., 4.7]
        table['theta'] = np.array([145., 20., 0., 60.]) * np.pi / 180.

        # make an image of the sources with Gaussian noise
        from photutils.datasets import make_gaussian_sources
        from photutils.datasets import make_noise_image
        shape = (100, 200)
        sources = make_gaussian_sources(shape, table)
        noise = make_noise_image(shape, type='gaussian', mean=0.,
                                 stddev=5., random_state=12345)
        image = sources + noise

        # detect the sources
        from photutils import detect_threshold, detect_sources
        threshold = detect_threshold(image, snr=3)
        from astropy.convolution import Gaussian2DKernel
        sigma = 3.0 / (2.0 * np.sqrt(2.0 * np.log(2.0)))   # FWHM = 3
        filter_kernel = Gaussian2DKernel(sigma, x_size=3, y_size=3)
        segm_image = detect_sources(image, threshold, npixels=5,
                                    filter_kernel=filter_kernel)

        # plot the image and the segmentation image
        import matplotlib.pyplot as plt
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))
        ax1.imshow(image, origin='lower', interpolation='nearest')
        ax2.imshow(segm_image, origin='lower', interpolation='nearest')
    """

    from scipy import ndimage

    if (npixels <= 0) or (int(npixels) != npixels):
        raise ValueError('npixels must be a positive integer, got '
                         '"{0}"'.format(npixels))

    conv_mode = 'constant'    # SExtractor mode
    conv_val = 0.0
    if filter_kernel is not None:
        if isinstance(filter_kernel, Kernel2D):
            image = ndimage.convolve(data, filter_kernel.array, mode=conv_mode,
                                     cval=conv_val)
        else:
            image = ndimage.convolve(data, filter_kernel, mode=conv_mode,
                                     cval=conv_val)
    else:
        image = data

    image = (image > threshold)
    if connectivity == 4:
        selem = ndimage.generate_binary_structure(2, 1)
    elif connectivity == 8:    # e.g., SExtractor
        selem = ndimage.generate_binary_structure(2, 2)
    objlabels, nobj = ndimage.label(image, structure=selem)
    objslices = ndimage.find_objects(objlabels)

    # remove objects with less than npixels
    for objslice in objslices:
        objlabel = objlabels[objslice]
        obj_npix = len(np.where(objlabel.ravel() != 0)[0])
        if obj_npix < npixels:
            objlabels[objslice] = 0

    # relabel to make sequential label indices
    objlabels, nobj = ndimage.label(objlabels, structure=selem)
    return objlabels

def find_peaks(data, threshold, min_separation=2, exclude_border=True,
               segment_image=None, npeaks=np.inf, footprint=None):
    """
    Find local peaks in an image that are above above a specified
    threshold value.

    Peaks are the local maxima above the ``threshold`` in a region of
    ``(2 * min_separation) + 1`` (i.e., peaks are separated by at least
    ``min_separation`` pixels).

    If peaks are flat (i.e., multiple adjacent pixels have identical
    intensities), then the coordinates of all such pixels are returned,
    even if they are not separated by at least ``min_separation``.

    Parameters
    ----------
    data : array_like
        The 2D array of the image.

    threshold : float
        The data value to be used for the detection threshold.

    min_separation : int, optional
        Minimum number of pixels separating peaks (i.e., peaks are
        separated by at least ``min_separation`` pixels).  To find the
        maximum number of peaks, use ``min_separation=1``.  If
        ``min_separation`` is not an integer, then it will be truncated.

    exclude_border : bool, optional
        If `True`, exclude peaks within ``min_separation`` from the
        border of the image as well as from each other.

    segment_image : `~numpy.ndarray` (int), optional
        If provided, then search for peaks located only within the
        labeled regions of a 2D segmentation image, where sources are
        marked by different positive integer values.  In the
        segmentation image a value of zero is reserved for the
        background.  ``segment_image`` must have the same shape as
        ``data``.

    npeaks : int, optional
        The maximum number of peaks to return.  When the number of
        detected peaks exceeds ``npeaks``, the peaks with the highest
        peak intensities will be returned.

    footprint : `~numpy.ndarray` of bools, optional
        A boolean array where `True` values describe the local footprint
        region within which to search for peaks at every point in
        ``data``.  Overrides ``min_separation``, except for border
        exclusion when ``exclude_border=True``.

    Returns
    -------
    output : `~numpy.ndarray`
        An ``Nx2`` array where the rows contain the ``(y, x)`` pixel
        coordinates of the local peaks.
    """

    if segment_image is not None:
        if segment_image.shape != data.shape:
            raise ValueError('segment_image and data must have the same '
                             'shape')

    from skimage.feature import peak_local_max
    coords = peak_local_max(data, min_distance=int(min_separation),
                            threshold_abs=threshold, threshold_rel=0.0,
                            exclude_border=exclude_border, indices=True,
                            num_peaks=npeaks, footprint=footprint,
                            labels=segment_image)
    if coords.shape[0] <= npeaks:
        return coords
    else:
        # NOTE: num_peaks is ignored by peak_local_max() if labels are input
        peak_values = data[coords[:, 0], coords[:, 1]]
        idx_maxsort = np.argsort(peak_values)[::-1]
        return coords[idx_maxsort][:npeaks]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

import os
import contextlib
import logging
import lzma
import tarfile
import shutil

from collections import Counter
from datetime import datetime
from textwrap import dedent
from urllib import urlretrieve

import click
import msgpack
import slugify

from bson import DBRef
from mongoengine import errors
from mongoengine.context_managers import switch_collection

from udata.commands import cli
from udata.core.dataset.models import Dataset
from udata.core.spatial import geoids
from udata.core.spatial.models import GeoLevel, GeoZone, SpatialCoverage
from udata.core.storages import logos, tmp

log = logging.getLogger(__name__)

DEFAULT_GEOZONES_FILE = 'https://github.com/etalab/geozones/releases/download/2019.0/geozones-countries-2019-0-msgpack.tar.xz'

def level_ref(level):
    return DBRef(GeoLevel._get_collection_name(), level)

@cli.group('spatial')
def grp():
    '''Geospatial related operations'''
    pass

def load_levels(col, path):
    with open(path) as fp:
        unpacker = msgpack.Unpacker(fp, encoding=str('utf-8'))
        for i, level in enumerate(unpacker, start=1):
            col.objects(id=level['id']).modify(
                upsert=True,
                set__name=level['label'],
                set__parents=[level_ref(p) for p in level['parents']],
                set__admin_level=level.get('admin_level')
            )
    return i

def load_zones(col, path):
    with open(path) as fp:
        unpacker = msgpack.Unpacker(fp, encoding=str('utf-8'))
        unpacker.next()  # Skip headers.
        for i, geozone in enumerate(unpacker):
            params = {
                'slug': slugify.slugify(geozone['name'], separator='-'),
                'level': geozone['level'],
                'code': geozone['code'],
                'name': geozone['name'],
                'keys': geozone.get('keys'),
                'parents': geozone.get('parents', []),
                'ancestors': geozone.get('ancestors', []),
                'successors': geozone.get('successors', []),
                'validity': geozone.get('validity'),
                'population': geozone.get('population'),
                'dbpedia': geozone.get('dbpedia'),
                'flag': geozone.get('flag'),
                'blazon': geozone.get('blazon'),
                'wikidata': geozone.get('wikidata'),
                'wikipedia': geozone.get('wikipedia'),
                'area': geozone.get('area'),
            }
            if geozone.get('geom') and (
                geozone['geom']['type'] != 'GeometryCollection' or
                    geozone['geom']['geometries']):
                params['geom'] = geozone['geom']
            try:
                col.objects(id=geozone['_id']).modify(upsert=True, **{
                    'set__{0}'.format(k): v for k, v in params.items()
                })
            except errors.ValidationError as e:
                log.warning('Validation error (%s) for %s with %s',
                            e, geozone['_id'], params)
                continue
    return i

@grp.command()
@click.argument('filename', metavar='<filename>', default=DEFAULT_GEOZONES_FILE)
@click.option('-d', '--drop', is_flag=True, help='Drop existing data')
def load(filename=DEFAULT_GEOZONES_FILE, drop=False):
    '''
    Load a geozones archive from <filename>

    <filename> can be either a local path or a remote URL.
    '''
    ts = datetime.now().isoformat().replace('-', '').replace(':', '').split('.')[0]
    if filename.startswith('http'):
        log.info('Downloading GeoZones bundle: %s', filename)
        filename, _ = urlretrieve(filename, tmp.path('geozones.tar.xz'))

    log.info('Extracting GeoZones bundle')
    with contextlib.closing(lzma.LZMAFile(filename)) as xz:
        with tarfile.open(fileobj=xz) as f:
            f.extractall(tmp.root)

    log.info('Loading GeoZones levels')

    log.info('Loading levels.msgpack')
    levels_filepath = tmp.path('levels.msgpack')
    if drop and GeoLevel.objects.count():
        name = '_'.join((GeoLevel._get_collection_name(), ts))
        target = GeoLevel._get_collection_name()
        with switch_collection(GeoLevel, name):
            total = load_levels(GeoLevel, levels_filepath)
            GeoLevel.objects._collection.rename(target, dropTarget=True)
    else:
        total = load_levels(GeoLevel, levels_filepath)
    os.remove(levels_filepath)
    log.info('Loaded {total} levels'.format(total=total))

    log.info('Loading zones.msgpack')
    zones_filepath = tmp.path('zones.msgpack')
    if drop and GeoZone.objects.count():
        name = '_'.join((GeoZone._get_collection_name(), ts))
        target = GeoZone._get_collection_name()
        with switch_collection(GeoZone, name):
            total = load_zones(GeoZone, zones_filepath)
            GeoZone.objects._collection.rename(target, dropTarget=True)
    else:
        total = load_zones(GeoZone, zones_filepath)
    os.remove(zones_filepath)
    log.info('Loaded {total} zones'.format(total=total))

    shutil.rmtree(tmp.path('translations'))  # Not in use for now.

def safe_tarinfo(tarinfo):
    '''make a tarinfo utf8-compatible'''
    tarinfo.name = tarinfo.name.decode('utf8')
    return tarinfo

@grp.command()
@click.argument('filename', metavar='<filename>')
def load_logos(filename):
    '''
    Load logos from a geologos archive from <filename>

    <filename> can be either a local path or a remote URL.
    '''
    if filename.startswith('http'):
        log.info('Downloading GeoLogos bundle: %s', filename)
        filename, _ = urlretrieve(filename, tmp.path('geologos.tar.xz'))

    log.info('Extracting GeoLogos bundle')
    with contextlib.closing(lzma.LZMAFile(filename)) as xz:
        with tarfile.open(fileobj=xz, encoding='utf8') as tar:
            decoded = (safe_tarinfo(t) for t in tar.getmembers())
            tar.extractall(tmp.root, members=decoded)

    log.info('Moving to the final location and cleaning up')
    if os.path.exists(logos.root):
        shutil.rmtree(logos.root)
    shutil.move(tmp.path('logos'), logos.root)
    log.info('Done')

@grp.command()
def migrate():
    '''
    Migrate zones from old to new ids in datasets.

    Should only be run once with the new version of geozones w/ geohisto.
    '''
    counter = Counter(['zones', 'datasets'])
    qs = GeoZone.objects.only('id', 'level', 'successors')
    # Iter over datasets with zones
    for dataset in Dataset.objects(spatial__zones__gt=[]):
        counter['datasets'] += 1
        new_zones = []
        for current_zone in dataset.spatial.zones:
            counter['zones'] += 1
            level, code, validity = geoids.parse(current_zone.id)
            zone = qs(level=level, code=code).valid_at(validity).first()
            if not zone:
                log.warning('No match for %s: skipped', current_zone.id)
                counter['skipped'] += 1
                continue
            previous = None
            while not zone.is_current and len(zone.successors) == 1 and zone.id != previous:
                previous = zone.id
                zone = qs(id=zone.successors[0]).first() or zone
            new_zones.append(zone.id)
            counter[zone.level] += 1
        dataset.update(
            spatial=SpatialCoverage(
                granularity=dataset.spatial.granularity,
                zones=list(new_zones)
            )
        )
    level_summary = '\n'.join([
        ' - {0}: {1}'.format(l.id, counter[l.id])
        for l in GeoLevel.objects.order_by('admin_level')
    ])
    summary = '\n'.join([dedent('''\
    Summary
    =======
    Processed {zones} zones in {datasets} datasets:\
    '''.format(level_summary, **counter)), level_summary])
    log.info(summary)
    log.info('Done')
<commit_msg>Cleanup after spatial load command<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

import os
import contextlib
import logging
import lzma
import tarfile
import shutil

from collections import Counter
from datetime import datetime
from textwrap import dedent
from urllib import urlretrieve

import click
import msgpack
import slugify

from bson import DBRef
from mongoengine import errors
from mongoengine.context_managers import switch_collection

from udata.commands import cli
from udata.core.dataset.models import Dataset
from udata.core.spatial import geoids
from udata.core.spatial.models import GeoLevel, GeoZone, SpatialCoverage
from udata.core.storages import logos, tmp

log = logging.getLogger(__name__)

DEFAULT_GEOZONES_FILE = 'https://github.com/etalab/geozones/releases/download/2019.0/geozones-countries-2019-0-msgpack.tar.xz'
GEOZONE_FILENAME = 'geozones.tar.xz'

def level_ref(level):
    return DBRef(GeoLevel._get_collection_name(), level)

@cli.group('spatial')
def grp():
    '''Geospatial related operations'''
    pass

def load_levels(col, path):
    with open(path) as fp:
        unpacker = msgpack.Unpacker(fp, encoding=str('utf-8'))
        for i, level in enumerate(unpacker, start=1):
            col.objects(id=level['id']).modify(
                upsert=True,
                set__name=level['label'],
                set__parents=[level_ref(p) for p in level['parents']],
                set__admin_level=level.get('admin_level')
            )
    return i

def load_zones(col, path):
    with open(path) as fp:
        unpacker = msgpack.Unpacker(fp, encoding=str('utf-8'))
        unpacker.next()  # Skip headers.
        for i, geozone in enumerate(unpacker):
            params = {
                'slug': slugify.slugify(geozone['name'], separator='-'),
                'level': geozone['level'],
                'code': geozone['code'],
                'name': geozone['name'],
                'keys': geozone.get('keys'),
                'parents': geozone.get('parents', []),
                'ancestors': geozone.get('ancestors', []),
                'successors': geozone.get('successors', []),
                'validity': geozone.get('validity'),
                'population': geozone.get('population'),
                'dbpedia': geozone.get('dbpedia'),
                'flag': geozone.get('flag'),
                'blazon': geozone.get('blazon'),
                'wikidata': geozone.get('wikidata'),
                'wikipedia': geozone.get('wikipedia'),
                'area': geozone.get('area'),
            }
            if geozone.get('geom') and (
                geozone['geom']['type'] != 'GeometryCollection' or
                    geozone['geom']['geometries']):
                params['geom'] = geozone['geom']
            try:
                col.objects(id=geozone['_id']).modify(upsert=True, **{
                    'set__{0}'.format(k): v for k, v in params.items()
                })
            except errors.ValidationError as e:
                log.warning('Validation error (%s) for %s with %s',
                            e, geozone['_id'], params)
                continue
    return i

@grp.command()
@click.argument('filename', metavar='<filename>', default=DEFAULT_GEOZONES_FILE)
@click.option('-d', '--drop', is_flag=True, help='Drop existing data')
def load(filename=DEFAULT_GEOZONES_FILE, drop=False):
    '''
    Load a geozones archive from <filename>

    <filename> can be either a local path or a remote URL.
    '''
    ts = datetime.now().isoformat().replace('-', '').replace(':', '').split('.')[0]
    prefix = 'geozones-{0}'.format(ts)
    if filename.startswith('http'):
        log.info('Downloading GeoZones bundle: %s', filename)
        filename, _ = urlretrieve(filename, tmp.path(GEOZONE_FILENAME))

    log.info('Extracting GeoZones bundle')
    with contextlib.closing(lzma.LZMAFile(filename)) as xz:
        with tarfile.open(fileobj=xz) as f:
            f.extractall(tmp.path(prefix))

    log.info('Loading GeoZones levels')

    log.info('Loading levels.msgpack')
    levels_filepath = tmp.path(prefix + '/levels.msgpack')
    if drop and GeoLevel.objects.count():
        name = '_'.join((GeoLevel._get_collection_name(), ts))
        target = GeoLevel._get_collection_name()
        with switch_collection(GeoLevel, name):
            total = load_levels(GeoLevel, levels_filepath)
            GeoLevel.objects._collection.rename(target, dropTarget=True)
    else:
        total = load_levels(GeoLevel, levels_filepath)
    log.info('Loaded {total} levels'.format(total=total))

    log.info('Loading zones.msgpack')
    zones_filepath = tmp.path(prefix + '/zones.msgpack')
    if drop and GeoZone.objects.count():
        name = '_'.join((GeoZone._get_collection_name(), ts))
        target = GeoZone._get_collection_name()
        with switch_collection(GeoZone, name):
            total = load_zones(GeoZone, zones_filepath)
            GeoZone.objects._collection.rename(target, dropTarget=True)
    else:
        total = load_zones(GeoZone, zones_filepath)
    log.info('Loaded {total} zones'.format(total=total))

    tmp.delete(prefix)
    if tmp.exists(GEOZONE_FILENAME):  # Has been downloaded
        tmp.delete(GEOZONE_FILENAME)

def safe_tarinfo(tarinfo):
    '''make a tarinfo utf8-compatible'''
    tarinfo.name = tarinfo.name.decode('utf8')
    return tarinfo

@grp.command()
@click.argument('filename', metavar='<filename>')
def load_logos(filename):
    '''
    Load logos from a geologos archive from <filename>

    <filename> can be either a local path or a remote URL.
    '''
    if filename.startswith('http'):
        log.info('Downloading GeoLogos bundle: %s', filename)
        filename, _ = urlretrieve(filename, tmp.path('geologos.tar.xz'))

    log.info('Extracting GeoLogos bundle')
    with contextlib.closing(lzma.LZMAFile(filename)) as xz:
        with tarfile.open(fileobj=xz, encoding='utf8') as tar:
            decoded = (safe_tarinfo(t) for t in tar.getmembers())
            tar.extractall(tmp.root, members=decoded)

    log.info('Moving to the final location and cleaning up')
    if os.path.exists(logos.root):
        shutil.rmtree(logos.root)
    shutil.move(tmp.path('logos'), logos.root)
    log.info('Done')

@grp.command()
def migrate():
    '''
    Migrate zones from old to new ids in datasets.

    Should only be run once with the new version of geozones w/ geohisto.
    '''
    counter = Counter(['zones', 'datasets'])
    qs = GeoZone.objects.only('id', 'level', 'successors')
    # Iter over datasets with zones
    for dataset in Dataset.objects(spatial__zones__gt=[]):
        counter['datasets'] += 1
        new_zones = []
        for current_zone in dataset.spatial.zones:
            counter['zones'] += 1
            level, code, validity = geoids.parse(current_zone.id)
            zone = qs(level=level, code=code).valid_at(validity).first()
            if not zone:
                log.warning('No match for %s: skipped', current_zone.id)
                counter['skipped'] += 1
                continue
            previous = None
            while not zone.is_current and len(zone.successors) == 1 and zone.id != previous:
                previous = zone.id
                zone = qs(id=zone.successors[0]).first() or zone
            new_zones.append(zone.id)
            counter[zone.level] += 1
        dataset.update(
            spatial=SpatialCoverage(
                granularity=dataset.spatial.granularity,
                zones=list(new_zones)
            )
        )
    level_summary = '\n'.join([
        ' - {0}: {1}'.format(l.id, counter[l.id])
        for l in GeoLevel.objects.order_by('admin_level')
    ])
    summary = '\n'.join([dedent('''\
    Summary
    =======
    Processed {zones} zones in {datasets} datasets:\
    '''.format(level_summary, **counter)), level_summary])
    log.info(summary)
    log.info('Done')
<|endoftext|>


<commit_before>from MySQLdb.constants import FIELD_TYPE
from MySQLdb.converters import conversions, Bool2Str
from MySQLdb import times, _mysql
import datetime
import inspect

CONVERSIONS = conversions

def _bytes_to_utf8(b):
    return b.decode('utf-8')

CONVERSIONS[FIELD_TYPE.STRING] = _bytes_to_utf8
CONVERSIONS[FIELD_TYPE.VAR_STRING] = _bytes_to_utf8
CONVERSIONS[FIELD_TYPE.VARCHAR] = _bytes_to_utf8
CONVERSIONS[FIELD_TYPE.BLOB] = _bytes_to_utf8

def _escape_bytes(b, c):
    return _mysql.string_literal(b, c).decode('utf-8')

CONVERSIONS[bytes] = _escape_bytes

def _escape_string(s, d):
    return _mysql.string_literal(s.encode('utf-8')).decode('utf-8')

CONVERSIONS[str] = _escape_string

def _escape_datetime(dt, c):
    return times.DateTime2literal(dt, c).decode('utf-8')
def _escape_timedelta(dt, c):
    return times.DateTimeDelta2literal(dt, c).decode('utf-8')

CONVERSIONS[datetime.datetime] = _escape_datetime
CONVERSIONS[datetime.timedelta] = _escape_timedelta

def _escape_bool(b, d):
    return Bool2Str(b, d).decode('utf-8')
CONVERSIONS[bool] = _escape_bool
<commit_msg>Use binary flag to determine if value should be bytes or utf8 (#24)<commit_after>from MySQLdb.constants import FIELD_TYPE, FLAG
from MySQLdb.converters import conversions, Bool2Str
from MySQLdb import times, _mysql
import datetime
import inspect

CONVERSIONS = conversions

def _bytes_to_utf8(b):
    return b.decode('utf-8')

CONVERSIONS[FIELD_TYPE.STRING] = ((FLAG.BINARY, bytes), (None, _bytes_to_utf8))
CONVERSIONS[FIELD_TYPE.VAR_STRING] = ((FLAG.BINARY, bytes), (None, _bytes_to_utf8))
CONVERSIONS[FIELD_TYPE.VARCHAR] = ((FLAG.BINARY, bytes), (None, _bytes_to_utf8))
CONVERSIONS[FIELD_TYPE.BLOB] = ((FLAG.BINARY, bytes), (None, _bytes_to_utf8))

def _escape_bytes(b, c):
    return _mysql.string_literal(b, c).decode('utf-8')

CONVERSIONS[bytes] = _escape_bytes

def _escape_string(s, d):
    return _mysql.string_literal(s.encode('utf-8')).decode('utf-8')

CONVERSIONS[str] = _escape_string

def _escape_datetime(dt, c):
    return times.DateTime2literal(dt, c).decode('utf-8')
def _escape_timedelta(dt, c):
    return times.DateTimeDelta2literal(dt, c).decode('utf-8')

CONVERSIONS[datetime.datetime] = _escape_datetime
CONVERSIONS[datetime.timedelta] = _escape_timedelta

def _escape_bool(b, d):
    return Bool2Str(b, d).decode('utf-8')
CONVERSIONS[bool] = _escape_bool
<|endoftext|>


<commit_before># stdlib
import sys
import time
from datetime import datetime, timedelta
import logging
from operator import attrgetter

# vendor
import tornado.ioloop

# project
from checks.check_status import ForwarderStatus

log = logging.getLogger(__name__)

def plural(count):
    if count > 1:
        return "s"
    return ""

class ImplementationError(Exception): pass

class Transaction(object):

    def __init__(self):

        self._id = None
        self._error_count = 0
        self._next_flush = datetime.now()
        self._size = None

    def get_id(self):
        return self._id

    def set_id(self, new_id):
        assert self._id is None
        self._id = new_id

    def inc_error_count(self):
        self._error_count = self._error_count + 1

    def get_error_count(self):
        return self._error_count

    def get_size(self):
        if self._size is None:
            self._size = sys.getsizeof(self)

        return self._size

    def get_next_flush(self):
        return self._next_flush

    def compute_next_flush(self,max_delay):
        # Transactions are replayed, try to send them faster for newer transactions
        # Send them every MAX_WAIT_FOR_REPLAY at most
        td = timedelta(seconds=self._error_count * 20)
        if td > max_delay:
            td = max_delay

        newdate = datetime.now() + td
        self._next_flush = newdate.replace(microsecond=0)

    def time_to_flush(self,now = datetime.now()):
        return self._next_flush < now

    def flush(self):
        raise ImplementationError("To be implemented in a subclass")

class TransactionManager(object):
    """Holds any transaction derived object list and make sure they
       are all commited, without exceeding parameters (throttling, memory consumption) """

    def __init__(self, max_wait_for_replay, max_queue_size, throttling_delay):
        self._MAX_WAIT_FOR_REPLAY = max_wait_for_replay
        self._MAX_QUEUE_SIZE = max_queue_size
        self._THROTTLING_DELAY = throttling_delay

        self._flush_without_ioloop = False # useful for tests

        self._transactions = [] #List of all non commited transactions
        self._total_count = 0 # Maintain size/count not to recompute it everytime
        self._total_size = 0
        self._flush_count = 0

        # Global counter to assign a number to each transaction: we may have an issue
        #  if this overlaps
        self._counter = 0

        self._trs_to_flush = None # Current transactions being flushed
        self._last_flush = datetime.now() # Last flush (for throttling)

        # Track an initial status message.
        ForwarderStatus().persist()

    def get_transactions(self):
        return self._transactions

    def print_queue_stats(self):
        log.debug("Queue size: at %s, %s transaction(s), %s KB" %
            (time.time(), self._total_count, (self._total_size/1024)))

    def get_tr_id(self):
        self._counter =  self._counter + 1
        return self._counter

    def append(self,tr):

        # Give the transaction an id
        tr.set_id(self.get_tr_id())

        # Check the size
        tr_size = tr.get_size()

        log.debug("New transaction to add, total size of queue would be: %s KB" %
            ((self._total_size + tr_size)/ 1024))

        if (self._total_size + tr_size) > self._MAX_QUEUE_SIZE:
            log.warn("Queue is too big, removing old transactions...")
            new_trs = sorted(self._transactions,key=attrgetter('_next_flush'), reverse = True)
            for tr2 in new_trs:
                if (self._total_size + tr_size) > self._MAX_QUEUE_SIZE:
                    self._transactions.remove(tr2)
                    self._total_count = self._total_count - 1
                    self._total_size = self._total_size - tr2.get_size()
                    log.warn("Removed transaction %s from queue" % tr2.get_id())

        # Done
        self._transactions.append(tr)
        self._total_count = self._total_count + 1
        self._total_size = self._total_size + tr_size

        log.debug("Transaction %s added" % (tr.get_id()))
        self.print_queue_stats()

    def flush(self):

        if self._trs_to_flush is not None:
            log.debug("A flush is already in progress, not doing anything")
            return

        to_flush = []
        # Do we have something to do ?
        now = datetime.now()
        for tr in self._transactions:
            if tr.time_to_flush(now):
                to_flush.append(tr)

        count = len(to_flush)
        if count > 0:
            log.debug("Flushing %s transaction%s" % (count,plural(count)))
            self._trs_to_flush = to_flush
            self.flush_next()
        self._flush_count += 1

        ForwarderStatus(
            queue_length=self._total_count,
            queue_size=self._total_size,
            flush_count=self._flush_count).persist()

    def flush_next(self):

        if len(self._trs_to_flush) > 0:

            td = self._last_flush + self._THROTTLING_DELAY - datetime.now()
            # Python 2.7 has this built in, python < 2.7 don't...
            if hasattr(td,'total_seconds'):
                delay = td.total_seconds()
            else:
                delay = (td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6) / 10.0**6

            if delay <= 0:
                tr = self._trs_to_flush.pop()
                self._last_flush = datetime.now()
                log.debug("Flushing transaction %d" % tr.get_id())
                try:
                    tr.flush()
                except Exception,e :
                    log.exception(e)
                    self.tr_error(tr)
                    self.flush_next()
            else:
                # Wait a little bit more
                if  tornado.ioloop.IOLoop.instance().running():
                    tornado.ioloop.IOLoop.instance().add_timeout(time.time() + delay,
                        lambda: self.flush_next())
                elif self._flush_without_ioloop:
                    # Tornado is no started (ie, unittests), do it manually: BLOCKING
                    time.sleep(delay)
                    self.flush_next()
        else:
            self._trs_to_flush = None

    def tr_error(self,tr):
        tr.inc_error_count()
        tr.compute_next_flush(self._MAX_WAIT_FOR_REPLAY)
        log.warn("Transaction %d in error (%s error%s), it will be replayed after %s" %
          (tr.get_id(), tr.get_error_count(), plural(tr.get_error_count()),
           tr.get_next_flush()))

    def tr_success(self,tr):
        log.debug("Transaction %d completed" % tr.get_id())
        self._transactions.remove(tr)
        self._total_count = self._total_count - 1
        self._total_size = self._total_size - tr.get_size()
        self.print_queue_stats()

<commit_msg>Fix for tornado 3.x<commit_after># stdlib
import sys
import time
from datetime import datetime, timedelta
import logging
from operator import attrgetter

# vendor
import tornado.ioloop

# project
from checks.check_status import ForwarderStatus

log = logging.getLogger(__name__)

def plural(count):
    if count > 1:
        return "s"
    return ""

class ImplementationError(Exception): pass

class Transaction(object):

    def __init__(self):

        self._id = None
        self._error_count = 0
        self._next_flush = datetime.now()
        self._size = None

    def get_id(self):
        return self._id

    def set_id(self, new_id):
        assert self._id is None
        self._id = new_id

    def inc_error_count(self):
        self._error_count = self._error_count + 1

    def get_error_count(self):
        return self._error_count

    def get_size(self):
        if self._size is None:
            self._size = sys.getsizeof(self)

        return self._size

    def get_next_flush(self):
        return self._next_flush

    def compute_next_flush(self,max_delay):
        # Transactions are replayed, try to send them faster for newer transactions
        # Send them every MAX_WAIT_FOR_REPLAY at most
        td = timedelta(seconds=self._error_count * 20)
        if td > max_delay:
            td = max_delay

        newdate = datetime.now() + td
        self._next_flush = newdate.replace(microsecond=0)

    def time_to_flush(self,now = datetime.now()):
        return self._next_flush < now

    def flush(self):
        raise ImplementationError("To be implemented in a subclass")

class TransactionManager(object):
    """Holds any transaction derived object list and make sure they
       are all commited, without exceeding parameters (throttling, memory consumption) """

    def __init__(self, max_wait_for_replay, max_queue_size, throttling_delay):
        self._MAX_WAIT_FOR_REPLAY = max_wait_for_replay
        self._MAX_QUEUE_SIZE = max_queue_size
        self._THROTTLING_DELAY = throttling_delay

        self._flush_without_ioloop = False # useful for tests

        self._transactions = [] #List of all non commited transactions
        self._total_count = 0 # Maintain size/count not to recompute it everytime
        self._total_size = 0
        self._flush_count = 0

        # Global counter to assign a number to each transaction: we may have an issue
        #  if this overlaps
        self._counter = 0

        self._trs_to_flush = None # Current transactions being flushed
        self._last_flush = datetime.now() # Last flush (for throttling)

        # Track an initial status message.
        ForwarderStatus().persist()

    def get_transactions(self):
        return self._transactions

    def print_queue_stats(self):
        log.debug("Queue size: at %s, %s transaction(s), %s KB" %
            (time.time(), self._total_count, (self._total_size/1024)))

    def get_tr_id(self):
        self._counter =  self._counter + 1
        return self._counter

    def append(self,tr):

        # Give the transaction an id
        tr.set_id(self.get_tr_id())

        # Check the size
        tr_size = tr.get_size()

        log.debug("New transaction to add, total size of queue would be: %s KB" %
            ((self._total_size + tr_size)/ 1024))

        if (self._total_size + tr_size) > self._MAX_QUEUE_SIZE:
            log.warn("Queue is too big, removing old transactions...")
            new_trs = sorted(self._transactions,key=attrgetter('_next_flush'), reverse = True)
            for tr2 in new_trs:
                if (self._total_size + tr_size) > self._MAX_QUEUE_SIZE:
                    self._transactions.remove(tr2)
                    self._total_count = self._total_count - 1
                    self._total_size = self._total_size - tr2.get_size()
                    log.warn("Removed transaction %s from queue" % tr2.get_id())

        # Done
        self._transactions.append(tr)
        self._total_count = self._total_count + 1
        self._total_size = self._total_size + tr_size

        log.debug("Transaction %s added" % (tr.get_id()))
        self.print_queue_stats()

    def flush(self):

        if self._trs_to_flush is not None:
            log.debug("A flush is already in progress, not doing anything")
            return

        to_flush = []
        # Do we have something to do ?
        now = datetime.now()
        for tr in self._transactions:
            if tr.time_to_flush(now):
                to_flush.append(tr)

        count = len(to_flush)
        if count > 0:
            log.debug("Flushing %s transaction%s" % (count,plural(count)))
            self._trs_to_flush = to_flush
            self.flush_next()
        self._flush_count += 1

        ForwarderStatus(
            queue_length=self._total_count,
            queue_size=self._total_size,
            flush_count=self._flush_count).persist()

    def flush_next(self):

        if len(self._trs_to_flush) > 0:

            td = self._last_flush + self._THROTTLING_DELAY - datetime.now()
            # Python 2.7 has this built in, python < 2.7 don't...
            if hasattr(td,'total_seconds'):
                delay = td.total_seconds()
            else:
                delay = (td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6) / 10.0**6

            if delay <= 0:
                tr = self._trs_to_flush.pop()
                self._last_flush = datetime.now()
                log.debug("Flushing transaction %d" % tr.get_id())
                try:
                    tr.flush()
                except Exception,e :
                    log.exception(e)
                    self.tr_error(tr)
                    self.flush_next()
            else:
                # Wait a little bit more
                if  tornado.ioloop.IOLoop.instance()._running:
                    tornado.ioloop.IOLoop.instance().add_timeout(time.time() + delay,
                        lambda: self.flush_next())
                elif self._flush_without_ioloop:
                    # Tornado is no started (ie, unittests), do it manually: BLOCKING
                    time.sleep(delay)
                    self.flush_next()
        else:
            self._trs_to_flush = None

    def tr_error(self,tr):
        tr.inc_error_count()
        tr.compute_next_flush(self._MAX_WAIT_FOR_REPLAY)
        log.warn("Transaction %d in error (%s error%s), it will be replayed after %s" %
          (tr.get_id(), tr.get_error_count(), plural(tr.get_error_count()),
           tr.get_next_flush()))

    def tr_success(self,tr):
        log.debug("Transaction %d completed" % tr.get_id())
        self._transactions.remove(tr)
        self._total_count = self._total_count - 1
        self._total_size = self._total_size - tr.get_size()
        self.print_queue_stats()

<|endoftext|>


<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import, print_function

import blessings
import prettytable

from .. import cli
from ..master import CURRENT as MASTER

try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

parser = cli.parser(
    description="frameworks list"
)

parser.add_argument(
    "-b", "--task-breakdown", action="store_true",
    help="show task breakdown per framework"
)

parser.add_argument(
    "-i", "--inactive", action="store_true",
    help="show inactive frameworks as well"
)

def task_stats(tasks, state):
    count = 0
    for task in tasks:
        if task['state'] == state:
            count += 1
    return count

@cli.init(parser)
def main(args):
    term = blessings.Terminal()

    table_generator = OrderedDict([
        ("ID", lambda x: x["id"]),
        ("name", lambda x: x["name"]),
        ("host", lambda x: x["hostname"]),
        ("active", lambda x: x["active"]),
        ("active tasks", lambda x: len(x["tasks"])),
    ])

    if args.task_breakdown:
        # active tasks
        table_generator['staging'] = lambda x: task_stats(x["tasks"], "TASK_STAGING")
        table_generator['starting'] = lambda x: task_stats(x["tasks"], "TASK_STARTING")
        table_generator['running'] = lambda x: task_stats(x["tasks"], "TASK_RUNNING")
        # terminal state tasks. The count here is approximate since the API caps the result
        # at a max of 1000 tasks
        table_generator['failed'] = lambda x: task_stats(x["completed_tasks"], "TASK_FAILED")
        table_generator['killed'] = lambda x: task_stats(x["completed_tasks"], "TASK_FAILED")
        table_generator['error'] = lambda x: task_stats(x["completed_tasks"], "TASK_ERROR")
        table_generator['lost'] = lambda x: task_stats(x["completed_tasks"], "TASK_ERROR")
        table_generator['finished'] = lambda x: task_stats(x["completed_tasks"], "TASK_FINISHED")

    tb = prettytable.PrettyTable(
        [x.upper() for x in table_generator.keys()],
        border=False,
        max_table_width=term.width,
        hrules=prettytable.NONE,
        vrules=prettytable.NONE,
        left_padding_width=0,
        right_padding_width=1
    )

    def format_row(framework):
        return [fn(framework) for fn in table_generator.values()]

    for framework in MASTER.state['frameworks']:
        if args.inactive or framework['active']:
            tb.add_row(format_row(framework))

    if tb.rowcount == 0:
        print("===>{0}You have no frameworks{1}<===".format(
            term.red, term.white))
    else:
        print(tb)
<commit_msg>made the task column a bit cleaner<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import, print_function

import blessings
import prettytable

from .. import cli
from ..master import CURRENT as MASTER

try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

parser = cli.parser(
    description="frameworks list"
)

parser.add_argument(
    "-b", "--task-breakdown", action="store_true",
    help="show task breakdown per framework"
)

parser.add_argument(
    "-i", "--inactive", action="store_true",
    help="show inactive frameworks as well"
)

def task_stats(tasks, state):
    count = 0
    for task in tasks:
        if task['state'] == state:
            count += 1
    return count

@cli.init(parser)
def main(args):
    term = blessings.Terminal()

    table_generator = OrderedDict([
        ("ID", lambda x: x["id"]),
        ("name", lambda x: x["name"]),
        ("host", lambda x: x["hostname"]),
        ("active", lambda x: x["active"]),
    ])

    # TODO: the task breakdown is not done in an efficient way (since we traverse the
    # task list on every framework). This can be optimized to do the traversal only
    # once per framework.
    if args.task_breakdown:
        # active tasks
        table_generator['staging'] = lambda x: task_stats(x["tasks"], "TASK_STAGING")
        table_generator['starting'] = lambda x: task_stats(x["tasks"], "TASK_STARTING")
        table_generator['running'] = lambda x: task_stats(x["tasks"], "TASK_RUNNING")
        # terminal state tasks. The count here is approximate since the API caps the result
        # at a max of 1000 tasks
        table_generator['failed'] = lambda x: task_stats(x["completed_tasks"], "TASK_FAILED")
        table_generator['killed'] = lambda x: task_stats(x["completed_tasks"], "TASK_FAILED")
        table_generator['error'] = lambda x: task_stats(x["completed_tasks"], "TASK_ERROR")
        table_generator['lost'] = lambda x: task_stats(x["completed_tasks"], "TASK_ERROR")
        table_generator['finished'] = lambda x: task_stats(x["completed_tasks"], "TASK_FINISHED")
    else:
        table_generator['active tasks'] = lambda x: len(x["tasks"])

    tb = prettytable.PrettyTable(
        [x.upper() for x in table_generator.keys()],
        border=False,
        max_table_width=term.width,
        hrules=prettytable.NONE,
        vrules=prettytable.NONE,
        left_padding_width=0,
        right_padding_width=1
    )

    def format_row(framework):
        return [fn(framework) for fn in table_generator.values()]

    for framework in MASTER.state['frameworks']:
        if args.inactive or framework['active']:
            tb.add_row(format_row(framework))

    if tb.rowcount == 0:
        # TODO: can this be possible ?
        print("===>{0}You have no frameworks{1}<===".format(
            term.red, term.white))
    else:
        print(tb)
<|endoftext|>


<commit_before>"""
Prepare a python program (tiny subset) for compliation into sed.
Transform into positive form:
- all operators and binary comparisons are replaced with call to functions
  x + y --> signed_add(x, y), idem -, *, //, ==, !=, <, <=, >, >=
- all operands of operators and binary comparisons inside signed_xxx functions
  must be positive integers.
- two functions handling negative values remain: is_positive and negative. They
  are treated separately when testing python positive form or compiling into sed.
- augmented assignments are replaced with simple assignments.

A testing mode of the transformer program generates code testing that arguments
of operators and comparisons are positive.
"""

from __future__ import division

import sys
import inspect
import ast
import operator
import codegen

# -- Transformer --------------------------------------------------------------

signed_func = {
    ast.Add: 'signed_add',
    ast.Sub: 'signed_sub',
    ast.Mult: 'signed_mult',
    ast.Div: None,
    ast.FloorDiv: 'signed_div',
    ast.Eq: 'signed_eq',
    ast.NotEq: 'signed_noteq',
    ast.Lt: 'signed_lt',
    ast.LtE: 'signed_lte',
    ast.Gt: 'signed_gt',
    ast.GtE: 'signed_gte'}

class NumsedAstTransformer(ast.NodeTransformer):

    def __init__(self, func):
        self.func = func
        self.required_func = set()

    def make_call(self, func, args):
        self.required_func.add(func)
        return ast.Call(func=ast.Name(id=func, ctx=ast.Load()),
                        args=args,
                        keywords=[], starargs=None, kwargs=None)

    def visit_BinOp(self, node):
        self.generic_visit(node)
        return self.make_call(self.func[type(node.op)], [node.left, node.right])

    def visit_Compare(self, node):
        self.generic_visit(node)
        left = node.left
        ops = node.ops
        comparators = node.comparators
        if len(ops) == 1:
            return self.make_call(self.func[type(node.ops[0])], [node.left, comparators[0]])
        else:
            list_compare = []
            while ops:
                compare = self.make_call(self.func[type(ops[0])], [left, comparators[0]])
                list_compare.append(compare)
                left = comparators[0]
                ops = ops[1:]
                comparators = comparators[1:]
            return ast.BoolOp(op=ast.And(), values=list_compare)

    def visit_AugAssign(self, node):
        # AugAssign(target=Name(id='x', ctx=Store()), op=Add(), value=Num(n=1)),
        # Assign(targets=[Name(id='x', ctx=Store())], value=BinOp(left=Name(id='x', ctx=Load()), op=Add(), right=Num(n=1)))])
        #self.generic_visit(node)
        target = node.target
        source = ast.Name(id=target.id, ctx=ast.Load())
        return ast.Assign(targets=[target], value=self.visit_BinOp(ast.BinOp(left=source, op=node.op, right=node.value)))

# -- Builtin functions -------------------------------------------------------

def signed_eq(x, y):
    if is_positive(x):
        if is_positive(y):
            return x == y
        else:
            return False
    else:
        if is_positive(y):
            return False
        else:
            return negative(x) == negative(y)

def signed_noteq(x, y):
    return not signed_eq(x, y)

def signed_lt(x, y):
    if is_positive(x):
        if is_positive(y):
            return x < y
        else:
            return True
    else:
        if is_positive(y):
            return True
        else:
            return negative(x) < negative(y)

def signed_gte(x, y):
    return not signed_lt(x, y)

def signed_add(x, y):
    if is_positive(x):
        if is_positive(y):
            r = x + y
        else:
            y = negative(y)
            if x > y:
                r = x - y
            else:
                r = negative(y - x)
    else:
        x = negative(x)
        if is_positive(y):
            if x > y:
                r = negative(x - y)
            else:
                r = y - x
        else:
            y = negative(y)
            r = negative(x + y)
    return r

def signed_sub(x, y):
    if is_positive(x):
        if is_positive(y):
            if x > y:
                return x - y
            else:
                return negative(y - x)
        else:
            return x + negative(y)
    else:
        abs_x = negative(x)
        if is_positive(y):
            return negative(abs_x + y)
        else:
            abs_y = negative(y)
            if abs_x > abs_y:
                return negative(abs_x - abs_y)
            else:
                return abs_y - abs_x

def signed_pos(x):
    return x

def signed_neg(x):
    return negative(x)

def euclide(a, b):
    # http://compoasso.free.fr/primelistweb/page/prime/euclide.php
    r = a
    q = 0
    n = 0
    aux = b

    while aux <= a:
        aux = aux * 2
        n += 1

    while n > 0:
        #aux = aux / 2
        aux *= 5
        aux /= 10
        n -= 1
        q = q * 2
        if r >= aux:
            r -= aux
            q += 1

    return q

def modulo(a, b):
    q = euclide(a, b)
    return a - b * q

def signed_mult(a, b):
    if (a >= 0 and b >= 0) or (a <= 0 and b <= 0):
        return a * b
    else:
        return negative(a * b)

def signed_div(a, b):
    if (a >= 0 and b >= 0) or (a <= 0 and b <= 0):
        return euclide(a, b)
    else:
        return negative(euclide(a, b))

def signed_mod(a, b):
    q = signed_divide(a, b)
    return a - b * q

# -- Primitives --------------------------------------------------------------

def is_positive(x):
    return operator.ge(x, 0)

def negative(x):
    return operator.neg(x)

# -- Testing transformation --------------------------------------------------

def unsigned_gte(x, y):
    assert x >= 0
    assert y >= 0
    return x > y

def unsigned_add(x, y):
    assert x >= 0
    assert y >= 0
    return x + y

# ... or rather visitor adds assertions before each operator and comparison

# -- Tests -------------------------------------------------------------------

class NumsedAstVisitor(ast.NodeVisitor):
    def generic_visit(self, node):
        print type(node).__name__
        ast.NodeVisitor.generic_visit(self, node)

    #def visit_BinOp(self, node):
    #    print ast.dump(node)

# -- Main --------------------------------------------------------------------

def transform(script_in, script_out):
    test = False

    tree = ast.parse(open(script_in).read())
    if test:
        print ast.dump(tree)
        exec(compile(tree, filename="<ast>", mode="exec"))
        print '--'

    numsed_ast_transformer = NumsedAstTransformer(signed_func)
    numsed_ast_transformer.visit(tree)

    if test:
        ast.fix_missing_locations(tree)
        print ast.dump(tree)
        # going to and from AST
        exec(compile(tree, filename="<ast>", mode="exec"))

    builtin = numsed_ast_transformer.required_func
    builtin = [globals()[x] for x in builtin]
    if signed_div in builtin or signed_mod in builtin:
        builtin.append(euclide)
    builtin += [is_positive, negative]
    #print builtin

    # add builtin functions to code to compile
    script = ''
    for func in builtin:
        script += '\n'
        script += ''.join(inspect.getsourcelines(func)[0])
    script += '\n'
    script += codegen.to_source(tree)

    with open(script_out, 'w') as f:
        f.writelines(script)

def main():
    print transform(sys.argv[1], sys.argv[2])

if __name__ == "__main__":
    main()
<commit_msg>Add transformer to assert positive operands.<commit_after>"""
Prepare a python program (tiny subset) for compliation into sed.
Transform into positive form:
- all operators and binary comparisons are replaced with call to functions
  x + y --> signed_add(x, y), idem -, *, //, ==, !=, <, <=, >, >=
- all operands of operators and binary comparisons inside signed_xxx functions
  must be positive integers.
- two functions handling negative values remain: is_positive and negative. They
  are treated separately when testing python positive form or compiling into sed.
- augmented assignments are replaced with simple assignments.

A testing mode of the transformer program generates code testing that arguments
of operators and comparisons are positive.
"""

from __future__ import division

import sys
import inspect
import ast
import operator
import codegen

# -- Transformer --------------------------------------------------------------

signed_func = {
    ast.Add: 'signed_add',
    ast.Sub: 'signed_sub',
    ast.Mult: 'signed_mult',
    ast.Div: None,
    ast.FloorDiv: 'signed_div',
    ast.Mod: 'signed_mod',
    ast.Eq: 'signed_eq',
    ast.NotEq: 'signed_noteq',
    ast.Lt: 'signed_lt',
    ast.LtE: 'signed_lte',
    ast.Gt: 'signed_gt',
    ast.GtE: 'signed_gte'}

unsigned_func = {
    ast.Add: 'unsigned_add',
    ast.Sub: 'unsigned_sub',
    ast.Mult: 'unsigned_mult',
    ast.Div: None,
    ast.FloorDiv: 'unsigned_div',
    ast.Mod: 'unsigned_mod',
    ast.Eq: 'unsigned_eq',
    ast.NotEq: 'unsigned_noteq',
    ast.Lt: 'unsigned_lt',
    ast.LtE: 'unsigned_lte',
    ast.Gt: 'unsigned_gt',
    ast.GtE: 'unsigned_gte'}

class NumsedAstTransformer(ast.NodeTransformer):

    def __init__(self, func):
        self.func = func
        self.required_func = set()

    def make_call(self, func, args):
        self.required_func.add(func)
        return ast.Call(func=ast.Name(id=func, ctx=ast.Load()),
                        args=args,
                        keywords=[], starargs=None, kwargs=None)

    def visit_BinOp(self, node):
        self.generic_visit(node)
        return self.make_call(self.func[type(node.op)], [node.left, node.right])

    def visit_Compare(self, node):
        self.generic_visit(node)
        left = node.left
        ops = node.ops
        comparators = node.comparators
        if len(ops) == 1:
            return self.make_call(self.func[type(node.ops[0])], [node.left, comparators[0]])
        else:
            list_compare = []
            while ops:
                compare = self.make_call(self.func[type(ops[0])], [left, comparators[0]])
                list_compare.append(compare)
                left = comparators[0]
                ops = ops[1:]
                comparators = comparators[1:]
            return ast.BoolOp(op=ast.And(), values=list_compare)

    def visit_AugAssign(self, node):
        # AugAssign(target=Name(id='x', ctx=Store()), op=Add(), value=Num(n=1)),
        # Assign(targets=[Name(id='x', ctx=Store())], value=BinOp(left=Name(id='x', ctx=Load()), op=Add(), right=Num(n=1)))])
        #self.generic_visit(node)
        target = node.target
        source = ast.Name(id=target.id, ctx=ast.Load())
        return ast.Assign(targets=[target], value=self.visit_BinOp(ast.BinOp(left=source, op=node.op, right=node.value)))

# -- Builtin functions -------------------------------------------------------

def signed_eq(x, y):
    if is_positive(x):
        if is_positive(y):
            return x == y
        else:
            return False
    else:
        if is_positive(y):
            return False
        else:
            return negative(x) == negative(y)

def signed_noteq(x, y):
    return not signed_eq(x, y)

def signed_lt(x, y):
    if is_positive(x):
        if is_positive(y):
            return x < y
        else:
            return True
    else:
        if is_positive(y):
            return True
        else:
            return negative(x) < negative(y)

def signed_gte(x, y):
    return not signed_lt(x, y)

def signed_add(x, y):
    if is_positive(x):
        if is_positive(y):
            r = x + y
        else:
            y = negative(y)
            if x > y:
                r = x - y
            else:
                r = negative(y - x)
    else:
        x = negative(x)
        if is_positive(y):
            if x > y:
                r = negative(x - y)
            else:
                r = y - x
        else:
            y = negative(y)
            r = negative(x + y)
    return r

def signed_sub(x, y):
    if is_positive(x):
        if is_positive(y):
            if x > y:
                return x - y
            else:
                return negative(y - x)
        else:
            return x + negative(y)
    else:
        abs_x = negative(x)
        if is_positive(y):
            return negative(abs_x + y)
        else:
            abs_y = negative(y)
            if abs_x > abs_y:
                return negative(abs_x - abs_y)
            else:
                return abs_y - abs_x

def signed_pos(x):
    return x

def signed_neg(x):
    return negative(x)

def euclide(a, b):
    # http://compoasso.free.fr/primelistweb/page/prime/euclide.php
    r = a
    q = 0
    n = 0
    aux = b

    while aux <= a:
        aux = aux * 2
        n += 1

    while n > 0:
        #aux = aux / 2
        aux *= 5
        aux //= 10 # TODO: //10 should be a primitive divide_by_10()
        n -= 1
        q = q * 2
        if r >= aux:
            r -= aux
            q += 1

    return q

def modulo(a, b):
    q = euclide(a, b)
    return a - b * q

def signed_mult(a, b):
    if (a >= 0 and b >= 0) or (a <= 0 and b <= 0):
        return a * b
    else:
        return negative(a * b)

def signed_div(a, b):
    if (a >= 0 and b >= 0) or (a <= 0 and b <= 0):
        return euclide(a, b)
    else:
        return negative(euclide(a, b))

def signed_mod(a, b):
    q = signed_divide(a, b)
    return a - b * q

# -- Primitives --------------------------------------------------------------

def is_positive(x):
    return operator.ge(x, 0)

def negative(x):
    return operator.neg(x)

# -- Testing transformation --------------------------------------------------

unsigned_func_pattern = """
def %s(x, y):
    assert x >= 0
    assert y >= 0
    return x %s y
"""

unsigned_op = {
    'unsigned_add': '+',
    'unsigned_sub': '-',
    'unsigned_mult': '*',
    'unsigned_div': '//',
    'unsigned_mod': '%',
    'unsigned_eq': '==',
    'unsigned_noteq': '!=',
    'unsigned_lt': '<',
    'unsigned_lte': '<=',
    'unsigned_gt': '>',
    'unsigned_gte': '>='}

def make_unsigned_func(name):
    return unsigned_func_pattern % (name, unsigned_op[name])

# -- Tests -------------------------------------------------------------------

class NumsedAstVisitor(ast.NodeVisitor):
    def generic_visit(self, node):
        print type(node).__name__
        ast.NodeVisitor.generic_visit(self, node)

    #def visit_BinOp(self, node):
    #    print ast.dump(node)

# -- Main --------------------------------------------------------------------

def transform(script_in, script_out):
    test = False

    tree = ast.parse(open(script_in).read())
    if test:
        print ast.dump(tree)
        exec(compile(tree, filename="<ast>", mode="exec"))
        print '--'

    numsed_ast_transformer = NumsedAstTransformer(signed_func)
    numsed_ast_transformer.visit(tree)

    if test:
        ast.fix_missing_locations(tree)
        print ast.dump(tree)
        # going to and from AST
        exec(compile(tree, filename="<ast>", mode="exec"))

    builtin = numsed_ast_transformer.required_func
    builtin = [globals()[x] for x in builtin]
    if signed_div in builtin or signed_mod in builtin:
        builtin.append(euclide)
    builtin += [is_positive, negative]
    #print builtin

    # add builtin functions to code to compile
    script = ''
    for func in builtin:
        script += '\n'
        script += ''.join(inspect.getsourcelines(func)[0])
    script += '\n'
    script += codegen.to_source(tree)

    with open(script_out, 'w') as f:
        f.writelines(script)

def transform_assert(script_in, script_out):
    test = False

    tree = ast.parse(open(script_in).read())
    if test:
        print ast.dump(tree)
        exec(compile(tree, filename="<ast>", mode="exec"))
        print '--'

    numsed_ast_transformer = NumsedAstTransformer(unsigned_func)
    numsed_ast_transformer.visit(tree)

    if test:
        ast.fix_missing_locations(tree)
        print ast.dump(tree)
        # going to and from AST
        exec(compile(tree, filename="<ast>", mode="exec"))

    # add builtin functions to code to compile
    builtin = numsed_ast_transformer.required_func
    script = ''
    for func in builtin:
        script += '\n'
        script += make_unsigned_func(func)
    script += '\n'
    script += 'import operator\n'
    script += codegen.to_source(tree)

    with open(script_out, 'w') as f:
        f.writelines(script)

def main():
    pass2 = True
    print transform(sys.argv[1], sys.argv[2])
    # pass #2
    if pass2:
        print transform_assert(sys.argv[2], 'tmpassert.py')

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


<commit_before>###
# Copyright (c) 2005, Jeremiah Fincher
# Copyright (c) 2009, James McCoy
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

import re
import HTMLParser
import htmlentitydefs

import supybot.conf as conf
import supybot.utils as utils
from supybot.commands import *
import supybot.plugins as plugins
import supybot.ircutils as ircutils
import supybot.callbacks as callbacks

class Title(HTMLParser.HTMLParser):
    entitydefs = htmlentitydefs.entitydefs.copy()
    entitydefs['nbsp'] = ' '
    entitydefs['apos'] = '\''
    def __init__(self):
        self.inTitle = False
        self.title = ''
        HTMLParser.HTMLParser.__init__(self)

    def handle_starttag(self, tag, attrs):
        if tag == 'title':
            self.inTitle = True

    def handle_endtag(self, tag):
        if tag == 'title':
            self.inTitle = False

    def handle_data(self, data):
        if self.inTitle:
            self.title += data

    def handle_entityref(self, name):
        if self.inTitle:
            if name in self.entitydefs:
                self.title += self.entitydefs[name]

class Web(callbacks.PluginRegexp):
    """Add the help for "@help Web" here."""
    threaded = True
    regexps = ['titleSnarfer']
    def callCommand(self, command, irc, msg, *args, **kwargs):
        try:
            super(Web, self).callCommand(command, irc, msg, *args, **kwargs)
        except utils.web.Error, e:
            irc.reply(str(e))

    def titleSnarfer(self, irc, msg, match):
        channel = msg.args[0]
        if not irc.isChannel(channel):
            return
        if callbacks.addressed(irc.nick, msg):
            return
        if self.registryValue('titleSnarfer', channel):
            url = match.group(0)
            r = self.registryValue('nonSnarfingRegexp', channel)
            if r and r.search(url):
                self.log.debug('Not titleSnarfing %q.', url)
                return
            try:
                size = conf.supybot.protocols.http.peekSize()
                text = utils.web.getUrl(url, size=size)
            except utils.web.Error, e:
                self.log.info('Couldn\'t snarf title of %u: %s.', url, e)
                return
            parser = Title()
            try:
                parser.feed(text)
            except HTMLParser.HTMLParseError:
                self.log.debug('Encountered a problem parsing %u.  Title may '
                               'already be set, though', url)
            if parser.title:
                domain = utils.web.getDomain(url)
                title = utils.web.htmlToText(parser.title.strip())
                s = format('Title: %s (at %s)', title, domain)
                irc.reply(s, prefixNick=False)
    titleSnarfer = urlSnarfer(titleSnarfer)
    titleSnarfer.__doc__ = utils.web._httpUrlRe

    def headers(self, irc, msg, args, url):
        """<url>

        Returns the HTTP headers of <url>.  Only HTTP urls are valid, of
        course.
        """
        fd = utils.web.getUrlFd(url)
        try:
            s = ', '.join([format('%s: %s', k, v)
                           for (k, v) in fd.headers.items()])
            irc.reply(s)
        finally:
            fd.close()
    headers = wrap(headers, ['httpUrl'])

    _doctypeRe = re.compile(r'(<!DOCTYPE[^>]+>)', re.M)
    def doctype(self, irc, msg, args, url):
        """<url>

        Returns the DOCTYPE string of <url>.  Only HTTP urls are valid, of
        course.
        """
        size = conf.supybot.protocols.http.peekSize()
        s = utils.web.getUrl(url, size=size)
        m = self._doctypeRe.search(s)
        if m:
            s = utils.str.normalizeWhitespace(m.group(0))
            irc.reply(s)
        else:
            irc.reply('That URL has no specified doctype.')
    doctype = wrap(doctype, ['httpUrl'])

    def size(self, irc, msg, args, url):
        """<url>

        Returns the Content-Length header of <url>.  Only HTTP urls are valid,
        of course.
        """
        fd = utils.web.getUrlFd(url)
        try:
            try:
                size = fd.headers['Content-Length']
                irc.reply(format('%u is %i bytes long.', url, size))
            except KeyError:
                size = conf.supybot.protocols.http.peekSize()
                s = fd.read(size)
                if len(s) != size:
                    irc.reply(format('%u is %i bytes long.', url, len(s)))
                else:
                    irc.reply(format('The server didn\'t tell me how long %u '
                                     'is but it\'s longer than %i bytes.',
                                     url, size))
        finally:
            fd.close()
    size = wrap(size, ['httpUrl'])

    def title(self, irc, msg, args, url):
        """<url>

        Returns the HTML <title>...</title> of a URL.
        """
        size = conf.supybot.protocols.http.peekSize()
        text = utils.web.getUrl(url, size=size)
        parser = Title()
        try:
            parser.feed(text)
        except HTMLParser.HTMLParseError:
            self.log.debug('Encountered a problem parsing %u.  Title may '
                           'already be set, though', url)
        if parser.title:
            irc.reply(utils.web.htmlToText(parser.title.strip()))
        elif len(text) < size:
            irc.reply('That URL appears to have no HTML title.')
        else:
            irc.reply(format('That URL appears to have no HTML title '
                             'within the first %i bytes.', size))
    title = wrap(title, ['httpUrl'])

    _netcraftre = re.compile(r'td align="left">\s+<a[^>]+>(.*?)<a href',
                             re.S | re.I)
    def netcraft(self, irc, msg, args, hostname):
        """<hostname|ip>

        Returns Netcraft.com's determination of what operating system and
        webserver is running on the host given.
        """
        url = 'http://uptime.netcraft.com/up/graph/?host=' + hostname
        html = utils.web.getUrl(url)
        m = self._netcraftre.search(html)
        if m:
            html = m.group(1)
            s = utils.web.htmlToText(html, tagReplace='').strip()
            s = s.rstrip('-').strip()
            irc.reply(s) # Snip off "the site"
        elif 'We could not get any results' in html:
            irc.reply('No results found for %s.' % hostname)
        else:
            irc.error('The format of page the was odd.')
    netcraft = wrap(netcraft, ['text'])

    def urlquote(self, irc, msg, args, text):
        """<text>

        Returns the URL quoted form of the text.
        """
        irc.reply(utils.web.urlquote(text))
    urlquote = wrap(urlquote, ['text'])

    def urlunquote(self, irc, msg, args, text):
        """<text>

        Returns the text un-URL quoted.
        """
        s = utils.web.urlunquote(text)
        irc.reply(s)
    urlunquote = wrap(urlunquote, ['text'])

    def fetch(self, irc, msg, args, url):
        """<url>

        Returns the contents of <url>, or as much as is configured in
        supybot.plugins.Web.fetch.maximum.  If that configuration variable is
        set to 0, this command will be effectively disabled.
        """
        max = self.registryValue('fetch.maximum')
        if not max:
            irc.error('This command is disabled '
                      '(supybot.plugins.Web.fetch.maximum is set to 0).',
                      Raise=True)
        timeout = self.registryValue('fetch.timeout')
        if timeout == 0:
            timeout = None
        fd = utils.web.getUrlFd(url, timeout=timeout)
        irc.reply(fd.read(max))
    fetch = wrap(fetch, ['url'])

Class = Web

# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79:
<commit_msg>Web: Fix fetch. Use getUrl instead of getUrlFd.<commit_after>###
# Copyright (c) 2005, Jeremiah Fincher
# Copyright (c) 2009, James McCoy
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

import re
import HTMLParser
import htmlentitydefs

import supybot.conf as conf
import supybot.utils as utils
from supybot.commands import *
import supybot.plugins as plugins
import supybot.ircutils as ircutils
import supybot.callbacks as callbacks

class Title(HTMLParser.HTMLParser):
    entitydefs = htmlentitydefs.entitydefs.copy()
    entitydefs['nbsp'] = ' '
    entitydefs['apos'] = '\''
    def __init__(self):
        self.inTitle = False
        self.title = ''
        HTMLParser.HTMLParser.__init__(self)

    def handle_starttag(self, tag, attrs):
        if tag == 'title':
            self.inTitle = True

    def handle_endtag(self, tag):
        if tag == 'title':
            self.inTitle = False

    def handle_data(self, data):
        if self.inTitle:
            self.title += data

    def handle_entityref(self, name):
        if self.inTitle:
            if name in self.entitydefs:
                self.title += self.entitydefs[name]

class Web(callbacks.PluginRegexp):
    """Add the help for "@help Web" here."""
    threaded = True
    regexps = ['titleSnarfer']
    def callCommand(self, command, irc, msg, *args, **kwargs):
        try:
            super(Web, self).callCommand(command, irc, msg, *args, **kwargs)
        except utils.web.Error, e:
            irc.reply(str(e))

    def titleSnarfer(self, irc, msg, match):
        channel = msg.args[0]
        if not irc.isChannel(channel):
            return
        if callbacks.addressed(irc.nick, msg):
            return
        if self.registryValue('titleSnarfer', channel):
            url = match.group(0)
            r = self.registryValue('nonSnarfingRegexp', channel)
            if r and r.search(url):
                self.log.debug('Not titleSnarfing %q.', url)
                return
            try:
                size = conf.supybot.protocols.http.peekSize()
                text = utils.web.getUrl(url, size=size)
            except utils.web.Error, e:
                self.log.info('Couldn\'t snarf title of %u: %s.', url, e)
                return
            parser = Title()
            try:
                parser.feed(text)
            except HTMLParser.HTMLParseError:
                self.log.debug('Encountered a problem parsing %u.  Title may '
                               'already be set, though', url)
            if parser.title:
                domain = utils.web.getDomain(url)
                title = utils.web.htmlToText(parser.title.strip())
                s = format('Title: %s (at %s)', title, domain)
                irc.reply(s, prefixNick=False)
    titleSnarfer = urlSnarfer(titleSnarfer)
    titleSnarfer.__doc__ = utils.web._httpUrlRe

    def headers(self, irc, msg, args, url):
        """<url>

        Returns the HTTP headers of <url>.  Only HTTP urls are valid, of
        course.
        """
        fd = utils.web.getUrlFd(url)
        try:
            s = ', '.join([format('%s: %s', k, v)
                           for (k, v) in fd.headers.items()])
            irc.reply(s)
        finally:
            fd.close()
    headers = wrap(headers, ['httpUrl'])

    _doctypeRe = re.compile(r'(<!DOCTYPE[^>]+>)', re.M)
    def doctype(self, irc, msg, args, url):
        """<url>

        Returns the DOCTYPE string of <url>.  Only HTTP urls are valid, of
        course.
        """
        size = conf.supybot.protocols.http.peekSize()
        s = utils.web.getUrl(url, size=size)
        m = self._doctypeRe.search(s)
        if m:
            s = utils.str.normalizeWhitespace(m.group(0))
            irc.reply(s)
        else:
            irc.reply('That URL has no specified doctype.')
    doctype = wrap(doctype, ['httpUrl'])

    def size(self, irc, msg, args, url):
        """<url>

        Returns the Content-Length header of <url>.  Only HTTP urls are valid,
        of course.
        """
        fd = utils.web.getUrlFd(url)
        try:
            try:
                size = fd.headers['Content-Length']
                irc.reply(format('%u is %i bytes long.', url, size))
            except KeyError:
                size = conf.supybot.protocols.http.peekSize()
                s = fd.read(size)
                if len(s) != size:
                    irc.reply(format('%u is %i bytes long.', url, len(s)))
                else:
                    irc.reply(format('The server didn\'t tell me how long %u '
                                     'is but it\'s longer than %i bytes.',
                                     url, size))
        finally:
            fd.close()
    size = wrap(size, ['httpUrl'])

    def title(self, irc, msg, args, url):
        """<url>

        Returns the HTML <title>...</title> of a URL.
        """
        size = conf.supybot.protocols.http.peekSize()
        text = utils.web.getUrl(url, size=size)
        parser = Title()
        try:
            parser.feed(text)
        except HTMLParser.HTMLParseError:
            self.log.debug('Encountered a problem parsing %u.  Title may '
                           'already be set, though', url)
        if parser.title:
            irc.reply(utils.web.htmlToText(parser.title.strip()))
        elif len(text) < size:
            irc.reply('That URL appears to have no HTML title.')
        else:
            irc.reply(format('That URL appears to have no HTML title '
                             'within the first %i bytes.', size))
    title = wrap(title, ['httpUrl'])

    _netcraftre = re.compile(r'td align="left">\s+<a[^>]+>(.*?)<a href',
                             re.S | re.I)
    def netcraft(self, irc, msg, args, hostname):
        """<hostname|ip>

        Returns Netcraft.com's determination of what operating system and
        webserver is running on the host given.
        """
        url = 'http://uptime.netcraft.com/up/graph/?host=' + hostname
        html = utils.web.getUrl(url)
        m = self._netcraftre.search(html)
        if m:
            html = m.group(1)
            s = utils.web.htmlToText(html, tagReplace='').strip()
            s = s.rstrip('-').strip()
            irc.reply(s) # Snip off "the site"
        elif 'We could not get any results' in html:
            irc.reply('No results found for %s.' % hostname)
        else:
            irc.error('The format of page the was odd.')
    netcraft = wrap(netcraft, ['text'])

    def urlquote(self, irc, msg, args, text):
        """<text>

        Returns the URL quoted form of the text.
        """
        irc.reply(utils.web.urlquote(text))
    urlquote = wrap(urlquote, ['text'])

    def urlunquote(self, irc, msg, args, text):
        """<text>

        Returns the text un-URL quoted.
        """
        s = utils.web.urlunquote(text)
        irc.reply(s)
    urlunquote = wrap(urlunquote, ['text'])

    def fetch(self, irc, msg, args, url):
        """<url>

        Returns the contents of <url>, or as much as is configured in
        supybot.plugins.Web.fetch.maximum.  If that configuration variable is
        set to 0, this command will be effectively disabled.
        """
        max = self.registryValue('fetch.maximum')
        if not max:
            irc.error('This command is disabled '
                      '(supybot.plugins.Web.fetch.maximum is set to 0).',
                      Raise=True)
        timeout = self.registryValue('fetch.timeout')
        if timeout == 0:
            timeout = None
        s = utils.web.getUrl(url, timeout=timeout, size=max)
        irc.reply(s)
    fetch = wrap(fetch, ['url'])

Class = Web

# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79:
<|endoftext|>