<commit_before>#################################################################
##
##  Almost done - May be divide growing and senescence
##  into separate functions.
##
##  'Field' concept is implemented for Single Species Vegetation.
##
##  Sai Nudurupati - 04Mar2014
#################################################################

from landlab import Component

import numpy as np

_VALID_METHODS = set(['Grid'])

def assert_method_is_valid(method):
    if method not in _VALID_METHODS:
        raise ValueError('%s: Invalid method name' % method)

class Vegetation( Component ):
    """
    Landlab component that implements 1D and 2D vegetation dynamics
    model.

    >>> grid = RasterModelGrid(5, 4, 1.e4)
    >>> veg = Vegetation(grid)
    >>> veg.name
    Vegetation
    """
    _name = 'Vegetation'

    _input_var_names = set([
        'ActualEvapotranspiration',
        'WaterStress',
        'PotentialEvapotranspiration',
    ])

    _output_var_names = set([
        'LiveLeafAreaIndex',
        'DeadLeafAreaIndex',
        'VegetationCover',
        'LiveBiomass',
        'DeadBiomass',
    ])

    _var_units = {
        'LiveLeafAreaIndex' : 'None',
        'DeadLeafAreaIndex' : 'None',
        'VegetationCover'  : 'None',
        'ActualEvapotranspiration': 'mm',
        'PotentialEvapotranspiration': 'mm',
        'WaterStress': 'Pa',
        'LiveBiomass' : 'g DM m^-2 d^-1',
        'DeadBiomass' : 'g DM m^-2 d^-1',
    }

    def __init__(self, grid, **kwds):
        self._method = kwds.pop('method', 'Grid')
        self._WUE = kwds.pop('WUE', 0.01)
        self._LAI_max = kwds.pop('LAI_MAX', 5)
        self._cb = kwds.pop('CB', 0.0047)
        self._cd = kwds.pop('CD', 0.009)
        self._ksg = kwds.pop('KSG', 0.008)
        self._kdd = kwds.pop('KDD', 0.01)
        self._kws = kwds.pop('KWS', 0.02)
        self._Blive_init = kwds.pop('BLIVE_INI', 4)
        self._Bdead_init = kwds.pop('BDEAD_INI', 450)
        self._ETthresholdup = kwds.pop('ETTup', 3.8)
        self._ETthresholddown = kwds.pop('ETTdwn', 6.5)
        self._ETdmax = kwds.pop('ETdmax', 10)

        assert_method_is_valid(self._method)

        super(Vegetation, self).__init__(grid, **kwds)

        for name in self._input_var_names:
            if not name in self.grid.at_cell:
                self.grid.add_zeros('cell', name, units=self._var_units[name])

        for name in self._output_var_names:
            if not name in self.grid.at_cell:
                self.grid.add_zeros('cell', name, units=self._var_units[name])

        self._cell_values = self.grid['cell']

        self._Blive_ini = self._Blive_init * np.ones(self.grid.number_of_cells)
        self._Bdead_ini = self._Bdead_init * np.ones(self.grid.number_of_cells)

    def update(self, **kwds):

        PETthreshold = kwds.pop('PotentialEvapotranspirationThreshold', 0)
        Tb = kwds.pop('Tb', 24.)
        Tr = kwds.pop('Tr', 0.01)
        PET = self._cell_values['PotentialEvapotranspiration']
        ActualET = self._cell_values['ActualEvapotranspiration']
        Water_stress = self._cell_values['WaterStress']

        self._LAIlive = self._cell_values['LiveLeafAreaIndex']
        self._LAIdead = self._cell_values['DeadLeafAreaIndex']
        self._Blive = self._cell_values['LiveBiomass']
        self._Bdead = self._cell_values['DeadBiomass']
        self._VegCov = self._cell_values['VegetationCover']

        for cell in range(0, self.grid.number_of_cells):

            LAIlive =  min (self._cb * self._Blive_ini[cell], self._LAI_max)
            LAIdead =  min (self._cd * self._Bdead_ini[cell], (self._LAI_max   \
                            - LAIlive))

            if PET[cell] > PETthreshold:  # Growing Season

                NPP = max((ActualET[cell]/(Tb+Tr))*                \
                            self._WUE*24*0.55*1000, 0.001)
                Bmax = (self._LAI_max - LAIdead)/self._cb
                Yconst = (1/((1/Bmax)+(((self._kws*Water_stress[cell])+        \
                            self._ksg)/NPP)))
                Blive = (self._Blive_ini[cell] - Yconst) *np.exp(-(NPP/Yconst)*\
                            ((Tb+Tr)/24)) + Yconst
                Bdead = (self._Bdead_ini[cell] + (Blive - max(Blive *          \
                            np.exp(-self._ksg * Tb/24),0.00001)))*       \
                            np.exp(-self._kdd * min( PET[cell]/10, 1 ) *       \
                            Tb/24)

            else:                                 # Senescense

                Blive = max(self._Blive_ini[cell] * np.exp((-2) * self._ksg *  \
                            Tb/24), 1 )
                Bdead = max( (self._Bdead_ini[cell] + ( self._Blive_ini[cell]  \
                            -max( self._Blive_ini[cell]*np.exp( (-2) *         \
                            self._ksg * Tb/24), 0.000001)))*             \
                            np.exp( (-1)*self._kdd *                           \
                            min( PET[cell]/10, 1 ) * Tb/24), 0 )

            LAIlive =  min( self._cb * Blive, self._LAI_max )
            LAIdead =  min( self._cd * Bdead, ( self._LAI_max - LAIlive ) )
            Vt = 1 - np.exp(-0.75 * (LAIlive + LAIdead))

            self._LAIlive[cell] = LAIlive
            self._LAIdead[cell] = LAIdead
            self._VegCov[cell] = Vt
            self._Blive[cell] = Blive
            self._Bdead[cell] = Bdead

        self._Blive_ini = self._Blive
        self._Bdead_ini = self._Bdead                       <commit_msg>Minor changes<commit_after>#################################################################
##
##  Almost done - May be divide growing and senescence
##  into separate functions.
##
##  'Field' concept is implemented for Single Species Vegetation.
##
##  Sai Nudurupati - 04Mar2014
#################################################################

from landlab import Component

import numpy as np

_VALID_METHODS = set(['Grid'])

def assert_method_is_valid(method):
    if method not in _VALID_METHODS:
        raise ValueError('%s: Invalid method name' % method)

class Vegetation( Component ):
    """
    Landlab component that implements 1D and 2D vegetation dynamics
    model.

    >>> grid = RasterModelGrid(5, 4, 1.e4)
    >>> veg = Vegetation(grid)
    >>> veg.name
    Vegetation
    """
    _name = 'Vegetation'

    _input_var_names = set([
        'ActualEvapotranspiration',
        'WaterStress',
        'PotentialEvapotranspiration',
    ])

    _output_var_names = set([
        'LiveLeafAreaIndex',
        'DeadLeafAreaIndex',
        'VegetationCover',
        'LiveBiomass',
        'DeadBiomass',
    ])

    _var_units = {
        'LiveLeafAreaIndex' : 'None',
        'DeadLeafAreaIndex' : 'None',
        'VegetationCover'  : 'None',
        'ActualEvapotranspiration': 'mm',
        'PotentialEvapotranspiration': 'mm',
        'WaterStress': 'Pa',
        'LiveBiomass' : 'g DM m^-2 d^-1',
        'DeadBiomass' : 'g DM m^-2 d^-1',
    }

    def __init__(self, grid, **kwds):
        self._method = kwds.pop('method', 'Grid')
        self._WUE = kwds.pop('WUE', 0.01)
        self._LAI_max = kwds.pop('LAI_MAX', 5)
        self._cb = kwds.pop('CB', 0.0047)
        self._cd = kwds.pop('CD', 0.009)
        self._ksg = kwds.pop('KSG', 0.008)
        self._kdd = kwds.pop('KDD', 0.01)
        self._kws = kwds.pop('KWS', 0.02)
        self._Blive_init = kwds.pop('BLIVE_INI', 4)
        self._Bdead_init = kwds.pop('BDEAD_INI', 450)
        self._ETthresholdup = kwds.pop('ETTup', 3.8)
        self._ETthresholddown = kwds.pop('ETTdwn', 6.5)
        self._ETdmax = kwds.pop('ETdmax', 10)

        assert_method_is_valid(self._method)

        super(Vegetation, self).__init__(grid, **kwds)

        for name in self._input_var_names:
            if not name in self.grid.at_cell:
                self.grid.add_zeros('cell', name, units=self._var_units[name])

        for name in self._output_var_names:
            if not name in self.grid.at_cell:
                self.grid.add_zeros('cell', name, units=self._var_units[name])

        self._cell_values = self.grid['cell']

        self._Blive_ini = self._Blive_init * np.ones(self.grid.number_of_cells)
        self._Bdead_ini = self._Bdead_init * np.ones(self.grid.number_of_cells)

    def update(self, **kwds):

        PETthreshold = kwds.pop('PotentialEvapotranspirationThreshold', 0)
        Tb = kwds.pop('Tb', 24.)
        Tr = kwds.pop('Tr', 0.01)
        PET = self._cell_values['PotentialEvapotranspiration']
        ActualET = self._cell_values['ActualEvapotranspiration']
        Water_stress = self._cell_values['WaterStress']

        self._LAIlive = self._cell_values['LiveLeafAreaIndex']
        self._LAIdead = self._cell_values['DeadLeafAreaIndex']
        self._Blive = self._cell_values['LiveBiomass']
        self._Bdead = self._cell_values['DeadBiomass']
        self._VegCov = self._cell_values['VegetationCover']

        for cell in range(0, self.grid.number_of_cells):

            LAIlive =  min (self._cb * self._Blive_ini[cell], self._LAI_max)
            LAIdead =  min (self._cd * self._Bdead_ini[cell], (self._LAI_max   \
                            - LAIlive))

            if PET[cell] > PETthreshold:  # Growing Season

                NPP = max((ActualET[cell]/(Tb+Tr))*                            \
                            self._WUE*24*0.55*1000, 0.001)
                Bmax = (self._LAI_max - LAIdead)/self._cb
                Yconst = (1/((1/Bmax)+(((self._kws*Water_stress[cell])+        \
                            self._ksg)/NPP)))
                Blive = (self._Blive_ini[cell] - Yconst) *np.exp(-(NPP/Yconst)*\
                            ((Tb+Tr)/24)) + Yconst
                Bdead = (self._Bdead_ini[cell] + (Blive - max(Blive *          \
                            np.exp(-self._ksg * Tb/24),0.00001)))*             \
                            np.exp(-self._kdd * min( PET[cell]/10, 1 ) *       \
                            Tb/24)

            else:                                 # Senescense

                Blive = max(self._Blive_ini[cell] * np.exp((-2) * self._ksg *  \
                            Tb/24), 1 )
                Bdead = max( (self._Bdead_ini[cell] + ( self._Blive_ini[cell]  \
                            -max( self._Blive_ini[cell]*np.exp( (-2) *         \
                            self._ksg * Tb/24), 0.000001)))*                   \
                            np.exp( (-1)*self._kdd *                           \
                            min( PET[cell]/10, 1 ) * Tb/24), 0 )

            LAIlive =  min( self._cb * Blive, self._LAI_max )
            LAIdead =  min( self._cd * Bdead, ( self._LAI_max - LAIlive ) )
            Vt = 1 - np.exp(-0.75 * (LAIlive + LAIdead))

            self._LAIlive[cell] = LAIlive
            self._LAIdead[cell] = LAIdead
            self._VegCov[cell] = Vt
            self._Blive[cell] = Blive
            self._Bdead[cell] = Bdead

        self._Blive_ini = self._Blive
        self._Bdead_ini = self._Bdead                       <|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

# =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2014 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

from email.mime.text import MIMEText
import types
import email.utils
import logging
import smtplib
import json

import requests
from flask_babel import gettext
from util import render_template2

LOGGER = logging.getLogger(__name__)

def do_email(config, resource, run, status_changed, result):
    # List of global email addresses to notify, may be list or
    # comma-separated str "To" needs comma-separated list,
    # while sendmail() requires list...

    if isinstance(config['GHC_NOTIFICATIONS_EMAIL'], types.StringTypes):
        config['GHC_NOTIFICATIONS_EMAIL'] = \
            config['GHC_NOTIFICATIONS_EMAIL'].split(',')

    # this should always be a list
    global_notifications = config['GHC_NOTIFICATIONS_EMAIL'] or []
    if not isinstance(global_notifications, (list, tuple, set,)):
        raise TypeError("Cannot use {} as list of emails".format(
                                       type(global_notifications)))

    notifications_email = global_notifications +\
        resource.get_recipients('email')

    if not notifications_email:
        LOGGER.warning("No emails for notification set for resource %s",
                       resource.identifier)
        return

    template_vars = {
        'result': result,
        'config': config,
        'resource': resource,
        'run': run
    }

    msgbody = render_template2('notification_email.txt', template_vars)
    msg = MIMEText(msgbody)
    msg['From'] = email.utils.formataddr((config['GHC_SITE_TITLE'],
                                          config['GHC_ADMIN_EMAIL']))
    msg['To'] = ','.join(notifications_email)
    msg['Subject'] = '[%s] %s: %s' % (config['GHC_SITE_TITLE'],
                                      result, resource.title)

    if not config.get('GHC_SMTP') or not\
        (any([config['GHC_SMTP'][k] for k in ('port',
                                              'server',
                                              'username',
                                              'password',)])):

        LOGGER.warning("No SMTP configuration. Not sending to %s",
                       notifications_email)
        print(msg.as_string())
        return

    server = smtplib.SMTP(config['GHC_SMTP']['server'],
                          config['GHC_SMTP']['port'])

    if config['DEBUG']:
        server.set_debuglevel(True)

    try:
        if config['GHC_SMTP']['tls']:
            server.starttls()
        else:
            server.connect()
    except Exception, err:
        LOGGER.exception("Cannot connect to smtp: %s[:%s]: %s",
                         config['GHC_SMTP']['server'],
                         config['GHC_SMTP']['port'],
                         err,
                         exc_info=err)
        return
    try:
        server.login(config['GHC_SMTP']['username'],
                     config['GHC_SMTP']['password'])
    except Exception, err:
        LOGGER.exception("Cannot log in to smtp: %s", err,
                         exc_info=err)
    try:
        server.sendmail(config['GHC_ADMIN_EMAIL'],
                        config['GHC_NOTIFICATIONS_EMAIL'],
                        msg.as_string())
    except Exception as err:
        LOGGER.exception(str(err), exc_info=err)
    finally:
        server.quit()

def _parse_line(_line):
    try:
        k, v = _line.split('=', 1)
        return {k: v}
    except (IndexError, ValueError,):
        raise ValueError("Invalid line: {}".format(_line))

def _parse_webhook_location(value):
    """
    Parse Recipient.location and returns tuple of url and params

    location should be in form

    URL

    PAYLOAD

    where PAYLOAD is a list of fields and values in form

    FIELD_NAME=FIELD_VALUE

    alternatively, it can be dictionary serialized as json
    """
    if not value.strip():
        raise ValueError("No payload")
    value = value.strip()
    url = None
    params = {}
    lines = value.splitlines()

    for idx, line in enumerate(lines):
        if idx == 0:
            url = line
        elif idx == 1:
            if line.strip():
                raise ValueError("Second line should be empty")
        elif idx == 2:
            try:
                params = json.loads('\n'.join(lines[2:]))
                break
            except (TypeError, ValueError,):
                params.update(_parse_line(line))
        else:
            params.update(_parse_line(line))

    if url is None:
        raise ValueError("Cannot parse url")
    return url, params,

def do_webhook(config, resource, run, status_changed, result):
    """
    Process webhook recipients for resource

    location should be in format:

    URL

    [PAYLOAD]

    There's blank line between URL and PAYLOAD. PAYLOAD
    should be either json or list of field=value items
    in each line.

    Webhook's request is POST send to url with payload containing
    PAYLOAD
    and fields:

    ghc.result=(result of test)
    ghc.resource.url=(url of resource)
    ghc.resource.title=(title of resource)

    """
    recipients = resource.get_recipients('webhook')
    if not recipients:
        return
    for rcp in recipients:
        try:
            url, params = _parse_webhook_location(rcp)
        except ValueError, err:
            LOGGER.warning("Cannot send to {}: {}"
                           .format(rcp, err), exc_info=err)

        resource_view = '{}/resource/{}'.format(
                                            config['GHC_SITE_URL'],
                                            resource.identifier)

        params['ghc.result'] = result
        params['ghc.resource.url'] = resource.url
        params['ghc.resource.title'] = resource.title
        params['ghc.resource.type'] = resource.resource_type
        params['ghc.resource.view'] = resource_view

        try:
            r = requests.post(url, params)
            LOGGER.info("webhook deployed, got %s as reposnse",
                        r)
        except requests.exceptions.RequestException, err:
            LOGGER.warning("cannot deploy webhook %s: %s",
                           rcp, err, exc_info=err)

def notify(config, resource, run, last_run_success):
    """execute a notification"""

    status_changed = False

    this_run_success = run.success

    if last_run_success and not this_run_success:
        result = gettext('Failing')
    elif not last_run_success and this_run_success:
        result = gettext('Fixed')
    elif not last_run_success and not this_run_success:
        result = gettext('Still Failing')
    elif last_run_success and this_run_success:
        result = gettext('Passing')

    if result != gettext('Passing'):
        status_changed = True

    # Check if still 'Still Failing' result should be notified
    if result == gettext('Still Failing') \
            and not config['GHC_NOTIFICATIONS_VERBOSITY']:
        # Receive just 'Failing' and 'Fixed' notifications
        status_changed = False

    if not status_changed:
        return

    print('Notifying: status changed: result=%s' % result)

    # run all channels, actual recipients will be filtered there
    for chann_handler in (do_email, do_webhook,):
        try:
            chann_handler(config, resource, run, status_changed, result)
        except Exception, err:
            LOGGER.warning("couldn't run notification for %s: %s",
                           chann_handler.func_name, err, exc_info=err)
<commit_msg>Fix connecting _localhost_ by calling `connect()` on an already initialized smtplib object, introduced in 5ef60b331bf5495e557b6c1c48e5fad5893edeb7. According to the docs, `smtplib.SMTP()` calls the `connect` method for the given host. Calling `connect` again but without any parameters will fall back to smtplibs default, which is localhost:25.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

# =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2014 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

from email.mime.text import MIMEText
import types
import email.utils
import logging
import smtplib
import json

import requests
from flask_babel import gettext
from util import render_template2

LOGGER = logging.getLogger(__name__)

def do_email(config, resource, run, status_changed, result):
    # List of global email addresses to notify, may be list or
    # comma-separated str "To" needs comma-separated list,
    # while sendmail() requires list...

    if isinstance(config['GHC_NOTIFICATIONS_EMAIL'], types.StringTypes):
        config['GHC_NOTIFICATIONS_EMAIL'] = \
            config['GHC_NOTIFICATIONS_EMAIL'].split(',')

    # this should always be a list
    global_notifications = config['GHC_NOTIFICATIONS_EMAIL'] or []
    if not isinstance(global_notifications, (list, tuple, set,)):
        raise TypeError("Cannot use {} as list of emails".format(
                                       type(global_notifications)))

    notifications_email = global_notifications +\
        resource.get_recipients('email')

    if not notifications_email:
        LOGGER.warning("No emails for notification set for resource %s",
                       resource.identifier)
        return

    template_vars = {
        'result': result,
        'config': config,
        'resource': resource,
        'run': run
    }

    msgbody = render_template2('notification_email.txt', template_vars)
    msg = MIMEText(msgbody)
    msg['From'] = email.utils.formataddr((config['GHC_SITE_TITLE'],
                                          config['GHC_ADMIN_EMAIL']))
    msg['To'] = ','.join(notifications_email)
    msg['Subject'] = '[%s] %s: %s' % (config['GHC_SITE_TITLE'],
                                      result, resource.title)

    if not config.get('GHC_SMTP') or not\
        (any([config['GHC_SMTP'][k] for k in ('port',
                                              'server',
                                              'username',
                                              'password',)])):

        LOGGER.warning("No SMTP configuration. Not sending to %s",
                       notifications_email)
        print(msg.as_string())
        return

    server = smtplib.SMTP(config['GHC_SMTP']['server'],
                          config['GHC_SMTP']['port'])

    if config['DEBUG']:
        server.set_debuglevel(True)

    try:
        if config['GHC_SMTP']['tls']:
            server.starttls()
    except Exception, err:
        LOGGER.exception("Cannot connect to smtp: %s[:%s]: %s",
                         config['GHC_SMTP']['server'],
                         config['GHC_SMTP']['port'],
                         err,
                         exc_info=err)
        return
    try:
        server.login(config['GHC_SMTP']['username'],
                     config['GHC_SMTP']['password'])
    except Exception, err:
        LOGGER.exception("Cannot log in to smtp: %s", err,
                         exc_info=err)
    try:
        server.sendmail(config['GHC_ADMIN_EMAIL'],
                        config['GHC_NOTIFICATIONS_EMAIL'],
                        msg.as_string())
    except Exception as err:
        LOGGER.exception(str(err), exc_info=err)
    finally:
        server.quit()

def _parse_line(_line):
    try:
        k, v = _line.split('=', 1)
        return {k: v}
    except (IndexError, ValueError,):
        raise ValueError("Invalid line: {}".format(_line))

def _parse_webhook_location(value):
    """
    Parse Recipient.location and returns tuple of url and params

    location should be in form

    URL

    PAYLOAD

    where PAYLOAD is a list of fields and values in form

    FIELD_NAME=FIELD_VALUE

    alternatively, it can be dictionary serialized as json
    """
    if not value.strip():
        raise ValueError("No payload")
    value = value.strip()
    url = None
    params = {}
    lines = value.splitlines()

    for idx, line in enumerate(lines):
        if idx == 0:
            url = line
        elif idx == 1:
            if line.strip():
                raise ValueError("Second line should be empty")
        elif idx == 2:
            try:
                params = json.loads('\n'.join(lines[2:]))
                break
            except (TypeError, ValueError,):
                params.update(_parse_line(line))
        else:
            params.update(_parse_line(line))

    if url is None:
        raise ValueError("Cannot parse url")
    return url, params,

def do_webhook(config, resource, run, status_changed, result):
    """
    Process webhook recipients for resource

    location should be in format:

    URL

    [PAYLOAD]

    There's blank line between URL and PAYLOAD. PAYLOAD
    should be either json or list of field=value items
    in each line.

    Webhook's request is POST send to url with payload containing
    PAYLOAD
    and fields:

    ghc.result=(result of test)
    ghc.resource.url=(url of resource)
    ghc.resource.title=(title of resource)

    """
    recipients = resource.get_recipients('webhook')
    if not recipients:
        return
    for rcp in recipients:
        try:
            url, params = _parse_webhook_location(rcp)
        except ValueError, err:
            LOGGER.warning("Cannot send to {}: {}"
                           .format(rcp, err), exc_info=err)

        resource_view = '{}/resource/{}'.format(
                                            config['GHC_SITE_URL'],
                                            resource.identifier)

        params['ghc.result'] = result
        params['ghc.resource.url'] = resource.url
        params['ghc.resource.title'] = resource.title
        params['ghc.resource.type'] = resource.resource_type
        params['ghc.resource.view'] = resource_view

        try:
            r = requests.post(url, params)
            LOGGER.info("webhook deployed, got %s as reposnse",
                        r)
        except requests.exceptions.RequestException, err:
            LOGGER.warning("cannot deploy webhook %s: %s",
                           rcp, err, exc_info=err)

def notify(config, resource, run, last_run_success):
    """execute a notification"""

    status_changed = False

    this_run_success = run.success

    if last_run_success and not this_run_success:
        result = gettext('Failing')
    elif not last_run_success and this_run_success:
        result = gettext('Fixed')
    elif not last_run_success and not this_run_success:
        result = gettext('Still Failing')
    elif last_run_success and this_run_success:
        result = gettext('Passing')

    if result != gettext('Passing'):
        status_changed = True

    # Check if still 'Still Failing' result should be notified
    if result == gettext('Still Failing') \
            and not config['GHC_NOTIFICATIONS_VERBOSITY']:
        # Receive just 'Failing' and 'Fixed' notifications
        status_changed = False

    if not status_changed:
        return

    print('Notifying: status changed: result=%s' % result)

    # run all channels, actual recipients will be filtered there
    for chann_handler in (do_email, do_webhook,):
        try:
            chann_handler(config, resource, run, status_changed, result)
        except Exception, err:
            LOGGER.warning("couldn't run notification for %s: %s",
                           chann_handler.func_name, err, exc_info=err)
<|endoftext|>


<commit_before>import unittest

import numpy as np

from vaspy.thermo import OsziCar

class TestOsziCar(unittest.TestCase):

    def setUp(self):
        #create an instance of OSZICAR file
        self.x = OsziCar('../testdata/OSZICAR')

    def test_attrs(self):
        "Make sure load() effects."
        self.assertTrue(hasattr(self.x, 'E0'))
        self.assertTrue(hasattr(self.x, 'dE'))
        self.assertTrue(hasattr(self.x, 'F'))
        self.assertTrue(hasattr(self.x, 'step'))
        self.assertTrue(hasattr(self.x, 'mag'))
        self.assertTrue(hasattr(self.x, 'content'))

        #should raise an exception for an AttributeError
        self.assertRaises(AttributeError)

    def test_esort(self):
        "Make sure the esort() effects"
        srted = self.x.esort(2)
        shouldbe = np.array([(-101.21186, 326), (-101.21116, 324)],
                            dtype=[('E0', '<f8'), ('step', '<i4')])
        self.assertTrue((srted == shouldbe).all())

if __name__ == '__main__':
    suite = unittest.TestLoader().loadTestsFromTestCase(TestOsziCar)
    unittest.TextTestRunner(verbosity=2).run(suite)
<commit_msg>check all vars in test_attrs()<commit_after># -*- coding:utf-8 -*-
'''
    OsziCar类单元测试.
'''
import unittest

import numpy as np

from vaspy.thermo import OsziCar

class TestOsziCar(unittest.TestCase):

    def setUp(self):
        #create an instance of OSZICAR file
        self.x = OsziCar('../testdata/OSZICAR')

    def test_attrs(self):
        "Make sure load() effects"
        for var in self.x.vars:
            self.assertTrue(hasattr(self.x, var))

        #should raise an exception for an AttributeError
        self.assertRaises(AttributeError)

    def test_esort(self):
        "Make sure the esort() effects"
        srted = self.x.esort(2)
        shouldbe = np.array([(-101.21186, 326), (-101.21116, 324)],
                            dtype=[('E0', '<f8'), ('step', '<i4')])
        self.assertTrue((srted == shouldbe).all())

if __name__ == '__main__':
    suite = unittest.TestLoader().loadTestsFromTestCase(TestOsziCar)
    unittest.TextTestRunner(verbosity=2).run(suite)
<|endoftext|>


<commit_before>from collections import deque
from untwisted.event import DUMPED, WRITE
from untwisted.errors import SEND_ERR_CODE
from untwisted.event import CLOSE, READ_ERR
import ssl
import socket

class Dump:
    def process_error(self, ssock, excpt):
        if not excpt.args[0] in SEND_ERR_CODE:
            if not ssock.dead:
                ssock.drive(CLOSE, excpt)

class DumpStr(Dump):
    __slots__ = 'data'

    def __init__(self, data=b''):
        self.data = memoryview(data)

    def process(self, ssock):
        try:
            size = ssock.send(self.data)
        except socket.error as excpt:
            self.process_error(ssock, excpt)
        else:
            self.data = self.data[size:]

    def __bool__(self):
        return bool(self.data)

class DumpFile(DumpStr):
    __slots__ = 'fd'
    BLOCK     = 1024 * 124

    def __init__(self, fd):
        self.fd = fd
        DumpStr.__init__(self)

    def process(self, ssock):
        if not self.data:
            self.process_file(ssock)
        DumpStr.process(self, ssock)

    def process_file(self, ssock):
        try:
            data = self.fd.read(DumpFile.BLOCK)
        except IOError as excpt:
            ssock.drive(READ_ERR, excpt)
        else:
            self.data = memoryview(data)

class DumpStrSSL(DumpStr):
    def process(self, ssock):
        # When ssl.SSLEOFError happens it shouldnt spawn CLOSE
        # because there may exist data to be read.
        # If it spawns CLOSE the socket is closed and no data
        # can be read from.
        try:
            size = ssock.send(self.data)
        except ssl.SSLWantReadError as excpt:
            pass
        except ssl.SSLWantWriteError as excpt:
            pass
        except ssl.SSLEOFError as excpt:
            pass
        except ssl.SSLError as excpt:
            ssock.drive(CLOSE, excpt)
        except socket.error as excpt:
            self.process_error(ssock, excpt)
        else:
            self.data = self.data[size:]

class DumpFileSSL(DumpStrSSL, DumpFile):
    pass

class SockWriter:
    """
    Used to write data through a socket asynchronously.
    """

    def __init__(self, ssock):
        """
        """

        self.queue = deque()
        self.data  = None
        ssock.dump = self.dump
        self.ssock = ssock
        ssock.dumpfile = self.dumpfile

    def update(self, ssock):
        """
        """

        if not self.data:
            self.process_queue(ssock)

        self.data.process(ssock)

    def process_queue(self, ssock):
        try:
            self.data = self.queue.popleft()
        except IndexError:
            self.stop()

    def stop(self):
        self.ssock.del_map(WRITE, self.update)
        self.ssock.drive(DUMPED)

    def start(self):
        if not self.queue:
            self.ssock.add_map(WRITE, self.update)

    def dump(self, data):
        """
        Send data through a SuperSocket instance.
        """

        self.start()
        dump = DumpStr(data)
        self.queue.append(dump)

    def dumpfile(self, fd):
        """
        Dump a file through a SuperSocket instance.
        """

        self.start()
        dump = DumpFile(fd)
        self.queue.append(dump)

class SockWriterSSL(SockWriter):
    """
    Used to dump data through a SuperSocket instance.
    """

    def dump(self, data):
        self.start()
        data = DumpStrSSL(data)
        self.queue.append(data)

    def dumpfile(self, fd):
        self.start()
        fd = DumpFileSSL(fd)
        self.queue.append(fd)

<commit_msg>Simplifying imports.<commit_after>from collections import deque
from untwisted.event import DUMPED, WRITE, CLOSE, READ_ERR
from untwisted.errors import SEND_ERR_CODE
import ssl
import socket

class Dump:
    def process_error(self, ssock, excpt):
        if not excpt.args[0] in SEND_ERR_CODE:
            if not ssock.dead:
                ssock.drive(CLOSE, excpt)

class DumpStr(Dump):
    __slots__ = 'data'

    def __init__(self, data=b''):
        self.data = memoryview(data)

    def process(self, ssock):
        try:
            size = ssock.send(self.data)
        except socket.error as excpt:
            self.process_error(ssock, excpt)
        else:
            self.data = self.data[size:]

    def __bool__(self):
        return bool(self.data)

class DumpFile(DumpStr):
    __slots__ = 'fd'
    BLOCK     = 1024 * 124

    def __init__(self, fd):
        self.fd = fd
        DumpStr.__init__(self)

    def process(self, ssock):
        if not self.data:
            self.process_file(ssock)
        DumpStr.process(self, ssock)

    def process_file(self, ssock):
        try:
            data = self.fd.read(DumpFile.BLOCK)
        except IOError as excpt:
            ssock.drive(READ_ERR, excpt)
        else:
            self.data = memoryview(data)

class DumpStrSSL(DumpStr):
    def process(self, ssock):
        # When ssl.SSLEOFError happens it shouldnt spawn CLOSE
        # because there may exist data to be read.
        # If it spawns CLOSE the socket is closed and no data
        # can be read from.
        try:
            size = ssock.send(self.data)
        except ssl.SSLWantReadError as excpt:
            pass
        except ssl.SSLWantWriteError as excpt:
            pass
        except ssl.SSLEOFError as excpt:
            pass
        except ssl.SSLError as excpt:
            ssock.drive(CLOSE, excpt)
        except socket.error as excpt:
            self.process_error(ssock, excpt)
        else:
            self.data = self.data[size:]

class DumpFileSSL(DumpStrSSL, DumpFile):
    pass

class SockWriter:
    """
    Used to write data through a socket asynchronously.
    """

    def __init__(self, ssock):
        """
        """

        self.queue = deque()
        self.data  = None
        ssock.dump = self.dump
        self.ssock = ssock
        ssock.dumpfile = self.dumpfile

    def update(self, ssock):
        """
        """

        if not self.data:
            self.process_queue(ssock)

        self.data.process(ssock)

    def process_queue(self, ssock):
        try:
            self.data = self.queue.popleft()
        except IndexError:
            self.stop()

    def stop(self):
        self.ssock.del_map(WRITE, self.update)
        self.ssock.drive(DUMPED)

    def start(self):
        if not self.queue:
            self.ssock.add_map(WRITE, self.update)

    def dump(self, data):
        """
        Send data through a SuperSocket instance.
        """

        self.start()
        dump = DumpStr(data)
        self.queue.append(dump)

    def dumpfile(self, fd):
        """
        Dump a file through a SuperSocket instance.
        """

        self.start()
        dump = DumpFile(fd)
        self.queue.append(dump)

class SockWriterSSL(SockWriter):
    """
    Used to dump data through a SuperSocket instance.
    """

    def dump(self, data):
        self.start()
        data = DumpStrSSL(data)
        self.queue.append(data)

    def dumpfile(self, fd):
        self.start()
        fd = DumpFileSSL(fd)
        self.queue.append(fd)

<|endoftext|>


<commit_before>#!/usr/bin/env python

from copy import deepcopy
import geojson
from pathlib import Path
import shapely.geometry as sgeom

from repo import BroadcastAreasRepository

package_path = Path(__file__).resolve().parent

def simplify_polygon(series):
    polygon, *_holes = series  # discard holes

    approx_metres_to_degree = 111320
    desired_resolution_metres = 10
    simplify_degrees = desired_resolution_metres / approx_metres_to_degree

    simplified_polygon = None
    num_polys = len(polygon)
    while True:
        simplified_polygon = sgeom.LineString(polygon)
        simplified_polygon = simplified_polygon.buffer(simplify_degrees)
        simplified_polygon = simplified_polygon.simplify(simplify_degrees)
        simplified_polygon = [[c[0], c[1]] for c in simplified_polygon.exterior.coords]

        num_polys = len(simplified_polygon)
        simplify_degrees *= 1.5
        if num_polys <= 125:
            break

    return [simplified_polygon]

def simplify_geometry(feature):
    if feature["type"] == "Polygon":
        feature["coordinates"] = simplify_polygon(feature["coordinates"])
        return feature
    elif feature["type"] == "MultiPolygon":
        feature["coordinates"] = [
            simplify_polygon(polygon)
            for polygon in feature["coordinates"]
        ]
        return feature
    else:
        raise Exception("Unknown type: {}".format(feature["type"]))

repo = BroadcastAreasRepository()

repo.delete_db()
repo.create_tables()

simple_datasets = [
    ("Countries", "ctry19nm"),
    ("Regions of England", "rgn18nm"),
    ("Counties and Unitary Authorities in England and Wales", "ctyua16nm"),
]
for dataset_name, name_field in simple_datasets:
    filepath = package_path / "{}.geojson".format(dataset_name)
    dataset_geojson = geojson.loads(filepath.read_text())

    repo.insert_broadcast_area_library(dataset_name)

    for feature in dataset_geojson["features"]:
        f_name = feature["properties"][name_field]

        simple_feature = deepcopy(feature)
        simple_feature["geometry"] = simplify_geometry(simple_feature["geometry"])
        repo.insert_broadcast_areas([[f_name, dataset_name, feature, simple_feature]])

# https://geoportal.statistics.gov.uk/datasets/wards-may-2020-boundaries-uk-bgc
# Converted to geojson manually from SHP because of GeoJSON download limits
wards_filepath = package_path / "Electoral Wards May 2020.geojson"

# http://geoportal.statistics.gov.uk/datasets/ward-to-westminster-parliamentary-constituency-to-local-authority-district-december-2019-lookup-in-the-united-kingdom/data
las_filepath = package_path / "Electoral Wards and Local Authorities 2020.geojson"

ward_code_to_la_mapping = {
    f["properties"]["WD19CD"]: f["properties"]["LAD19NM"]
    for f in geojson.loads(las_filepath.read_text())["features"]
}

dataset_name = "Electoral Wards of the United Kingdom"
repo.insert_broadcast_area_library(dataset_name)

areas_to_add = []

for f in geojson.loads(wards_filepath.read_text())["features"]:

    ward_code = f["properties"]["wd20cd"]
    ward_name = f["properties"]["wd20nm"]

    try:
        la_name = ward_code_to_la_mapping[ward_code]

        f_name = "{} - {}".format(la_name, ward_name)

        sf = deepcopy(f)
        sf["geometry"] = simplify_geometry(sf["geometry"])

        areas_to_add.append([f_name, dataset_name, f, sf])

    except KeyError:
        print("Skipping", ward_code, ward_name)  # noqa: T001

repo.insert_broadcast_areas(areas_to_add)
<commit_msg>broadcast-areas: create broadcast areas in groups<commit_after>#!/usr/bin/env python

from copy import deepcopy
import sqlite3
import geojson
from pathlib import Path
import shapely.geometry as sgeom

from notifications_utils.safe_string import make_string_safe_for_id

from repo import BroadcastAreasRepository

package_path = Path(__file__).resolve().parent

def convert_shape_to_feature(shape):
    return {
        "type": "Feature",
        "properties": {},
        "geometry": sgeom.mapping(shape),
    }

def simplify_polygon(series):
    polygon, *_holes = series  # discard holes

    approx_metres_to_degree = 111320
    desired_resolution_metres = 10
    simplify_degrees = desired_resolution_metres / approx_metres_to_degree

    simplified_polygon = None
    num_polys = len(polygon)
    while True:
        simplified_polygon = sgeom.LineString(polygon)
        simplified_polygon = simplified_polygon.buffer(simplify_degrees)
        simplified_polygon = simplified_polygon.simplify(simplify_degrees)
        simplified_polygon = [[c[0], c[1]] for c in simplified_polygon.exterior.coords]

        num_polys = len(simplified_polygon)
        simplify_degrees *= 1.75
        if num_polys <= 125:
            break

    return [simplified_polygon]

def simplify_geometry(feature):
    if feature["type"] == "Polygon":
        feature["coordinates"] = simplify_polygon(feature["coordinates"])
        return feature
    elif feature["type"] == "MultiPolygon":
        feature["coordinates"] = [
            simplify_polygon(polygon)
            for polygon in feature["coordinates"]
        ]
        return feature
    else:
        raise Exception("Unknown type: {}".format(feature["type"]))

repo = BroadcastAreasRepository()

repo.delete_db()
repo.create_tables()

simple_datasets = [
    ("Countries", "ctry19cd", "ctry19nm"),
    ("Regions of England", "rgn18cd", "rgn18nm"),
    ("Counties and Unitary Authorities in England and Wales", "ctyua16cd", "ctyua16nm"),
]
for dataset_name, id_field, name_field in simple_datasets:
    filepath = package_path / "{}.geojson".format(dataset_name)

    dataset_id = make_string_safe_for_id(dataset_name)
    dataset_geojson = geojson.loads(filepath.read_text())

    repo.insert_broadcast_area_library(dataset_id, dataset_name, False)

    for feature in dataset_geojson["features"]:
        f_id = dataset_id + "-" + feature["properties"][id_field]
        f_name = feature["properties"][name_field]

        simple_feature = deepcopy(feature)
        simple_feature["geometry"] = simplify_geometry(simple_feature["geometry"])

        repo.insert_broadcast_areas([[
            f_id, f_name,
            dataset_id, None,
            feature, simple_feature,
        ]])

# https://geoportal.statistics.gov.uk/datasets/wards-may-2020-boundaries-uk-bgc
# Converted to geojson manually from SHP because of GeoJSON download limits
wards_filepath = package_path / "Electoral Wards May 2020.geojson"

# http://geoportal.statistics.gov.uk/datasets/ward-to-westminster-parliamentary-constituency-to-local-authority-district-december-2019-lookup-in-the-united-kingdom/data
las_filepath = package_path / "Electoral Wards and Local Authorities 2020.geojson"

ward_code_to_la_mapping = {
    f["properties"]["WD19CD"]: f["properties"]["LAD19NM"]
    for f in geojson.loads(las_filepath.read_text())["features"]
}
ward_code_to_la_id_mapping = {
    f["properties"]["WD19CD"]: f["properties"]["LAD19CD"]
    for f in geojson.loads(las_filepath.read_text())["features"]
}

dataset_name = "Electoral Wards of the United Kingdom"
dataset_id = make_string_safe_for_id(dataset_name)
repo.insert_broadcast_area_library(dataset_id, dataset_name, True)

for f in geojson.loads(las_filepath.read_text())["features"]:
    group_id = dataset_id + "-" + f["properties"]["LAD19CD"]
    group_name = f["properties"]["LAD19NM"]

    try:
        repo.insert_broadcast_area_library_group(group_id, group_name, dataset_id)
    except sqlite3.IntegrityError:
        # Already exists
        pass

areas_to_add = []

for f in geojson.loads(wards_filepath.read_text())["features"]:
    ward_code = f["properties"]["wd20cd"]
    ward_name = f["properties"]["wd20nm"]
    ward_id = dataset_id + "-" + ward_code

    try:
        la_id = dataset_id + "-" + ward_code_to_la_id_mapping[ward_code]
        la_name = ward_code_to_la_mapping[ward_code]

        sf = deepcopy(f)
        sf["geometry"] = simplify_geometry(sf["geometry"])

        areas_to_add.append([
            ward_id, ward_name,
            dataset_id, la_id,
            f, sf
        ])

    except KeyError:
        print("Skipping", ward_code, ward_name)  # noqa: T001

repo.insert_broadcast_areas(areas_to_add)
areas_to_add = []

for group_id, group_name in repo.get_all_groups_for_library(dataset_id):
    areas = repo.get_all_areas_for_group(group_id)

    features = [geojson.loads(f) for _, _, f, _ in areas]
    shapes = [sgeom.shape(f["geometry"]) for f in features]

    if len(shapes) == 0:
        continue

    print(group_id, group_name, "union")
    shape = shapes[0]
    for other in shapes[1:]:
        shape = shape.buffer(0).union(other.buffer(0))

    feature = convert_shape_to_feature(shape)
    simple_feature = deepcopy(feature)
    simple_feature["geometry"] = simplify_geometry(simple_feature["geometry"])

    areas_to_add.append([
        group_id, group_name,
        dataset_id, None,
        feature, simple_feature
    ])

repo.insert_broadcast_areas(areas_to_add)
<|endoftext|>


<commit_before>"""Main entry point for untz."""
import argparse
import logging
from multiprocessing import cpu_count
import threading
import os
from queue import Queue
import signal
import sys

from . import __version__
from .encoder import apply_gain, encode_file
from .utils import preflight_checks, recursive_file_search

LOGGER = logging.getLogger(__name__)
threads = [] # pylint: disable=C0103
event_flag = threading.Event() # pylint: disable=C0103

def get_args():
    """Parse and return arguments."""
    parser = argparse.ArgumentParser(description='Convert FLACs to OGGs and '
                                                 'sort into a sensible folder hierarchy.')
    parser.add_argument('-b', '--base',
                        dest='base_dir',
                        help='Base directory to store output files to.', required=True)
    parser.add_argument('-n', '--names',
                        default='%n - %t',
                        dest='pattern',
                        help='Produce filenames as this string, with %%g, %%a, %%l, %%n, %%t, %%d '
                             'replaced by genre, artist, album, track number, title, '
                             'and date, respectively. Also, %%%% gives a literal %%. '
                             'Defaults to %%n - %%t.')
    parser.add_argument('-p', '--passthrough',
                        dest='ogg_cli_parameters',
                        help='Additional CLI arguments to pass through to oggenc.',
                        default='')
    parser.add_argument('-q', '--quality',
                        dest='quality',
                        type=float,
                        default=10,
                        help='Sets encoding quality to n, between -1 (low) and 10 (high). '
                             'Fractional quality levels such as 2.5 are permitted.')
    parser.add_argument('-r', '--replaygain',
                        dest='replaygain',
                        action='store_true',
                        help='Apply replaygain tags with vorbisgain.')
    parser.add_argument('-t', '--threads',
                        dest='threads',
                        type=int,
                        default=cpu_count(),
                        help='Thread pool size to spawn.')
    parser.add_argument('-v', '--verbose',
                        action='store_true',
                        dest='verbose',
                        help='Set verbose logging.')
    parser.add_argument('-V', '--version',
                        action='version',
                        version=__version__,
                        help='Print version and exit.')
    parser.add_argument(dest='inputs', help='List of folder/file inputs.', nargs='+')

    return parser.parse_args()

def _encode_on_filter(args, queue):
    while event_flag.is_set():
        file_entry = queue.get()

        # FIFO queue so we can do this
        if file_entry is None:
            LOGGER.info("PID %d received poison signal. Breaking.", os.getpid())
            break

        if file_entry.lower().endswith('.flac'):
            encode_file(file_entry,
                        args.base_dir.rstrip('/'),
                        args.pattern,
                        args.quality,
                        args.ogg_cli_parameters)

def _shutdown(signum, _):
    # Finish after current transcode is done per thread.
    event_flag.clear()

    for thread in threads:
        thread.join()

    # Abnormal exit
    sys.exit(signum)

def main():
    """Main logic for untz."""
    args = get_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    preflight_checks()

    queue = Queue()

    LOGGER.info('Starting %d threads.', args.threads)

    for i in range(args.threads):
        threads.append(threading.Thread(target=_encode_on_filter,
                                        args=(args, queue)))

    event_flag.set()

    for thread in threads:
        thread.start()

    # Register signal
    signal.signal(signal.SIGINT, _shutdown)
    signal.signal(signal.SIGTERM, _shutdown)

    for i in args.inputs:
        if os.path.isdir(i):
            for file_entry in recursive_file_search(i):
                LOGGER.debug('Encoding "%s"', file_entry)
                queue.put_nowait(file_entry)
        else:
            LOGGER.debug('Encoding "%s"', i)
            queue.put(i)

    # Send stop signal
    LOGGER.debug('Sending poison signal.')
    for _ in range(args.threads):
        queue.put(None)

    LOGGER.info('Waiting for processes to die.')
    for thread in threads:
        thread.join()

    if args.replaygain:
        apply_gain(args.base_dir)

    LOGGER.warning('Program exiting now.')
<commit_msg>Use ThreadPool<commit_after>"""Main entry point for untz."""
import argparse
from concurrent.futures import ThreadPoolExecutor
import logging
import os

from . import __version__
from .encoder import apply_gain, encode_file
from .utils import preflight_checks, recursive_file_search

ARGS = get_args()
LOGGER = logging.getLogger(__name__)
threads = [] # pylint: disable=C0103

def get_args():
    """Parse and return arguments."""
    parser = argparse.ArgumentParser(description='Convert FLACs to OGGs and '
                                                 'sort into a sensible folder hierarchy.')
    parser.add_argument('-b', '--base',
                        dest='base_dir',
                        help='Base directory to store output files to.', required=True)
    parser.add_argument('-n', '--names',
                        default='%n - %t',
                        dest='pattern',
                        help='Produce filenames as this string, with %%g, %%a, %%l, %%n, %%t, %%d '
                             'replaced by genre, artist, album, track number, title, '
                             'and date, respectively. Also, %%%% gives a literal %%. '
                             'Defaults to %%n - %%t.')
    parser.add_argument('-p', '--passthrough',
                        dest='ogg_cli_parameters',
                        help='Additional CLI arguments to pass through to oggenc.',
                        default='')
    parser.add_argument('-q', '--quality',
                        dest='quality',
                        type=float,
                        default=10,
                        help='Sets encoding quality to n, between -1 (low) and 10 (high). '
                             'Fractional quality levels such as 2.5 are permitted.')
    parser.add_argument('-r', '--replaygain',
                        dest='replaygain',
                        action='store_true',
                        help='Apply replaygain tags with vorbisgain.')
    parser.add_argument('-t', '--threads',
                        dest='threads',
                        type=int,
                        default=cpu_count(),
                        help='Thread pool size to spawn.')
    parser.add_argument('-v', '--verbose',
                        action='store_true',
                        dest='verbose',
                        help='Set verbose logging.')
    parser.add_argument('-V', '--version',
                        action='version',
                        version=__version__,
                        help='Print version and exit.')
    parser.add_argument(dest='inputs', help='List of folder/file inputs.', nargs='+')

    return parser.parse_args()

def _encode_on_filter(args, queue):
    if file_entry.lower().endswith('.flac'):
        encode_file(file_entry,
                    ARGS.base_dir.rstrip('/'),
                    ARGS.pattern,
                    ARGS.quality,
                    ARGS.ogg_cli_parameters)

def main():
    """Main logic for untz."""

    if ARGS.verbose:
        logging.basicConfig(level=logging.DEBUG)

    preflight_checks()

    LOGGER.info('Starting %d threads.', ARGS.threads)
    with ThreadPoolExecutor(max_workers=ARGS.threads) as tpe:
        for i in ARGS.inputs:
            if os.path.isdir(i):
                tpe.map(_encode_on_filter, recursive_file_search(i))
            else:
                tpe.submit(_encode_on_filter, i)

    if ARGS.replaygain:
        apply_gain(ARGS.base_dir)

    LOGGER.warning('Program exiting now.')
<|endoftext|>


<commit_before>"""
Functions for querying the Solr server.

    results = query_solr( db, { 'q': 'obama' } )

    sentences = results['response']['docs']
    for sentence in sentences:
        print(f"Found sentence ID: {sentence['story_sentences_id']}")

More information about Solr integration at docs/solr.markdown.
"""

import abc
import copy
import re
from typing import Union, List, Dict, Any, Optional

from mediawords.db import DatabaseHandler
from mediawords.solr.params import SolrParams
from mediawords.solr.request import solr_request
from mediawords.util.log import create_logger
from mediawords.util.parse_json import decode_json
from mediawords.util.perl import decode_object_from_bytes_if_needed

log = create_logger(__name__)

class _AbstractSolrException(Exception, metaclass=abc.ABCMeta):
    """Abstract .solr exception."""
    pass

class _AbstractSolrInternalErrorException(_AbstractSolrException):
    """Internal code error (most likely a bug in the code)."""
    pass

class McUppercaseBooleanOperatorsInvalidTypeException(_AbstractSolrInternalErrorException):
    """Exception thrown when weird stuff gets passed to _uppercase_boolean_operators()."""
    pass

class McQuerySolrInternalErrorException(_AbstractSolrInternalErrorException):
    """Exception thrown when query_solr() receives something that it didn't expect."""
    pass

class _AbstractSolrInvalidQueryException(_AbstractSolrException):
    """Invalid Solr query."""
    pass

class McInsertCollectionMediaIDsInvalidQueryException(_AbstractSolrInvalidQueryException):
    """Invalid Solr query encountered in _insert_collection_media_ids()."""
    pass

class McQuerySolrInvalidQueryException(_AbstractSolrInvalidQueryException):
    """Invalid Solr query in query_solr()."""
    pass

class McQuerySolrRangeQueryException(_AbstractSolrInvalidQueryException):
    """Exception thrown on range queries which are not allowed."""
    pass

def _uppercase_boolean_operators(query: Union[List[str], str]) -> Union[List[str], str]:
    """
    Convert any "and", "or", or "not" operations in the argument to uppercase.

    If the argument is a list, call ourselves on all elements of the list.
    """
    query = decode_object_from_bytes_if_needed(query)

    if not query:
        return query

    def upper_repl(match) -> str:
        return match.group(1).upper()

    if isinstance(query, list):
        query = [_uppercase_boolean_operators(_) for _ in query]
    elif isinstance(query, str):
        query = re.sub(r'\b(and|or|not)\b', upper_repl, query)
    else:
        raise McUppercaseBooleanOperatorsInvalidTypeException(f"Invalid query type: {query}")

    return query

def _insert_collection_media_ids(db: DatabaseHandler, q: str) -> str:
    """
    Transform any "tags_id_media:" or "collections_id:" clauses into "media_id:" clauses with the "media_ids" that
    corresponds to the given tags.
    """

    q = decode_object_from_bytes_if_needed(q)

    def get_media_ids_clause(match) -> str:
        """Given the argument of "tags_id_media:" or "collections_id:" clause, return the corresponding "media_ids"."""
        arg = match.group(2)

        tags_ids = []
        if re.search(r'^\d+', arg):
            tags_ids.append(arg)

        else:
            parens_match = re.search(r'^\((.*)\)$', arg)

            if parens_match:
                parens = parens_match.group(1)

                parens = re.sub(r'or', ' ', parens, flags=re.IGNORECASE)
                parens = parens.strip()

                if re.search(r'[^\d\s]', parens):
                    raise McInsertCollectionMediaIDsInvalidQueryException((
                        f'Only "or" clauses allowed inside "tags_id_media:" or "collections_id:" clauses: {parens}; '
                        f'full match: {arg}'
                    ))

                for tags_id in re.split(r'\s+', parens):
                    tags_ids.append(tags_id)

            elif re.search(r'^\[', arg):
                raise McQuerySolrRangeQueryException(
                    'Range queries not allowed for "tags_id_media:" or "collections_id:" clauses'
                )

            else:
                raise McInsertCollectionMediaIDsInvalidQueryException(
                    f'Unrecognized format of "tags_id_media:" or "collections_id:" clause: {arg}'
                )

        media_ids = db.query("""
            SELECT media_id
            FROM media_tags_map
            WHERE tags_id IN %(tags_ids)s
            ORDER BY media_id
        """, {'tags_ids': tuple(tags_ids)}).flat()

        # Replace empty list with an id that will always return nothing from Solr
        if not media_ids:
            media_ids = [-1]

        media_clause = f"media_id:({' '.join([str(_) for _ in media_ids])})"

        return media_clause

    q = re.sub(r'(tags_id_media|collections_id):(\d+|\([^)]*\)|\[[^\]]*\])', get_media_ids_clause, q)

    return q

def query_solr(db: DatabaseHandler, params: SolrParams) -> Dict[str, Any]:
    """
    Execute a query on the Solr server using the given parameters. Return a maximum of 1 million sentences.

    The "params" argument is a dictionary of query parameters to Solr, detailed here:

        https://lucene.apache.org/solr/guide/6_6/common-query-parameters.html.

    The query ("params['q']") is transformed: lower case boolean operators are made uppercase to make Solr recognize
    them as boolean queries.

    Return decoded response in the format described here:

        https://lucene.apache.org/solr/guide/6_6/response-writers.html#ResponseWriters-JSONResponseWriter
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    if not params:
        raise McQuerySolrInternalErrorException('Parameters must be set.')

    if not isinstance(params, dict):
        raise McQuerySolrInternalErrorException('Parameters must be a dictionary.')

    params['wt'] = 'json'

    if 'rows' in params:
        params['rows'] = int(params['rows'])
    else:
        params['rows'] = 1000

    if 'df' not in params:
        params['df'] = 'text'

    params['rows'] = min(params['rows'], 10_000_000)

    if 'q' not in params:
        params['q'] = ''

    if 'fq' not in params:
        params['fq'] = []

    if not isinstance(params['fq'], list):
        params['fq'] = [params['fq']]

    if ':[' in params['q']:
        raise McQuerySolrRangeQueryException(
            "Range queries are not allowed in the main query. Please use a filter query instead for range queries."
        )

    # if params['q']:
    #     params['q'] = f"{{!complexphrase inOrder=false}} {params['q']}"

    params['q'] = _uppercase_boolean_operators(params['q'])
    params['fq'] = _uppercase_boolean_operators(params['fq'])

    params['q'] = _insert_collection_media_ids(db=db, q=params['q'])
    params['fq'] = [_insert_collection_media_ids(db=db, q=_) for _ in params['fq']]

    response_json = solr_request(
        path='select',
        params={},
        content=params,
        content_type='application/x-www-form-urlencoded; charset=utf-8',
    )

    try:
        response = decode_json(response_json)
    except Exception as ex:
        raise McQuerySolrInternalErrorException(f"Error parsing Solr JSON: {ex}\nJSON: {response_json}")

    if 'error' in response:
        raise McQuerySolrInvalidQueryException(f"Error received from Solr: {response_json}")

    return response

def _get_intersection_of_lists(lists: List[List[int]]) -> List[int]:
    """Given a list of lists, each of which points to a list of IDs, return an intersection between them."""
    lists = decode_object_from_bytes_if_needed(lists)

    if not lists:
        log.error("Lists are empty")
        return []

    intersection = set(lists[0])

    for cur_list in lists[1:]:
        intersection = intersection.intersection(set(cur_list))

    return sorted(list(intersection))

def _get_stories_ids_from_stories_only_q(q: str) -> Optional[List[int]]:
    """
    Transform the pseudoquery fields in the query and then run a simple pattern to detect queries that consists of one
    or more AND'ed "stories_id:..." clauses.

    For those cases, just return the story IDs list rather than running it through Solr.

    Return None if the query does not match.
    """
    q = decode_object_from_bytes_if_needed(q)

    if not q:
        return None

    q = re.sub(r'^\s*\(\s*(?P<inside_parens>.*)\s*\)\s*$', r'\g<inside_parens>', q)
    q = q.strip()

    if 'and' in q.lower():
        p = q.lower().index('and')
    else:
        p = -1

    if p > 0:
        q_a = q[:p - 1]
        a_stories_ids = _get_stories_ids_from_stories_only_q(q=q_a)
        if a_stories_ids is None:
            return None

        q_b = q[p + 4:]
        b_stories_ids = _get_stories_ids_from_stories_only_q(q=q_b)
        if b_stories_ids is None:
            return None

        r = _get_intersection_of_lists(lists=[a_stories_ids, b_stories_ids])
        return r

    story_match = re.search(r'^stories_id:(\d+)$', q)
    if story_match:
        stories_id = int(story_match.group(1))
        r = [stories_id]
        return r

    if re.search(r'^stories_id:\([\s\d]+\)$', q):
        stories_ids = []
        for story_match in re.findall(r'(\d+)', q):
            stories_id = int(story_match)
            stories_ids.append(stories_id)
        return stories_ids

    return None

def _get_stories_ids_from_stories_only_params(params: SolrParams) -> Optional[List[int]]:
    """
    Transform the pseudoquery fields in the "q" and "fq" params and then run a simple pattern to detect queries that
    consists of one or more AND'ed "stories_id:..." clauses in the "q" param and all "fq" params.

    Return None if either the "q" or any of the "fq" params do not match.
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    q = params.get('q', '')
    fqs = params.get('fq', [])
    start = params.get('start', None)
    rows = params.get('rows', None)

    if start is not None:
        start = int(start)
    if rows is not None:
        rows = int(rows)

    # Return None if there are any unrecognized params
    param_keys = set(list(params.keys()))
    if not param_keys.issubset({'q', 'fq', 'start', 'rows'}):
        log.error(f"Parameters have unrecognized keys; all keys: {param_keys}")
        return None

    if not q:
        log.error("'q' is unset.")
        return None

    stories_ids_lists = []

    if fqs:
        if not isinstance(fqs, list):
            fqs = [fqs]

        for fq in fqs:
            stories_ids = _get_stories_ids_from_stories_only_q(q=fq)
            if stories_ids is None:
                return None
            else:
                stories_ids_lists.append(stories_ids)

    # If there are stories_ids only "fqs" and a '*:*' "q", just use the "fqs"
    if stories_ids_lists and q == '*:*':
        r = _get_intersection_of_lists(lists=stories_ids_lists)

    # If there were no "fqs" and a '*:*' "q", return None
    elif q == '*:*':
        return None

    # Otherwise, combine "q" and "fqs"
    else:
        stories_ids = _get_stories_ids_from_stories_only_q(q=q)
        if stories_ids is None:
            return None

        r = _get_intersection_of_lists(lists=[stories_ids] + stories_ids_lists)

    if start is not None:
        r = r[start:]

    if rows is not None:
        r = r[:rows]

    return r

def get_solr_num_found(db: DatabaseHandler, params: SolrParams) -> int:
    """Execute the query and return only the number of documents found."""
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    params['rows'] = 0

    res = query_solr(db=db, params=params)

    num_found = res['response']['numFound']

    return num_found

def search_solr_for_stories_ids(db: DatabaseHandler, params: SolrParams) -> List[int]:
    """
    Return a list of all of the "stories_ids" that match the Solr query.

    Using Solr side grouping on the "stories_id" field.
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    stories_ids = _get_stories_ids_from_stories_only_params(params)
    if stories_ids:
        return stories_ids

    params['fl'] = 'stories_id'

    response = query_solr(db=db, params=params)

    stories_ids = [_['stories_id'] for _ in response['response']['docs']]

    return stories_ids

def search_solr_for_processed_stories_ids(db: DatabaseHandler,
                                          q: str,
                                          fq: Optional[Union[str, List[str]]],
                                          last_ps_id: int,
                                          num_stories: int,
                                          sort_by_random: bool = False) -> List[int]:
    """
    Return the first "num_stories" "processed_stories_id" that match the given query, sorted by "processed_stories_id"
    and with "processed_stories_id" greater than "last_ps_id".

    Returns at most "num_stories" stories.

    If "sort_by_random" is True, tell Solr to sort results by random order.
    """

    q = decode_object_from_bytes_if_needed(q)
    fq = decode_object_from_bytes_if_needed(fq)
    if isinstance(last_ps_id, bytes):
        last_ps_id = decode_object_from_bytes_if_needed(last_ps_id)
    if isinstance(num_stories, bytes):
        num_stories = decode_object_from_bytes_if_needed(num_stories)
    if isinstance(sort_by_random, bytes):
        sort_by_random = decode_object_from_bytes_if_needed(sort_by_random)

    last_ps_id = int(last_ps_id)
    num_stories = int(num_stories)
    sort_by_random = bool(int(sort_by_random))

    if not num_stories:
        return []

    if fq:
        if not isinstance(fq, list):
            fq = [fq]
    else:
        fq = []

    if last_ps_id:
        min_ps_id = last_ps_id + 1
        fq.append(f"processed_stories_id:[{min_ps_id} TO *]")

    params = {
        'q': q,
        'fq': fq,
        'fl': 'processed_stories_id',
        'rows': num_stories,
        'sort': 'random_1 asc' if sort_by_random else 'processed_stories_id asc',
    }

    response = query_solr(db=db, params=params)

    ps_ids = [_['processed_stories_id'] for _ in response['response']['docs']]

    return ps_ids

def search_solr_for_media_ids(db: DatabaseHandler, params: SolrParams) -> List[int]:
    """Return all of the media IDs that match the Solr query."""
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    params['fl'] = 'media_id'
    params['facet'] = 'true'
    params['facet.limit'] = 1_000_000
    params['facet.field'] = 'media_id'
    params['facet.mincount'] = 1
    params['rows'] = 0

    response = query_solr(db=db, params=params)

    counts = response['facet_counts']['facet_fields']['media_id']

    # Every second element (?)
    media_ids = counts[::2]

    return media_ids
<commit_msg>Return q if q is empty<commit_after>"""
Functions for querying the Solr server.

    results = query_solr( db, { 'q': 'obama' } )

    sentences = results['response']['docs']
    for sentence in sentences:
        print(f"Found sentence ID: {sentence['story_sentences_id']}")

More information about Solr integration at docs/solr.markdown.
"""

import abc
import copy
import re
from typing import Union, List, Dict, Any, Optional

from mediawords.db import DatabaseHandler
from mediawords.solr.params import SolrParams
from mediawords.solr.request import solr_request
from mediawords.util.log import create_logger
from mediawords.util.parse_json import decode_json
from mediawords.util.perl import decode_object_from_bytes_if_needed

log = create_logger(__name__)

class _AbstractSolrException(Exception, metaclass=abc.ABCMeta):
    """Abstract .solr exception."""
    pass

class _AbstractSolrInternalErrorException(_AbstractSolrException):
    """Internal code error (most likely a bug in the code)."""
    pass

class McUppercaseBooleanOperatorsInvalidTypeException(_AbstractSolrInternalErrorException):
    """Exception thrown when weird stuff gets passed to _uppercase_boolean_operators()."""
    pass

class McQuerySolrInternalErrorException(_AbstractSolrInternalErrorException):
    """Exception thrown when query_solr() receives something that it didn't expect."""
    pass

class _AbstractSolrInvalidQueryException(_AbstractSolrException):
    """Invalid Solr query."""
    pass

class McInsertCollectionMediaIDsInvalidQueryException(_AbstractSolrInvalidQueryException):
    """Invalid Solr query encountered in _insert_collection_media_ids()."""
    pass

class McQuerySolrInvalidQueryException(_AbstractSolrInvalidQueryException):
    """Invalid Solr query in query_solr()."""
    pass

class McQuerySolrRangeQueryException(_AbstractSolrInvalidQueryException):
    """Exception thrown on range queries which are not allowed."""
    pass

def _uppercase_boolean_operators(query: Union[List[str], str]) -> Union[List[str], str]:
    """
    Convert any "and", "or", or "not" operations in the argument to uppercase.

    If the argument is a list, call ourselves on all elements of the list.
    """
    query = decode_object_from_bytes_if_needed(query)

    if not query:
        return query

    def upper_repl(match) -> str:
        return match.group(1).upper()

    if isinstance(query, list):
        query = [_uppercase_boolean_operators(_) for _ in query]
    elif isinstance(query, str):
        query = re.sub(r'\b(and|or|not)\b', upper_repl, query)
    else:
        raise McUppercaseBooleanOperatorsInvalidTypeException(f"Invalid query type: {query}")

    return query

def _insert_collection_media_ids(db: DatabaseHandler, q: str) -> str:
    """
    Transform any "tags_id_media:" or "collections_id:" clauses into "media_id:" clauses with the "media_ids" that
    corresponds to the given tags.
    """

    q = decode_object_from_bytes_if_needed(q)

    def get_media_ids_clause(match) -> str:
        """Given the argument of "tags_id_media:" or "collections_id:" clause, return the corresponding "media_ids"."""
        arg = match.group(2)

        tags_ids = []
        if re.search(r'^\d+', arg):
            tags_ids.append(arg)

        else:
            parens_match = re.search(r'^\((.*)\)$', arg)

            if parens_match:
                parens = parens_match.group(1)

                parens = re.sub(r'or', ' ', parens, flags=re.IGNORECASE)
                parens = parens.strip()

                if re.search(r'[^\d\s]', parens):
                    raise McInsertCollectionMediaIDsInvalidQueryException((
                        f'Only "or" clauses allowed inside "tags_id_media:" or "collections_id:" clauses: {parens}; '
                        f'full match: {arg}'
                    ))

                for tags_id in re.split(r'\s+', parens):
                    tags_ids.append(tags_id)

            elif re.search(r'^\[', arg):
                raise McQuerySolrRangeQueryException(
                    'Range queries not allowed for "tags_id_media:" or "collections_id:" clauses'
                )

            else:
                raise McInsertCollectionMediaIDsInvalidQueryException(
                    f'Unrecognized format of "tags_id_media:" or "collections_id:" clause: {arg}'
                )

        media_ids = db.query("""
            SELECT media_id
            FROM media_tags_map
            WHERE tags_id IN %(tags_ids)s
            ORDER BY media_id
        """, {'tags_ids': tuple(tags_ids)}).flat()

        # Replace empty list with an id that will always return nothing from Solr
        if not media_ids:
            media_ids = [-1]

        media_clause = f"media_id:({' '.join([str(_) for _ in media_ids])})"

        return media_clause

    if not q:
        return q

    q = re.sub(r'(tags_id_media|collections_id):(\d+|\([^)]*\)|\[[^\]]*\])', get_media_ids_clause, q)

    return q

def query_solr(db: DatabaseHandler, params: SolrParams) -> Dict[str, Any]:
    """
    Execute a query on the Solr server using the given parameters. Return a maximum of 1 million sentences.

    The "params" argument is a dictionary of query parameters to Solr, detailed here:

        https://lucene.apache.org/solr/guide/6_6/common-query-parameters.html.

    The query ("params['q']") is transformed: lower case boolean operators are made uppercase to make Solr recognize
    them as boolean queries.

    Return decoded response in the format described here:

        https://lucene.apache.org/solr/guide/6_6/response-writers.html#ResponseWriters-JSONResponseWriter
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    if not params:
        raise McQuerySolrInternalErrorException('Parameters must be set.')

    if not isinstance(params, dict):
        raise McQuerySolrInternalErrorException('Parameters must be a dictionary.')

    params['wt'] = 'json'

    if 'rows' in params:
        params['rows'] = int(params['rows'])
    else:
        params['rows'] = 1000

    if 'df' not in params:
        params['df'] = 'text'

    params['rows'] = min(params['rows'], 10_000_000)

    if 'q' not in params:
        params['q'] = ''

    if 'fq' not in params:
        params['fq'] = []

    if not isinstance(params['fq'], list):
        params['fq'] = [params['fq']]

    if ':[' in params['q']:
        raise McQuerySolrRangeQueryException(
            "Range queries are not allowed in the main query. Please use a filter query instead for range queries."
        )

    # if params['q']:
    #     params['q'] = f"{{!complexphrase inOrder=false}} {params['q']}"

    params['q'] = _uppercase_boolean_operators(params['q'])
    params['fq'] = _uppercase_boolean_operators(params['fq'])

    if params['q']:
        params['q'] = _insert_collection_media_ids(db=db, q=params['q'])
    if params['fq']:
        params['fq'] = [_insert_collection_media_ids(db=db, q=_) for _ in params['fq']]

    response_json = solr_request(
        path='select',
        params={},
        content=params,
        content_type='application/x-www-form-urlencoded; charset=utf-8',
    )

    try:
        response = decode_json(response_json)
    except Exception as ex:
        raise McQuerySolrInternalErrorException(f"Error parsing Solr JSON: {ex}\nJSON: {response_json}")

    if 'error' in response:
        raise McQuerySolrInvalidQueryException(f"Error received from Solr: {response_json}")

    return response

def _get_intersection_of_lists(lists: List[List[int]]) -> List[int]:
    """Given a list of lists, each of which points to a list of IDs, return an intersection between them."""
    lists = decode_object_from_bytes_if_needed(lists)

    if not lists:
        log.error("Lists are empty")
        return []

    intersection = set(lists[0])

    for cur_list in lists[1:]:
        intersection = intersection.intersection(set(cur_list))

    return sorted(list(intersection))

def _get_stories_ids_from_stories_only_q(q: str) -> Optional[List[int]]:
    """
    Transform the pseudoquery fields in the query and then run a simple pattern to detect queries that consists of one
    or more AND'ed "stories_id:..." clauses.

    For those cases, just return the story IDs list rather than running it through Solr.

    Return None if the query does not match.
    """
    q = decode_object_from_bytes_if_needed(q)

    if not q:
        return None

    q = re.sub(r'^\s*\(\s*(?P<inside_parens>.*)\s*\)\s*$', r'\g<inside_parens>', q)
    q = q.strip()

    if 'and' in q.lower():
        p = q.lower().index('and')
    else:
        p = -1

    if p > 0:
        q_a = q[:p - 1]
        a_stories_ids = _get_stories_ids_from_stories_only_q(q=q_a)
        if a_stories_ids is None:
            return None

        q_b = q[p + 4:]
        b_stories_ids = _get_stories_ids_from_stories_only_q(q=q_b)
        if b_stories_ids is None:
            return None

        r = _get_intersection_of_lists(lists=[a_stories_ids, b_stories_ids])
        return r

    story_match = re.search(r'^stories_id:(\d+)$', q)
    if story_match:
        stories_id = int(story_match.group(1))
        r = [stories_id]
        return r

    if re.search(r'^stories_id:\([\s\d]+\)$', q):
        stories_ids = []
        for story_match in re.findall(r'(\d+)', q):
            stories_id = int(story_match)
            stories_ids.append(stories_id)
        return stories_ids

    return None

def _get_stories_ids_from_stories_only_params(params: SolrParams) -> Optional[List[int]]:
    """
    Transform the pseudoquery fields in the "q" and "fq" params and then run a simple pattern to detect queries that
    consists of one or more AND'ed "stories_id:..." clauses in the "q" param and all "fq" params.

    Return None if either the "q" or any of the "fq" params do not match.
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    q = params.get('q', '')
    fqs = params.get('fq', [])
    start = params.get('start', None)
    rows = params.get('rows', None)

    if start is not None:
        start = int(start)
    if rows is not None:
        rows = int(rows)

    # Return None if there are any unrecognized params
    param_keys = set(list(params.keys()))
    if not param_keys.issubset({'q', 'fq', 'start', 'rows'}):
        log.error(f"Parameters have unrecognized keys; all keys: {param_keys}")
        return None

    if not q:
        log.error("'q' is unset.")
        return None

    stories_ids_lists = []

    if fqs:
        if not isinstance(fqs, list):
            fqs = [fqs]

        for fq in fqs:
            stories_ids = _get_stories_ids_from_stories_only_q(q=fq)
            if stories_ids is None:
                return None
            else:
                stories_ids_lists.append(stories_ids)

    # If there are stories_ids only "fqs" and a '*:*' "q", just use the "fqs"
    if stories_ids_lists and q == '*:*':
        r = _get_intersection_of_lists(lists=stories_ids_lists)

    # If there were no "fqs" and a '*:*' "q", return None
    elif q == '*:*':
        return None

    # Otherwise, combine "q" and "fqs"
    else:
        stories_ids = _get_stories_ids_from_stories_only_q(q=q)
        if stories_ids is None:
            return None

        r = _get_intersection_of_lists(lists=[stories_ids] + stories_ids_lists)

    if start is not None:
        r = r[start:]

    if rows is not None:
        r = r[:rows]

    return r

def get_solr_num_found(db: DatabaseHandler, params: SolrParams) -> int:
    """Execute the query and return only the number of documents found."""
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    params['rows'] = 0

    res = query_solr(db=db, params=params)

    num_found = res['response']['numFound']

    return num_found

def search_solr_for_stories_ids(db: DatabaseHandler, params: SolrParams) -> List[int]:
    """
    Return a list of all of the "stories_ids" that match the Solr query.

    Using Solr side grouping on the "stories_id" field.
    """
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    stories_ids = _get_stories_ids_from_stories_only_params(params)
    if stories_ids:
        return stories_ids

    params['fl'] = 'stories_id'

    response = query_solr(db=db, params=params)

    stories_ids = [_['stories_id'] for _ in response['response']['docs']]

    return stories_ids

def search_solr_for_processed_stories_ids(db: DatabaseHandler,
                                          q: str,
                                          fq: Optional[Union[str, List[str]]],
                                          last_ps_id: int,
                                          num_stories: int,
                                          sort_by_random: bool = False) -> List[int]:
    """
    Return the first "num_stories" "processed_stories_id" that match the given query, sorted by "processed_stories_id"
    and with "processed_stories_id" greater than "last_ps_id".

    Returns at most "num_stories" stories.

    If "sort_by_random" is True, tell Solr to sort results by random order.
    """

    q = decode_object_from_bytes_if_needed(q)
    fq = decode_object_from_bytes_if_needed(fq)
    if isinstance(last_ps_id, bytes):
        last_ps_id = decode_object_from_bytes_if_needed(last_ps_id)
    if isinstance(num_stories, bytes):
        num_stories = decode_object_from_bytes_if_needed(num_stories)
    if isinstance(sort_by_random, bytes):
        sort_by_random = decode_object_from_bytes_if_needed(sort_by_random)

    last_ps_id = int(last_ps_id)
    num_stories = int(num_stories)
    sort_by_random = bool(int(sort_by_random))

    if not num_stories:
        return []

    if fq:
        if not isinstance(fq, list):
            fq = [fq]
    else:
        fq = []

    if last_ps_id:
        min_ps_id = last_ps_id + 1
        fq.append(f"processed_stories_id:[{min_ps_id} TO *]")

    params = {
        'q': q,
        'fq': fq,
        'fl': 'processed_stories_id',
        'rows': num_stories,
        'sort': 'random_1 asc' if sort_by_random else 'processed_stories_id asc',
    }

    response = query_solr(db=db, params=params)

    ps_ids = [_['processed_stories_id'] for _ in response['response']['docs']]

    return ps_ids

def search_solr_for_media_ids(db: DatabaseHandler, params: SolrParams) -> List[int]:
    """Return all of the media IDs that match the Solr query."""
    params = decode_object_from_bytes_if_needed(params)

    # Avoid editing the dictionary itself
    params = copy.deepcopy(params)

    params['fl'] = 'media_id'
    params['facet'] = 'true'
    params['facet.limit'] = 1_000_000
    params['facet.field'] = 'media_id'
    params['facet.mincount'] = 1
    params['rows'] = 0

    response = query_solr(db=db, params=params)

    counts = response['facet_counts']['facet_fields']['media_id']

    # Every second element (?)
    media_ids = counts[::2]

    return media_ids
<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Interpreter version: python 2.7
#
"""
Definition of structures, which are used to hold informations about
catalogization process.
"""
# Imports =====================================================================
from collections import namedtuple

# Structures ==================================================================
class SemanticInfo(namedtuple("SemanticInfo", ["hasAcquisitionFields",
                                               "hasISBNAgencyFields",
                                               "hasDescriptiveCatFields",
                                               "hasDescriptiveCatReviewFields",
                                               "hasSubjectCatFields",
                                               "hasSubjectCatReviewFields",
                                               "isClosed"])):
    """
    This structure is used to represent informations about export progress in
    Aleph.

    It contains informations about state of the record, so it can be tracked
    from edeposit project.

    See :func:`.toSemanticInfo` for details of parsing of those attributes.

    Attributes:
        hasAcquisitionFields (bool): Was the record aproved by acquisition?
        hasISBNAgencyFields (bool):  Was the record approved by ISBN agency?
        hasDescriptiveCatFields (bool): Did the record get thru name
                                        description (jmenný popis).
        hasDescriptiveCatReviewFields (bool): Did the record get thru name
                                              revision (jmenná revize).
        hasSubjectCatFields (bool): Did the record get thru subject
                                    description (věcný popis).
        hasSubjectCatReviewFields (bool): Did the record get thru subject
                                          revision (věcná revize).
        isClosed (bool): Was the record closed? This sometimes happen when bad
                         ISBN is given by creator of the record, but different
                         is in the book.
    """
    pass
<commit_msg>#38: Added .summaryRecordSysNumber and .parsedSummaryRecordSysNumber.<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Interpreter version: python 2.7
#
"""
Definition of structures, which are used to hold informations about
catalogization process.
"""
# Imports =====================================================================
from collections import namedtuple

# Structures ==================================================================
class SemanticInfo(namedtuple("SemanticInfo", ["hasAcquisitionFields",
                                               "hasISBNAgencyFields",
                                               "hasDescriptiveCatFields",
                                               "hasDescriptiveCatReviewFields",
                                               "hasSubjectCatFields",
                                               "hasSubjectCatReviewFields",
                                               "isClosed",
                                               "parsedSummaryRecordSysNumber",
                                               "summaryRecordSysNumber"])):
    """
    This structure is used to represent informations about export progress in
    Aleph.

    It contains informations about state of the record, so it can be tracked
    from edeposit project.

    See :func:`.toSemanticInfo` for details of parsing of those attributes.

    Attributes:
        hasAcquisitionFields (bool): Was the record aproved by acquisition?
        hasISBNAgencyFields (bool):  Was the record approved by ISBN agency?
        hasDescriptiveCatFields (bool): Did the record get thru name
                                        description (jmenný popis).
        hasDescriptiveCatReviewFields (bool): Did the record get thru name
                                              revision (jmenná revize).
        hasSubjectCatFields (bool): Did the record get thru subject
                                    description (věcný popis).
        hasSubjectCatReviewFields (bool): Did the record get thru subject
                                          revision (věcná revize).
        isClosed (bool): Was the record closed? This sometimes happen when bad
                         ISBN is given by creator of the record, but different
                         is in the book.
        parsedSummaryRecordSysNumber (str): Same as
              :attr:`summaryRecordSysNumber` but without natural language
              details.
        summaryRecordSysNumber (str): Identificator of the new record if
                               `.isClosed` is True. Format of the string is
                               not specified and can be different for each
                               record.
    """
    pass
<|endoftext|>