<commit_before>"""Parser responsible for reading the ``*dep.m`` files"""
import re
from warnings import warn

from matplotlib import pyplot
from six import iteritems

from serpentTools.utils import (
    magicPlotDocDecorator, formatPlot, DEPLETION_PLOT_LABELS,
    convertVariableName, str2vec,
)
from ._collections import DEPLETION_VARIABLES
from serpentTools.engines import KeywordParser
from serpentTools.parsers.base import MaterialReader
from serpentTools.objects.materials import DepletedMaterial
from serpentTools.utils.core import deconvertVariableName
from serpentTools.messages import (
    warning, debug, error, SerpentToolsException,
)
from serpentTools.utils import (
    getKeyMatchingShapes,
    logDirectCompare,
    compareDocDecorator,
    DEF_COMP_LOWER,
    DEF_COMP_UPPER,
    DEF_COMP_SIGMA,
)
from serpentTools.io import MatlabConverter

METADATA_KEYS = {'ZAI', 'NAMES', 'BU', 'DAYS'}

class DepPlotMixin(object):

    @magicPlotDocDecorator
    def plot(self, xUnits, yUnits=None, timePoints=None, names=None, zai=None,
             materials=None, ax=None, legend=None, logx=False, logy=False,
             loglog=False, labelFmt=None, xlabel=None, ylabel=None, ncol=1,
             **kwargs):
        """
        Plot properties for all materials in this file together.

        Parameters
        ----------
        xUnits: str
            If ``xUnits`` is given and ``yUnits`` is ``None``, then
            the plotted data will be ``xUnits`` against ``'days'``
            name of x value to obtain, e.g. ``'days'``, ``'burnup'``
        yUnits: str
            name of y value to return, e.g. ``'adens'``, ``'ingTox'``
        timePoints: list or None
            If given, select the time points according to those
            specified here. Otherwise, select all points

            .. deprecated:: 0.7.0
               Will plot against all time points

        names: str or list or None
            If given, plot  values corresponding to these isotope
            names. Otherwise, plot values for all isotopes.
        zai: int or list or None
            If given, plot values corresponding to these
            isotope ``ZZAAAI`` values. Otherwise, plot for all isotopes

            .. versionadded:: 0.5.1

        materials: None or list
            Selection of materials from ``self.materials`` to plot.
            If None, plot all materials, potentially including ``tot``
        {ax}
        {legend}
        {xlabel} Otherwise, use ``xUnits``
        {ylabel} Otherwise, use ``yUnits``
        {logx}
        {logy}
        {loglog}
        {matLabelFmt}
        {ncol}
        {kwargs} :py:func:`matplotlib.pyplot.plot`

        Returns
        -------
        {rax}

        See Also
        --------
        * :py:func:`~serpentTools.objects.materials.DepletedMaterial.getValues`
        * :py:func:`matplotlib.pyplot.plot`
        * :py:meth:`str.format` - used for formatting labels
        * :py:func:`~serpentTools.objects.materials.DepletedMaterial.plot`

        Raises
        ------
        KeyError
            If x axis units are not ``'days'`` nor ``'burnup'``
        SerpentToolsException
            If the materials dictionary does not contain any items
        """
        if yUnits is None:
            yUnits = xUnits
            xUnits = 'days'

        if not self.materials:
            raise SerpentToolsException("Material dictionary is empty")

        if xUnits not in ('days', 'burnup'):
            raise KeyError("Plot method only uses x-axis data from <days> and "
                           "<burnup>, not {}".format(xUnits))
        missing = set()
        ax = ax or pyplot.gca()
        materials = materials or self.materials.keys()
        labelFmt = labelFmt or '{mat} {iso}'
        for mat in materials:
            if mat not in self.materials:
                missing.add(mat)
                continue

            ax = self.materials[mat].plot(
                xUnits, yUnits, timePoints, names,
                zai, ax, legend=False, xlabel=xlabel, ylabel=ylabel,
                logx=False, logy=False, loglog=False, labelFmt=labelFmt,
                **kwargs)
        if missing:
            warning("The following materials were not found in materials "
                    "dictionary: {}".format(', '.join(missing)))
        formatPlot(ax, legend=legend, ncol=ncol, logx=logx, logy=logy,
                   loglog=loglog,
                   xlabel=xlabel or DEPLETION_PLOT_LABELS[xUnits],
                   ylabel=ylabel or DEPLETION_PLOT_LABELS[yUnits],
                   )

        return ax

class DepletionReader(DepPlotMixin, MaterialReader):
    """
    Parser responsible for reading and working with depletion files.

    Parameters
    ----------
    filePath: str
        path to the depletion file

    Attributes
    ----------
    {attrs:s}
    settings: dict
        names and values of the settings used to control operations
        of this reader
    """
    docAttrs = """materials: dict
        Dictionary with material names as keys and the corresponding
        :py:class:`~serpentTools.objects.materials.DepletedMaterial` class
        for that material as values
    metadata: dict
        Dictionary with file-wide data names as keys and the
        corresponding data, e.g. ``'zai'``: [list of zai numbers]"""
    __doc__ = __doc__.format(attrs=docAttrs)

    def __init__(self, filePath):
        MaterialReader.__init__(self, filePath, 'depletion')
        patterns = self.settings['materials'] or ['.*']
        # match all materials if nothing given
        self._matPatterns = [re.compile(mat) for mat in patterns]
        DepPlotMixin.__init__(self)

    def __getitem__(self, name):
        """Retrieve a material from :attr:`materials`."""
        return self.materials[name]

    def _read(self):
        """Read through the depletion file and store requested data."""
        keys = ['MAT', 'TOT'] if self.settings['processTotal'] else ['MAT']
        keys.extend(self.settings['metadataKeys'])
        separators = ['\n', '];', '\r\n']
        with KeywordParser(self.filePath, keys, separators) as parser:
            for chunk in parser.yieldChunks():
                mvar = getMatlabVarName(chunk)
                if mvar[:3] in ['MAT', 'TOT']:
                    name, variable = getMaterialNameAndVariable(mvar)
                    self._checkAddData(chunk, name, variable)
                    continue
                self._addMetadata(chunk)
        if 'days' in self.metadata:
            for mKey in self.materials:
                self.materials[mKey].days = self.metadata['days']

    def _addMetadata(self, chunk):
        for varName in METADATA_KEYS:
            if varName not in chunk[0]:
                continue
            if varName in ['ZAI', 'NAMES']:
                cleaned = [line.strip() for line in chunk[1:]]
                if varName == 'NAMES':
                    values = [item[1:item.find(" ")] for item in cleaned]
                else:
                    values = str2vec(cleaned, int, list)
            else:
                line = self._cleanSingleLine(chunk)
                values = str2vec(line)
            self.metadata[convertVariableName(varName)] = values
            return
        warning("Unsure about how to process metadata chunk {}"
                .format(chunk[0]))

    @staticmethod
    def _cleanSingleLine(chunk):
        return chunk[0][chunk[0].index('[') + 1:chunk[0].index(']')]

    def _checkAddData(self, chunk, name, variable):
        if name == 'total':
            if not self.settings['processTotal']:
                return
            return self._processChunk(chunk, name, variable)
        for pattern in self._matPatterns:
            if pattern.match(name):
                return self._processChunk(chunk, name, variable)

    def _processChunk(self, chunk, name, variable):
        if (self.settings['materialVariables']
                and variable not in self.settings['materialVariables']):
            return
        if name not in self.materials:
            debug('Adding material {}...'.format(name))
            self.materials[name] = DepletedMaterial(name, self.metadata)
        if len(chunk) == 1:  # single line values, e.g. volume or burnup
            cleaned = self._cleanSingleLine(chunk)
        else:
            cleaned = []
            for line in chunk:
                if '[' in line or ']' in line:
                    continue
                cleaned.append(line[:line.index('%')])
        self.materials[name].addData(variable, cleaned)

    def _precheck(self):
        """do a quick scan to ensure this looks like a material file."""
        with open(self.filePath) as fh:
            for line in fh:
                sline = line.split()
                if not sline:
                    continue
                elif 'MAT' == line.split('_')[0]:
                    return
        error("No materials found in {}".format(self.filePath))

    def _postcheck(self):
        """ensure the parser grabbed expected materials."""
        if not self.materials:
            error("No materials obtained from {}".format(self.filePath))
            return

        for mKey, mat in iteritems(self.materials):
            assert isinstance(mat, DepletedMaterial), (
                'Unexpected key {}: {} in materials dictionary'.format(
                    mKey, type(mat))
            )
        debug('  found {} materials'.format(len(self.materials)))

        if 'bu' in self.metadata:
            self.metadata['burnup'] = self.metadata.pop('bu')

    def _compare(self, other, lower, upper, _sigma):

        similar = self._compareMetadata(other, lower, upper, _sigma)
        if not self._comparePrecheckMetadata(other):
            return False
        similar &= self._compareMaterials(other, lower, upper, _sigma)
        return similar

    def _comparePrecheckMetadata(self, other):
        for key, myVec in iteritems(self.metadata):
            otherVec = other.metadata[key]
            if len(myVec) != len(otherVec):
                error("Stopping comparison early due to mismatched {} vectors"
                      "\n\t>{}\n\t<{}".format(key, myVec, otherVec))
                return False
        return True

    @compareDocDecorator
    def compareMaterials(self, other, lower=DEF_COMP_LOWER,
                         upper=DEF_COMP_UPPER, sigma=DEF_COMP_SIGMA):
        """
        Return the result of comparing all materials on two readers

        Parameters
        ----------
        other: :class:`DepletionReader`
            Reader to compare against
        {compLimits}
        {sigma}

        Returns
        -------
        bool:
            ``True`` if all materials agree to the given tolerances

        Raises
        ------
        {compTypeErr}
        """
        self._checkCompareObj(other)
        self._compareLogPreMsg(other, lower, upper, quantity='materials')

        if not self._comparePrecheckMetadata(other):
            return False

        return self._compareMaterials(other, lower, upper, sigma)

    def _compareMaterials(self, other, lower, upper, sigma):
        """Private method for going directly into the comparison."""
        commonMats = getKeyMatchingShapes(
            self.materials, other.materials, 'materials')
        similar = (
            len(self.materials) == len(other.materials) == len(commonMats))

        for matName in sorted(commonMats):
            myMat = self[matName]
            otherMat = other[matName]
            similar &= myMat.compare(otherMat, lower, upper, sigma)
        return similar

    @compareDocDecorator
    def compareMetadata(self, other, lower=DEF_COMP_LOWER,
                        upper=DEF_COMP_UPPER, sigma=DEF_COMP_SIGMA):
        """
        Return the result of comparing metadata on two readers

        Parameters
        ----------
        other: :class:`DepletionReader`
            Object to compare against
        {compLimits}
        {header}

        Returns
        -------
        bool
            True if the metadata agree within the given tolerances

        Raises
        ------
        {compTypeErr}
        """

        self._checkCompareObj(other)

        self._compareLogPreMsg(other, lower, upper, quantity='metadata')

        return self._compareMetadata(other, lower, upper, sigma)

    def _compareMetadata(self, other, lower, upper, _sigma):
        """Private method for comparing metadata"""

        similar = logDirectCompare(
            self.metadata['names'], other.metadata['names'],
            0, 0, 'names')
        similar &= logDirectCompare(
            self.metadata['zai'], other.metadata['zai'],
            0, 0, 'zai')
        similar &= logDirectCompare(
            self.metadata['days'], other.metadata['days'],
            lower, upper, 'days')
        similar &= logDirectCompare(
            self.metadata['burnup'], other.metadata['burnup'],
            lower, upper, 'burnup')
        return similar

    def saveAsMatlab(self, fileP, reconvert=True, metadata=None,
                     append=True, format='5', longNames=True,
                     compress=True, oned='row'):
        """
        Write a binary MATLAB file from the contents of this reader

        Parameters
        ----------
        fileP: str or file-like object
            Name of the file to write
        reconvert: bool
            If this evaluates to true, convert values back into their
            original form as they appear in the output file, e.g.
            ``MAT_TOTAL_ING_TOX``. Otherwise, maintain the ``mixedCase``
            style, ``total_ingTox``.
        metadata: bool or str or list of strings
            If this evaluates to true, then write all metadata to the
            file as well.
        append: bool
            If true and a file exists under ``output``, append to that file.
            Otherwise the file will be overwritten
        format: {'5', '4'}
            Format of file to write. ``'5'`` for MATLAB 5 to 7.2, ``'4'`` for
            MATLAB 4
        longNames: bool
            If true, allow variable names to reach 63 characters,
            which works with MATLAB 7.6+. Otherwise, maximum length is
            31 characters
        compress: bool
            If true, compress matrices on write
        oned: {'row', 'col'}:
            Write one-dimensional arrays as row vectors if
            ``oned=='row'`` (default), or column vectors

        Raises
        ------
        ImportError:
            If :term:`scipy` is not installed

        See Also
        --------
        :func:`scipy.io.savemat`
        """
        if metadata is not None:
            # need this path to perserve selecting not to write
            # metadata
            from scipy.io import savemat
            data = self._gather_matlab(reconvert, metadata)
            return savemat(fileP, data, append, format, longNames,
                           compress, oned)
        converter = MatlabConverter(self, fileP)
        return converter.convert(reconvert, append, format, longNames,
                                 compress, oned)

    def _gather_matlab(self, reconvert, metadata=None):
        """
        Return a dictionary with all data

        Should only be used by developers for converting to new
        output file types
        """

        if metadata is not None:
            warn("version 0.7.0+ will not use metadata argument and write "
                 "all by default.", FutureWarning)
        # set these here to reduce number of if/elses

        if reconvert:
            converter = deconvert
        else:
            converter = prepToMatlab

        # remove for v0.7.0
        if metadata and self.metadata:
            data = {k.upper() if reconvert else k: v
                    for k, v in iteritems(self.metadata)}
        else:
            data = {}

        for matName, material in iteritems(self.materials):
            for varName, varData in iteritems(material.data):
                data[converter(matName, varName)] = varData

        return data

def getMaterialNameAndVariable(mlabName):
    for serpVar in DEPLETION_VARIABLES:
        leng = len(serpVar)
        if serpVar == mlabName[-len(serpVar):]:
            matName = mlabName[4:-(leng + 1)]
            if not matName:
                if 'TOT' == mlabName[:3]:
                    return 'total', serpVar
                break
            return matName, serpVar
    raise ValueError(
        "Could not find a match for {}. "
        "Please alert developers".format(mlabName))

def getMatlabVarName(chunk):
    if isinstance(chunk, list):
        return getMatlabVarName(chunk[0])
    return chunk[:chunk.index(' ')]

_MAT_FMT_ORIG = "MAT_{}_{}"
_MAT_FMT_CC = "{}_{}"

def deconvert(material, variable):
    """Restore the original name as present in the depletion file"""
    return _MAT_FMT_ORIG.format(material, deconvertVariableName(variable))

def prepToMatlab(material, variable):
    """Create the name of a single variable for MATLAB"""
    return _MAT_FMT_CC.format(material, variable)
<commit_msg>API Deprecate DepletionReader.saveAsMatlab<commit_after>"""Parser responsible for reading the ``*dep.m`` files"""
import re
from warnings import warn

from matplotlib import pyplot
from six import iteritems

from serpentTools.utils import (
    magicPlotDocDecorator, formatPlot, DEPLETION_PLOT_LABELS,
    convertVariableName, str2vec,
)
from ._collections import DEPLETION_VARIABLES
from serpentTools.engines import KeywordParser
from serpentTools.parsers.base import MaterialReader
from serpentTools.objects.materials import DepletedMaterial
from serpentTools.utils.core import deconvertVariableName
from serpentTools.messages import (
    warning, debug, error, SerpentToolsException, deprecated,
)
from serpentTools.utils import (
    getKeyMatchingShapes,
    logDirectCompare,
    compareDocDecorator,
    DEF_COMP_LOWER,
    DEF_COMP_UPPER,
    DEF_COMP_SIGMA,
)
from serpentTools.io import MatlabConverter

METADATA_KEYS = {'ZAI', 'NAMES', 'BU', 'DAYS'}

class DepPlotMixin(object):

    @magicPlotDocDecorator
    def plot(self, xUnits, yUnits=None, timePoints=None, names=None, zai=None,
             materials=None, ax=None, legend=None, logx=False, logy=False,
             loglog=False, labelFmt=None, xlabel=None, ylabel=None, ncol=1,
             **kwargs):
        """
        Plot properties for all materials in this file together.

        Parameters
        ----------
        xUnits: str
            If ``xUnits`` is given and ``yUnits`` is ``None``, then
            the plotted data will be ``xUnits`` against ``'days'``
            name of x value to obtain, e.g. ``'days'``, ``'burnup'``
        yUnits: str
            name of y value to return, e.g. ``'adens'``, ``'ingTox'``
        timePoints: list or None
            If given, select the time points according to those
            specified here. Otherwise, select all points

            .. deprecated:: 0.7.0
               Will plot against all time points

        names: str or list or None
            If given, plot  values corresponding to these isotope
            names. Otherwise, plot values for all isotopes.
        zai: int or list or None
            If given, plot values corresponding to these
            isotope ``ZZAAAI`` values. Otherwise, plot for all isotopes

            .. versionadded:: 0.5.1

        materials: None or list
            Selection of materials from ``self.materials`` to plot.
            If None, plot all materials, potentially including ``tot``
        {ax}
        {legend}
        {xlabel} Otherwise, use ``xUnits``
        {ylabel} Otherwise, use ``yUnits``
        {logx}
        {logy}
        {loglog}
        {matLabelFmt}
        {ncol}
        {kwargs} :py:func:`matplotlib.pyplot.plot`

        Returns
        -------
        {rax}

        See Also
        --------
        * :py:func:`~serpentTools.objects.materials.DepletedMaterial.getValues`
        * :py:func:`matplotlib.pyplot.plot`
        * :py:meth:`str.format` - used for formatting labels
        * :py:func:`~serpentTools.objects.materials.DepletedMaterial.plot`

        Raises
        ------
        KeyError
            If x axis units are not ``'days'`` nor ``'burnup'``
        SerpentToolsException
            If the materials dictionary does not contain any items
        """
        if yUnits is None:
            yUnits = xUnits
            xUnits = 'days'

        if not self.materials:
            raise SerpentToolsException("Material dictionary is empty")

        if xUnits not in ('days', 'burnup'):
            raise KeyError("Plot method only uses x-axis data from <days> and "
                           "<burnup>, not {}".format(xUnits))
        missing = set()
        ax = ax or pyplot.gca()
        materials = materials or self.materials.keys()
        labelFmt = labelFmt or '{mat} {iso}'
        for mat in materials:
            if mat not in self.materials:
                missing.add(mat)
                continue

            ax = self.materials[mat].plot(
                xUnits, yUnits, timePoints, names,
                zai, ax, legend=False, xlabel=xlabel, ylabel=ylabel,
                logx=False, logy=False, loglog=False, labelFmt=labelFmt,
                **kwargs)
        if missing:
            warning("The following materials were not found in materials "
                    "dictionary: {}".format(', '.join(missing)))
        formatPlot(ax, legend=legend, ncol=ncol, logx=logx, logy=logy,
                   loglog=loglog,
                   xlabel=xlabel or DEPLETION_PLOT_LABELS[xUnits],
                   ylabel=ylabel or DEPLETION_PLOT_LABELS[yUnits],
                   )

        return ax

class DepletionReader(DepPlotMixin, MaterialReader):
    """
    Parser responsible for reading and working with depletion files.

    Parameters
    ----------
    filePath: str
        path to the depletion file

    Attributes
    ----------
    {attrs:s}
    settings: dict
        names and values of the settings used to control operations
        of this reader
    """
    docAttrs = """materials: dict
        Dictionary with material names as keys and the corresponding
        :py:class:`~serpentTools.objects.materials.DepletedMaterial` class
        for that material as values
    metadata: dict
        Dictionary with file-wide data names as keys and the
        corresponding data, e.g. ``'zai'``: [list of zai numbers]"""
    __doc__ = __doc__.format(attrs=docAttrs)

    def __init__(self, filePath):
        MaterialReader.__init__(self, filePath, 'depletion')
        patterns = self.settings['materials'] or ['.*']
        # match all materials if nothing given
        self._matPatterns = [re.compile(mat) for mat in patterns]
        DepPlotMixin.__init__(self)

    def __getitem__(self, name):
        """Retrieve a material from :attr:`materials`."""
        return self.materials[name]

    def _read(self):
        """Read through the depletion file and store requested data."""
        keys = ['MAT', 'TOT'] if self.settings['processTotal'] else ['MAT']
        keys.extend(self.settings['metadataKeys'])
        separators = ['\n', '];', '\r\n']
        with KeywordParser(self.filePath, keys, separators) as parser:
            for chunk in parser.yieldChunks():
                mvar = getMatlabVarName(chunk)
                if mvar[:3] in ['MAT', 'TOT']:
                    name, variable = getMaterialNameAndVariable(mvar)
                    self._checkAddData(chunk, name, variable)
                    continue
                self._addMetadata(chunk)
        if 'days' in self.metadata:
            for mKey in self.materials:
                self.materials[mKey].days = self.metadata['days']

    def _addMetadata(self, chunk):
        for varName in METADATA_KEYS:
            if varName not in chunk[0]:
                continue
            if varName in ['ZAI', 'NAMES']:
                cleaned = [line.strip() for line in chunk[1:]]
                if varName == 'NAMES':
                    values = [item[1:item.find(" ")] for item in cleaned]
                else:
                    values = str2vec(cleaned, int, list)
            else:
                line = self._cleanSingleLine(chunk)
                values = str2vec(line)
            self.metadata[convertVariableName(varName)] = values
            return
        warning("Unsure about how to process metadata chunk {}"
                .format(chunk[0]))

    @staticmethod
    def _cleanSingleLine(chunk):
        return chunk[0][chunk[0].index('[') + 1:chunk[0].index(']')]

    def _checkAddData(self, chunk, name, variable):
        if name == 'total':
            if not self.settings['processTotal']:
                return
            return self._processChunk(chunk, name, variable)
        for pattern in self._matPatterns:
            if pattern.match(name):
                return self._processChunk(chunk, name, variable)

    def _processChunk(self, chunk, name, variable):
        if (self.settings['materialVariables']
                and variable not in self.settings['materialVariables']):
            return
        if name not in self.materials:
            debug('Adding material {}...'.format(name))
            self.materials[name] = DepletedMaterial(name, self.metadata)
        if len(chunk) == 1:  # single line values, e.g. volume or burnup
            cleaned = self._cleanSingleLine(chunk)
        else:
            cleaned = []
            for line in chunk:
                if '[' in line or ']' in line:
                    continue
                cleaned.append(line[:line.index('%')])
        self.materials[name].addData(variable, cleaned)

    def _precheck(self):
        """do a quick scan to ensure this looks like a material file."""
        with open(self.filePath) as fh:
            for line in fh:
                sline = line.split()
                if not sline:
                    continue
                elif 'MAT' == line.split('_')[0]:
                    return
        error("No materials found in {}".format(self.filePath))

    def _postcheck(self):
        """ensure the parser grabbed expected materials."""
        if not self.materials:
            error("No materials obtained from {}".format(self.filePath))
            return

        for mKey, mat in iteritems(self.materials):
            assert isinstance(mat, DepletedMaterial), (
                'Unexpected key {}: {} in materials dictionary'.format(
                    mKey, type(mat))
            )
        debug('  found {} materials'.format(len(self.materials)))

        if 'bu' in self.metadata:
            self.metadata['burnup'] = self.metadata.pop('bu')

    def _compare(self, other, lower, upper, _sigma):

        similar = self._compareMetadata(other, lower, upper, _sigma)
        if not self._comparePrecheckMetadata(other):
            return False
        similar &= self._compareMaterials(other, lower, upper, _sigma)
        return similar

    def _comparePrecheckMetadata(self, other):
        for key, myVec in iteritems(self.metadata):
            otherVec = other.metadata[key]
            if len(myVec) != len(otherVec):
                error("Stopping comparison early due to mismatched {} vectors"
                      "\n\t>{}\n\t<{}".format(key, myVec, otherVec))
                return False
        return True

    @compareDocDecorator
    def compareMaterials(self, other, lower=DEF_COMP_LOWER,
                         upper=DEF_COMP_UPPER, sigma=DEF_COMP_SIGMA):
        """
        Return the result of comparing all materials on two readers

        Parameters
        ----------
        other: :class:`DepletionReader`
            Reader to compare against
        {compLimits}
        {sigma}

        Returns
        -------
        bool:
            ``True`` if all materials agree to the given tolerances

        Raises
        ------
        {compTypeErr}
        """
        self._checkCompareObj(other)
        self._compareLogPreMsg(other, lower, upper, quantity='materials')

        if not self._comparePrecheckMetadata(other):
            return False

        return self._compareMaterials(other, lower, upper, sigma)

    def _compareMaterials(self, other, lower, upper, sigma):
        """Private method for going directly into the comparison."""
        commonMats = getKeyMatchingShapes(
            self.materials, other.materials, 'materials')
        similar = (
            len(self.materials) == len(other.materials) == len(commonMats))

        for matName in sorted(commonMats):
            myMat = self[matName]
            otherMat = other[matName]
            similar &= myMat.compare(otherMat, lower, upper, sigma)
        return similar

    @compareDocDecorator
    def compareMetadata(self, other, lower=DEF_COMP_LOWER,
                        upper=DEF_COMP_UPPER, sigma=DEF_COMP_SIGMA):
        """
        Return the result of comparing metadata on two readers

        Parameters
        ----------
        other: :class:`DepletionReader`
            Object to compare against
        {compLimits}
        {header}

        Returns
        -------
        bool
            True if the metadata agree within the given tolerances

        Raises
        ------
        {compTypeErr}
        """

        self._checkCompareObj(other)

        self._compareLogPreMsg(other, lower, upper, quantity='metadata')

        return self._compareMetadata(other, lower, upper, sigma)

    def _compareMetadata(self, other, lower, upper, _sigma):
        """Private method for comparing metadata"""

        similar = logDirectCompare(
            self.metadata['names'], other.metadata['names'],
            0, 0, 'names')
        similar &= logDirectCompare(
            self.metadata['zai'], other.metadata['zai'],
            0, 0, 'zai')
        similar &= logDirectCompare(
            self.metadata['days'], other.metadata['days'],
            lower, upper, 'days')
        similar &= logDirectCompare(
            self.metadata['burnup'], other.metadata['burnup'],
            lower, upper, 'burnup')
        return similar

    @deprecated("toMatlab")
    def saveAsMatlab(self, fileP, reconvert=True, metadata=None,
                     append=True, format='5', longNames=True,
                     compress=True, oned='row'):
        """
        Write a binary MATLAB file from the contents of this reader

        .. deprecated:: 0.7.0
            Use :meth:`~serpentTools.DepletionReader.toMatlab`

        Parameters
        ----------
        fileP: str or file-like object
            Name of the file to write
        reconvert: bool
            If this evaluates to true, convert values back into their
            original form as they appear in the output file, e.g.
            ``MAT_TOTAL_ING_TOX``. Otherwise, maintain the ``mixedCase``
            style, ``total_ingTox``.
        metadata: bool or str or list of strings
            If this evaluates to true, then write all metadata to the
            file as well.
        append: bool
            If true and a file exists under ``output``, append to that file.
            Otherwise the file will be overwritten
        format: {'5', '4'}
            Format of file to write. ``'5'`` for MATLAB 5 to 7.2, ``'4'`` for
            MATLAB 4
        longNames: bool
            If true, allow variable names to reach 63 characters,
            which works with MATLAB 7.6+. Otherwise, maximum length is
            31 characters
        compress: bool
            If true, compress matrices on write
        oned: {'row', 'col'}:
            Write one-dimensional arrays as row vectors if
            ``oned=='row'`` (default), or column vectors

        Raises
        ------
        ImportError:
            If :term:`scipy` is not installed

        See Also
        --------
        :func:`scipy.io.savemat`
        """
        if metadata is not None:
            warn("metadata argument is deprecated. All metadata written",
                 FutureWarning)
            # need this path to perserve selecting not to write
            # metadata
        converter = MatlabConverter(self, fileP)
        return converter.convert(reconvert, append, format, longNames,
                                 compress, oned)

    def _gather_matlab(self, reconvert, metadata=None):
        """
        Return a dictionary with all data

        Should only be used by developers for converting to new
        output file types
        """

        # set these here to reduce number of if/elses

        if reconvert:
            converter = deconvert
        else:
            converter = prepToMatlab

        data = {k.upper() if reconvert else k: v
                for k, v in iteritems(self.metadata)}

        for matName, material in iteritems(self.materials):
            for varName, varData in iteritems(material.data):
                data[converter(matName, varName)] = varData

        return data

def getMaterialNameAndVariable(mlabName):
    for serpVar in DEPLETION_VARIABLES:
        leng = len(serpVar)
        if serpVar == mlabName[-len(serpVar):]:
            matName = mlabName[4:-(leng + 1)]
            if not matName:
                if 'TOT' == mlabName[:3]:
                    return 'total', serpVar
                break
            return matName, serpVar
    raise ValueError(
        "Could not find a match for {}. "
        "Please alert developers".format(mlabName))

def getMatlabVarName(chunk):
    if isinstance(chunk, list):
        return getMatlabVarName(chunk[0])
    return chunk[:chunk.index(' ')]

_MAT_FMT_ORIG = "MAT_{}_{}"
_MAT_FMT_CC = "{}_{}"

def deconvert(material, variable):
    """Restore the original name as present in the depletion file"""
    return _MAT_FMT_ORIG.format(material, deconvertVariableName(variable))

def prepToMatlab(material, variable):
    """Create the name of a single variable for MATLAB"""
    return _MAT_FMT_CC.format(material, variable)
<|endoftext|>


<commit_before>##
# Copyright (c) 2008-2010 Sprymix Inc.
# All rights reserved.
#
# See LICENSE for details.
##

import datetime
import dateutil.parser
import dateutil.relativedelta
import re

class DateTime(datetime.datetime):
    def __new__(cls, value=None):
        if isinstance(value, datetime.datetime):
            d = value
        elif isinstance(value, str):
            try:
                d = dateutil.parser.parse(value)
            except ValueError as e:
                raise ValueError("invalid value for DateTime object: %s" % value) from e
        else:
            raise ValueError("invalid value for DateTime object: %s" % value)

        return super().__new__(cls, d.year, d.month, d.day, d.hour, d.minute, d.second,
                                    d.microsecond, d.tzinfo)

class TimeDelta(dateutil.relativedelta.relativedelta):
    _interval_tokens = {'year': 'years', 'years': 'years',
                        'month': 'months', 'months': 'months',
                        'week': 'weeks', 'weeks': 'weeks',
                        'day': 'days', 'days': 'days',
                        'hour': 'hours', 'hours': 'hours',
                        'minute': 'minutes', 'minutes': 'minutes',
                        'second': 'seconds', 'seconds': 'seconds',
                        'microsecond': 'microseconds', 'microseconds': 'microseconds'}

    def __new__(cls, value=None, *, dt1=None, dt2=None, years=0, months=0, days=0, leapdays=0,
                      weeks=0, hours=0, minutes=0, seconds=0, microseconds=0, year=None,
                      month=None, day=None, weekday=None, yearday=None, nlyearday=None, hour=None,
                      minute=None, second=None, microsecond=None):

        if value is None:
            result = super().__new__(cls)
            super().__init__(result,
                             dt1=dt1, dt2=dt2, years=years, months=months, days=days,
                             hours=hours, minutes=minutes,
                             seconds=seconds, microseconds=microseconds,
                             leapdays=leapdays, year=year, weeks=weeks,
                             month=month, day=day, weekday=weekday,
                             hour=hour, minute=minute, second=second,
                             microsecond=microsecond)

        elif isinstance(value, dateutil.relativedelta.relativedelta):
            result = super().__new__(cls)
            super().__init__(result,
                             years=value.years, months=value.months, days=value.days,
                             hours=value.hours, minutes=value.minutes,
                             seconds=value.seconds, microseconds=value.microseconds,
                             leapdays=value.leapdays, year=value.year,
                             month=value.month, day=value.day, weekday=value.weekday,
                             hour=value.hour, minute=value.minute,
                             second=value.second, microsecond=value.microsecond)

        elif isinstance(value, datetime.timedelta):
            result = super().__new__(cls)
            super().__init__(result,
                             days=value.days, seconds=value.seconds,
                             microseconds=value.microseconds)

        elif isinstance(value, str):
            expecting_number = True

            intervals = {}
            interval_tokens = cls._interval_tokens

            try:
                number = None

                for token in re.split('\s+', value):
                    if expecting_number:
                        number = int(token)
                    else:
                        interval = interval_tokens[token]
                        if interval in intervals:
                            raise ValueError
                        intervals[interval] = number
                    expecting_number = not expecting_number
                if not expecting_number:
                    raise ValueError
            except (ValueError, KeyError):
                raise ValueError("invalid timedelta value: '%s'" % value)

            result = super().__new__(cls)
            super().__init__(result, **intervals)

        return result

    def __init__(self, *args, **kwargs):
        # We initialize in __new__, so nothing to do here
        pass

    def __str__(self):
        s = []
        for k in ('years', 'months', 'days', 'hours', 'minutes', 'seconds', 'microseconds'):
            v = getattr(self, k)
            if v:
                s.append('%s %s' % (v, k if v != 1 else k[:-1]))
        return ' '.join(s)

    def __hash__(self):
        return hash((self.weekday, self.year, self.months, self.days, self.hours,
                     self.minutes, self.seconds, self.leapdays, self.year, self.month,
                     self.day, self.hour, self.minute, self.second, self.microsecond))

    def __gt__(self, other):
        if not isinstance(other, TimeDelta):
            return NotImplemented

        if self.weekday or other.weekday:
            raise TypeError('TimeDeltas with non-empty weekday are not orderable')

        return self.years > other.years or self.years == other.years and \
               (self.months > other.months or self.months == other.months and \
               (self.days > other.days or self.days == other.days and \
               (self.hours > other.hours or self.hours == other.hours and \
               (self.minutes > other.minutes or self.minutes == other.minutes and \
               (self.seconds > other.seconds or self.seconds == other.seconds and \
               (self.microseconds > other.microseconds))))))

    def __ge__(self, other):
        return self == other or self > other

    def __lt__(self, other):
        return not self >= other

    def __le__(self, other):
        return not self > other
<commit_msg>utils.datetime: Always produce timezone-aware DateTime objects<commit_after>##
# Copyright (c) 2008-2010 Sprymix Inc.
# All rights reserved.
#
# See LICENSE for details.
##

import datetime
import dateutil.parser
import dateutil.relativedelta
import dateutil.tz
import re

from semantix.utils import config

@config.configurable
class DateTime(datetime.datetime):
    local_timezone = config.cvalue(type=str, default=None, doc='Default local time-zone')
    local_tz = None

    def __new__(cls, value=None):
        if isinstance(value, datetime.datetime):
            d = value
        elif isinstance(value, str):
            try:
                d = dateutil.parser.parse(value)
            except ValueError as e:
                raise ValueError("invalid value for DateTime object: %s" % value) from e
        else:
            raise ValueError("invalid value for DateTime object: %s" % value)

        if cls.local_tz is None:
            cls.local_tz = dateutil.tz.gettz(name=cls.local_timezone)

        tzinfo = d.tzinfo or cls.local_tz

        return super().__new__(cls, d.year, d.month, d.day, d.hour, d.minute, d.second,
                                    d.microsecond, tzinfo)

class TimeDelta(dateutil.relativedelta.relativedelta):
    _interval_tokens = {'year': 'years', 'years': 'years',
                        'month': 'months', 'months': 'months',
                        'week': 'weeks', 'weeks': 'weeks',
                        'day': 'days', 'days': 'days',
                        'hour': 'hours', 'hours': 'hours',
                        'minute': 'minutes', 'minutes': 'minutes',
                        'second': 'seconds', 'seconds': 'seconds',
                        'microsecond': 'microseconds', 'microseconds': 'microseconds'}

    def __new__(cls, value=None, *, dt1=None, dt2=None, years=0, months=0, days=0, leapdays=0,
                      weeks=0, hours=0, minutes=0, seconds=0, microseconds=0, year=None,
                      month=None, day=None, weekday=None, yearday=None, nlyearday=None, hour=None,
                      minute=None, second=None, microsecond=None):

        if value is None:
            result = super().__new__(cls)
            super().__init__(result,
                             dt1=dt1, dt2=dt2, years=years, months=months, days=days,
                             hours=hours, minutes=minutes,
                             seconds=seconds, microseconds=microseconds,
                             leapdays=leapdays, year=year, weeks=weeks,
                             month=month, day=day, weekday=weekday,
                             hour=hour, minute=minute, second=second,
                             microsecond=microsecond)

        elif isinstance(value, dateutil.relativedelta.relativedelta):
            result = super().__new__(cls)
            super().__init__(result,
                             years=value.years, months=value.months, days=value.days,
                             hours=value.hours, minutes=value.minutes,
                             seconds=value.seconds, microseconds=value.microseconds,
                             leapdays=value.leapdays, year=value.year,
                             month=value.month, day=value.day, weekday=value.weekday,
                             hour=value.hour, minute=value.minute,
                             second=value.second, microsecond=value.microsecond)

        elif isinstance(value, datetime.timedelta):
            result = super().__new__(cls)
            super().__init__(result,
                             days=value.days, seconds=value.seconds,
                             microseconds=value.microseconds)

        elif isinstance(value, str):
            expecting_number = True

            intervals = {}
            interval_tokens = cls._interval_tokens

            try:
                number = None

                for token in re.split('\s+', value):
                    if expecting_number:
                        number = int(token)
                    else:
                        interval = interval_tokens[token]
                        if interval in intervals:
                            raise ValueError
                        intervals[interval] = number
                    expecting_number = not expecting_number
                if not expecting_number:
                    raise ValueError
            except (ValueError, KeyError):
                raise ValueError("invalid timedelta value: '%s'" % value)

            result = super().__new__(cls)
            super().__init__(result, **intervals)

        return result

    def __init__(self, *args, **kwargs):
        # We initialize in __new__, so nothing to do here
        pass

    def __str__(self):
        s = []
        for k in ('years', 'months', 'days', 'hours', 'minutes', 'seconds', 'microseconds'):
            v = getattr(self, k)
            if v:
                s.append('%s %s' % (v, k if v != 1 else k[:-1]))
        return ' '.join(s)

    def __hash__(self):
        return hash((self.weekday, self.year, self.months, self.days, self.hours,
                     self.minutes, self.seconds, self.leapdays, self.year, self.month,
                     self.day, self.hour, self.minute, self.second, self.microsecond))

    def __gt__(self, other):
        if not isinstance(other, TimeDelta):
            return NotImplemented

        if self.weekday or other.weekday:
            raise TypeError('TimeDeltas with non-empty weekday are not orderable')

        return self.years > other.years or self.years == other.years and \
               (self.months > other.months or self.months == other.months and \
               (self.days > other.days or self.days == other.days and \
               (self.hours > other.hours or self.hours == other.hours and \
               (self.minutes > other.minutes or self.minutes == other.minutes and \
               (self.seconds > other.seconds or self.seconds == other.seconds and \
               (self.microseconds > other.microseconds))))))

    def __ge__(self, other):
        return self == other or self > other

    def __lt__(self, other):
        return not self >= other

    def __le__(self, other):
        return not self > other
<|endoftext|>


<commit_before>##
# Copyright (c) 2008-2010 Sprymix Inc.
# All rights reserved.
#
# See LICENSE for details.
##

import numbers
import postgresql.string
from semantix.caos.backends.pgsql import common
from semantix.utils.ast import codegen

class SQLSourceGeneratorError(Exception): pass

class SQLSourceGenerator(codegen.SourceGenerator):
    def generic_visit(self, node):
        raise SQLSourceGeneratorError('No method to generate code for %s' % node.__class__.__name__)

    def visit_CTENode(self, node):
        self.write(common.quote_ident(node.alias))

    def gen_ctes(self, ctes):
        self.write('WITH')
        count = len(ctes)
        for i, cte in enumerate(ctes):
            self.new_lines = 1
            self.write(common.quote_ident(cte.alias))
            self.write(' AS ')
            self.indentation += 2
            alias = cte.alias
            self.visit_SelectQueryNode(cte, use_alias=False)
            if i != count - 1:
                self.write(',')
            self.indentation -= 2

        self.new_lines = 1

    def visit_SelectQueryNode(self, node, use_alias=True):
        self.new_lines = 1

        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        self.write('SELECT')
        if node.distinct is not None:
            self.write(' DISTINCT')

        self.new_lines = 1
        self.indentation += 2

        count = len(node.targets)
        for i, target in enumerate(node.targets):
            self.new_lines = 1
            self.visit(target)
            if i != count -1:
                self.write(',')

        self.indentation -= 2

        if node.fromlist:
            self.indentation += 1
            self.new_lines = 1
            self.write('FROM')
            if node.from_only:
                self.write(' ONLY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.fromlist)
            for i, source in enumerate(node.fromlist):
                self.new_lines = 1
                self.visit(source)
                if i != count - 1:
                    self.write(',')

            self.indentation -= 2

        if node.where:
            self.indentation += 1
            self.new_lines = 1
            self.write('WHERE')
            self.new_lines = 1
            self.indentation += 1
            self.visit(node.where)
            self.indentation -= 2

        if node.groupby:
            self.indentation += 1
            self.new_lines = 1
            self.write('GROUP BY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.groupby)
            for i, groupexpr in enumerate(node.groupby):
                self.new_lines = 1
                self.visit(groupexpr)
                if i != count - 1:
                    self.write(',')
            self.indentation -= 2

        if node.having:
            self.indentation += 1
            self.new_lines = 1
            self.write('HAVING')
            self.new_lines = 1
            self.indentation += 1
            self.visit(node.having)
            self.indentation -= 2

        if node.orderby:
            self.indentation += 1
            self.new_lines = 1
            self.write('ORDER BY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.orderby)
            for i, sortexpr in enumerate(node.orderby):
                self.new_lines = 1
                self.visit(sortexpr)
                if i != count - 1:
                    self.write(',')
            self.indentation -= 2

        if node.offset:
            self.indentation += 1
            self.new_lines = 1
            self.write('OFFSET ')
            self.visit(node.offset)
            self.indentation -= 1

        if node.limit:
            self.indentation += 1
            self.new_lines = 1
            self.write('LIMIT ')
            self.visit(node.limit)
            self.indentation -= 1

        self.new_lines = 1
        self.write(')')

        if node.alias and use_alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_UnionNode(self, node):
        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        count = len(node.queries)
        for i, query in enumerate(node.queries):
            self.new_lines = 1
            self.visit(query)
            if i != count - 1:
                self.write(' UNION ')
                if not node.distinct:
                    self.write(' ALL ')

        self.write(')')
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_IntersectNode(self, node):
        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        count = len(node.queries)
        for i, query in enumerate(node.queries):
            self.new_lines = 1
            self.visit(query)
            if i != count - 1:
                self.write(' INTERSECT ')

        self.write(')')
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_ExistsNode(self, node):
        self.write('EXISTS (')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.expr)
        self.indentation -= 1
        self.new_lines = 1
        self.write(')')

    def visit_UpdateQueryNode(self, node):
        self.write('UPDATE ')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.fromexpr)
        self.indentation -= 1
        self.new_lines = 1
        self.write('SET')

        self.indentation += 1
        count = len(node.values)
        for i, expr in enumerate(node.values):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

        self.new_lines = 1
        self.write('WHERE')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.where)
        self.new_lines = 1
        self.indentation -= 1
        self.write('RETURNING')
        self.new_lines = 1
        self.indentation += 1

        count = len(node.targets)
        for i, expr in enumerate(node.targets):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

    def visit_DeleteQueryNode(self, node):
        self.write('DELETE FROM ')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.fromexpr)
        self.indentation -= 1
        self.new_lines = 1
        self.write('WHERE')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.where)
        self.new_lines = 1
        self.indentation -= 1
        self.write('RETURNING')
        self.new_lines = 1
        self.indentation += 1

        count = len(node.targets)
        for i, expr in enumerate(node.targets):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

    def visit_SelectExprNode(self, node):
        self.visit(node.expr)
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_UpdateExprNode(self, node):
        # Do no call visit_FieldRefNode here, since fields in UPDATE clause
        # must not have table qualifiers
        #
        self.write(common.qname(str(node.expr.field)))
        self.write(' = ')
        self.visit(node.value)

    def visit_FieldRefNode(self, node):
        if node.field == '*':
            self.write(common.quote_ident(node.table.alias) + '.' + str(node.field))
        else:
            if node.table:
                self.write(common.qname(node.table.alias, str(node.field)))
            else:
                self.write(common.quote_ident(str(node.field)))

    def visit_FromExprNode(self, node):
        self.visit(node.expr)
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_JoinNode(self, node):
        self.visit(node.left)
        self.new_lines = 1
        self.write(node.type.upper() + ' JOIN ')
        self.visit(node.right)
        self.write(' ON ')
        self.visit(node.condition)

    def visit_TableNode(self, node):
        self.write(common.qname(node.schema, node.name))
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_BinOpNode(self, node):
        self.write('(')
        self.visit(node.left)
        self.write(' ' + str(node.op).upper() + ' ')
        self.visit(node.right)
        self.write(')')

    def visit_PredicateNode(self, node):
        self.visit(node.expr)

    def visit_ConstantNode(self, node):
        if node.expr is not None:
            self.visit(node.expr)
        elif node.index is not None:
            self.write('$%d' % (node.index + 1))
            if node.type is not None:
                self.write('::%s' % node.type)
        else:
            if node.value is None:
                self.write('NULL')
            elif isinstance(node.value, (bool, numbers.Number)):
                self.write(str(node.value))
            else:
                self.write(postgresql.string.quote_literal(str(node.value)))

    def visit_SequenceNode(self, node):
        self.write('(')
        count = len(node.elements)
        for i, e in enumerate(node.elements):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_RowExprNode(self, node):
        self.write('(')
        count = len(node.args)
        for i, e in enumerate(node.args):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_ArrayNode(self, node):
        self.write('ARRAY[')
        count = len(node.elements)
        for i, e in enumerate(node.elements):
            self.visit(e)
            if i != count - 1:
                self.write(', ')
        self.write(']')

    def visit_FunctionCallNode(self, node):
        self.write(node.name)
        self.write('(')
        count = len(node.args)
        for i, e in enumerate(node.args):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_SortExprNode(self, node):
        self.visit(node.expr)
        if node.direction:
            self.write(' ' + node.direction)

    def visit_TypeCastNode(self, node):
        self.visit(node.expr)
        self.write('::')
        self.visit(node.type)

    def visit_TypeNode(self, node):
        self.write(node.name)

    def visit_StarIndirectionNode(self, node):
        self.write('*')
<commit_msg>caos.backends.pgsql.codegen: Plase NULLS LAST when sorting in desc order<commit_after>##
# Copyright (c) 2008-2010 Sprymix Inc.
# All rights reserved.
#
# See LICENSE for details.
##

import numbers
import postgresql.string
from semantix.caos.backends.pgsql import common
from semantix.utils.ast import codegen

class SQLSourceGeneratorError(Exception): pass

class SQLSourceGenerator(codegen.SourceGenerator):
    def generic_visit(self, node):
        raise SQLSourceGeneratorError('No method to generate code for %s' % node.__class__.__name__)

    def visit_CTENode(self, node):
        self.write(common.quote_ident(node.alias))

    def gen_ctes(self, ctes):
        self.write('WITH')
        count = len(ctes)
        for i, cte in enumerate(ctes):
            self.new_lines = 1
            self.write(common.quote_ident(cte.alias))
            self.write(' AS ')
            self.indentation += 2
            alias = cte.alias
            self.visit_SelectQueryNode(cte, use_alias=False)
            if i != count - 1:
                self.write(',')
            self.indentation -= 2

        self.new_lines = 1

    def visit_SelectQueryNode(self, node, use_alias=True):
        self.new_lines = 1

        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        self.write('SELECT')
        if node.distinct is not None:
            self.write(' DISTINCT')

        self.new_lines = 1
        self.indentation += 2

        count = len(node.targets)
        for i, target in enumerate(node.targets):
            self.new_lines = 1
            self.visit(target)
            if i != count -1:
                self.write(',')

        self.indentation -= 2

        if node.fromlist:
            self.indentation += 1
            self.new_lines = 1
            self.write('FROM')
            if node.from_only:
                self.write(' ONLY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.fromlist)
            for i, source in enumerate(node.fromlist):
                self.new_lines = 1
                self.visit(source)
                if i != count - 1:
                    self.write(',')

            self.indentation -= 2

        if node.where:
            self.indentation += 1
            self.new_lines = 1
            self.write('WHERE')
            self.new_lines = 1
            self.indentation += 1
            self.visit(node.where)
            self.indentation -= 2

        if node.groupby:
            self.indentation += 1
            self.new_lines = 1
            self.write('GROUP BY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.groupby)
            for i, groupexpr in enumerate(node.groupby):
                self.new_lines = 1
                self.visit(groupexpr)
                if i != count - 1:
                    self.write(',')
            self.indentation -= 2

        if node.having:
            self.indentation += 1
            self.new_lines = 1
            self.write('HAVING')
            self.new_lines = 1
            self.indentation += 1
            self.visit(node.having)
            self.indentation -= 2

        if node.orderby:
            self.indentation += 1
            self.new_lines = 1
            self.write('ORDER BY')
            self.new_lines = 1
            self.indentation += 1
            count = len(node.orderby)
            for i, sortexpr in enumerate(node.orderby):
                self.new_lines = 1
                self.visit(sortexpr)
                if i != count - 1:
                    self.write(',')
            self.indentation -= 2

        if node.offset:
            self.indentation += 1
            self.new_lines = 1
            self.write('OFFSET ')
            self.visit(node.offset)
            self.indentation -= 1

        if node.limit:
            self.indentation += 1
            self.new_lines = 1
            self.write('LIMIT ')
            self.visit(node.limit)
            self.indentation -= 1

        self.new_lines = 1
        self.write(')')

        if node.alias and use_alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_UnionNode(self, node):
        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        count = len(node.queries)
        for i, query in enumerate(node.queries):
            self.new_lines = 1
            self.visit(query)
            if i != count - 1:
                self.write(' UNION ')
                if not node.distinct:
                    self.write(' ALL ')

        self.write(')')
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_IntersectNode(self, node):
        self.write('(')
        if node.ctes:
            self.gen_ctes(node.ctes)

        count = len(node.queries)
        for i, query in enumerate(node.queries):
            self.new_lines = 1
            self.visit(query)
            if i != count - 1:
                self.write(' INTERSECT ')

        self.write(')')
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_ExistsNode(self, node):
        self.write('EXISTS (')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.expr)
        self.indentation -= 1
        self.new_lines = 1
        self.write(')')

    def visit_UpdateQueryNode(self, node):
        self.write('UPDATE ')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.fromexpr)
        self.indentation -= 1
        self.new_lines = 1
        self.write('SET')

        self.indentation += 1
        count = len(node.values)
        for i, expr in enumerate(node.values):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

        self.new_lines = 1
        self.write('WHERE')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.where)
        self.new_lines = 1
        self.indentation -= 1
        self.write('RETURNING')
        self.new_lines = 1
        self.indentation += 1

        count = len(node.targets)
        for i, expr in enumerate(node.targets):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

    def visit_DeleteQueryNode(self, node):
        self.write('DELETE FROM ')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.fromexpr)
        self.indentation -= 1
        self.new_lines = 1
        self.write('WHERE')
        self.new_lines = 1
        self.indentation += 1
        self.visit(node.where)
        self.new_lines = 1
        self.indentation -= 1
        self.write('RETURNING')
        self.new_lines = 1
        self.indentation += 1

        count = len(node.targets)
        for i, expr in enumerate(node.targets):
            self.new_lines = 1
            self.visit(expr)
            if i != count - 1:
                self.write(',')
        self.indentation -= 1

    def visit_SelectExprNode(self, node):
        self.visit(node.expr)
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_UpdateExprNode(self, node):
        # Do no call visit_FieldRefNode here, since fields in UPDATE clause
        # must not have table qualifiers
        #
        self.write(common.qname(str(node.expr.field)))
        self.write(' = ')
        self.visit(node.value)

    def visit_FieldRefNode(self, node):
        if node.field == '*':
            self.write(common.quote_ident(node.table.alias) + '.' + str(node.field))
        else:
            if node.table:
                self.write(common.qname(node.table.alias, str(node.field)))
            else:
                self.write(common.quote_ident(str(node.field)))

    def visit_FromExprNode(self, node):
        self.visit(node.expr)
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_JoinNode(self, node):
        self.visit(node.left)
        self.new_lines = 1
        self.write(node.type.upper() + ' JOIN ')
        self.visit(node.right)
        self.write(' ON ')
        self.visit(node.condition)

    def visit_TableNode(self, node):
        self.write(common.qname(node.schema, node.name))
        if node.alias:
            self.write(' AS ' + common.quote_ident(node.alias))

    def visit_BinOpNode(self, node):
        self.write('(')
        self.visit(node.left)
        self.write(' ' + str(node.op).upper() + ' ')
        self.visit(node.right)
        self.write(')')

    def visit_PredicateNode(self, node):
        self.visit(node.expr)

    def visit_ConstantNode(self, node):
        if node.expr is not None:
            self.visit(node.expr)
        elif node.index is not None:
            self.write('$%d' % (node.index + 1))
            if node.type is not None:
                self.write('::%s' % node.type)
        else:
            if node.value is None:
                self.write('NULL')
            elif isinstance(node.value, (bool, numbers.Number)):
                self.write(str(node.value))
            else:
                self.write(postgresql.string.quote_literal(str(node.value)))

    def visit_SequenceNode(self, node):
        self.write('(')
        count = len(node.elements)
        for i, e in enumerate(node.elements):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_RowExprNode(self, node):
        self.write('(')
        count = len(node.args)
        for i, e in enumerate(node.args):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_ArrayNode(self, node):
        self.write('ARRAY[')
        count = len(node.elements)
        for i, e in enumerate(node.elements):
            self.visit(e)
            if i != count - 1:
                self.write(', ')
        self.write(']')

    def visit_FunctionCallNode(self, node):
        self.write(node.name)
        self.write('(')
        count = len(node.args)
        for i, e in enumerate(node.args):
            self.visit(e)
            if i != count - 1:
                self.write(', ')

        self.write(')')

    def visit_SortExprNode(self, node):
        self.visit(node.expr)
        if node.direction:
            self.write(' ' + node.direction)
            if node.direction == 'desc':
                self.write(' NULLS LAST')

    def visit_TypeCastNode(self, node):
        self.visit(node.expr)
        self.write('::')
        self.visit(node.type)

    def visit_TypeNode(self, node):
        self.write(node.name)

    def visit_StarIndirectionNode(self, node):
        self.write('*')
<|endoftext|>


<commit_before>#The MIT License (MIT)
#
#Copyright (C) 2014 OpenBet Limited
#
#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,
#ITNESS 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 shutit_module import ShutItModule

#From: https://groups.google.com/forum/#!topic/docker-user/D0n-lURDn0o
# Expose port 5901 and 6080
class vnc(ShutItModule):

	def check_ready(self,shutit):
		# Only apt-based systems are supported support atm
		return shutit.cfg['container']['install_type'] == 'apt'

	def is_installed(self,shutit):
		return shutit.file_exists('/root/start_vnc.sh')

	def build(self,shutit):
		# TODO: distr-independence
		shutit.install('lsb-release')
		shutit.install('lsb-release')
		shutit.send_and_expect('lsb_release -s -c')
		release=shutit.get_output().split('\n')[1]
		if shutit.cfg['container']['install_type'] == 'apt':
			shutit.send_and_expect("""echo "deb http://archive.ubuntu.com/ubuntu """ + release + """ main universe multiverse" > /etc/apt/sources.list""")
			shutit.add_line_to_file("""deb http://archive.ubuntu.com/ubuntu/ """ + release + """-updates main restricted""","""/etc/apt/sources.list""")
			shutit.send_and_expect('apt-get update -qq',timeout=10000)
			shutit.send_and_expect('apt-get upgrade -y')
		shutit.install('gnome-terminal')
		shutit.install('openjdk-6-jre')
		shutit.install('xserver-xorg')
		shutit.install('vnc4server')
		shutit.install('novnc')
		shutit.install('apt-utils')
		# apt-utils?
		if shutit.cfg['container']['install_type'] == 'apt':
			send = 'apt-get install -qq -y --no-install-recommends ubuntu-desktop > /tmp/ubuntu-desktop && rm -f /tmp/ubuntu-desktop'
		while True:
			res = shutit.send_and_expect(send,['Unpacking','Setting up',shutit.cfg['expect_prompts']['root_prompt']],timeout=9999,check_exit=False)
			if res == 2:
				break
			elif res == 0 or res == 1:
				send = ''
		send = 'vncserver'
		while True:
			res = shutit.send_and_expect(send,expect=['assword:','erify',shutit.cfg['expect_prompts']['root_prompt']],check_exit=False,fail_on_empty_before=False,record_command=False)
			if res == 0 or res == 1:
				send = shutit.cfg['shutit.tk.vnc.vnc']['password']
			elif res == 2:
				break
		shutit.add_line_to_file('# start vnc','/root/start_vnc.sh')
		shutit.add_line_to_file('rm -rf /tmp/.X*','/root/start_vnc.sh')
		shutit.add_line_to_file("""vncserver << END
""" + shutit.cfg['shutit.tk.vnc.vnc']['password'] + """
""" + shutit.cfg['shutit.tk.vnc.vnc']['password'] + """
END""",'/root/start_vnc.sh')
		shutit.add_line_to_file('echo "Did you expose ports 5901 and 6080?"','/root/start_vnc.sh',match_regexp='echo .Did you expose ports 5901 and 6080..')
		shutit.add_line_to_file('echo "If so, then vncviewer localhost:1 should work."','/root/start_vnc.sh',match_regexp='echo .If so, then vncviewer localhost:1 should work..')
		shutit.add_line_to_file('# stop vnc','/root/stop_vnc.sh')
		shutit.add_line_to_file("""ps -ef | grep Xvnc4 | grep -v grep | awk '{print $2}' | xargs kill""",'/root/stop_vnc.sh')
		shutit.add_line_to_file('sleep 10','/root/stop_vnc.sh')
		shutit.add_line_to_file('rm -rf /tmp/.X*-lock','/root/stop_vnc.sh')
		shutit.send_and_expect('chmod +x /root/start_vnc.sh')
		shutit.send_and_expect('chmod +x /root/stop_vnc.sh')
		return True

	def start(self,shutit):
		shutit.send_and_expect('/root/start_vnc.sh',check_exit=False)
		return True

	def stop(self,shutit):
		shutit.send_and_expect('/root/stop_vnc.sh',check_exit=False)
		return True

	def get_config(self,shutit):
		shutit.get_config('shutit.tk.vnc.vnc','password','vncpass')
		# Set build to privileged
		shutit.cfg['build']['privileged'] = True
		return True

def module():
	return vnc(
		'shutit.tk.vnc.vnc', 0.322,
		description='vnc server. contains instructions for use within /root/start_vnc.sh output.',
		depends=['shutit.tk.setup']
	)

<commit_msg>simplify<commit_after>#The MIT License (MIT)
#
#Copyright (C) 2014 OpenBet Limited
#
#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,
#ITNESS 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 shutit_module import ShutItModule

#From: https://groups.google.com/forum/#!topic/docker-user/D0n-lURDn0o
# Expose port 5901 and 6080
class vnc(ShutItModule):

	def check_ready(self,shutit):
		# Only apt-based systems are supported support atm
		return shutit.cfg['container']['install_type'] == 'apt'

	def is_installed(self,shutit):
		return shutit.file_exists('/root/start_vnc.sh')

	def build(self,shutit):
		# TODO: distr-independence
		shutit.install('lsb-release')
		shutit.install('lsb-release')
		shutit.send_and_expect('lsb_release -s -c')
		release=shutit.get_output().split('\n')[1]
		if shutit.cfg['container']['install_type'] == 'apt':
			shutit.send_and_expect("""echo "deb http://archive.ubuntu.com/ubuntu """ + release + """ main universe multiverse" > /etc/apt/sources.list""")
			shutit.add_line_to_file("""deb http://archive.ubuntu.com/ubuntu/ """ + release + """-updates main restricted""","""/etc/apt/sources.list""")
			shutit.send_and_expect('apt-get update -qq',timeout=10000)
			shutit.send_and_expect('apt-get upgrade -y')
		shutit.install('gnome-terminal')
		shutit.install('openjdk-6-jre')
		shutit.install('xserver-xorg')
		shutit.install('vnc4server')
		shutit.install('novnc')
		shutit.install('apt-utils')
		# apt-utils?
		if shutit.cfg['container']['install_type'] == 'apt':
			send = 'apt-get install -qq -y --no-install-recommends ubuntu-desktop > /tmp/ubuntu-desktop && rm -f /tmp/ubuntu-desktop'
		while True:
			res = shutit.send_and_expect(send,expect=['Unpacking','Setting up',shutit.cfg['expect_prompts']['root_prompt']],timeout=9999)
			if res == 2:
				break
			elif res == 0 or res == 1:
				send = ''
		send = 'vncserver'
		while True:
			res = shutit.send_and_expect(send,expect=['assword:','erify',shutit.cfg['expect_prompts']['root_prompt']],fail_on_empty_before=False,record_command=False)
			if res == 0 or res == 1:
				send = shutit.cfg['shutit.tk.vnc.vnc']['password']
			elif res == 2:
				break
		shutit.add_line_to_file('# start vnc','/root/start_vnc.sh')
		shutit.add_line_to_file('rm -rf /tmp/.X*','/root/start_vnc.sh')
		shutit.add_line_to_file("""vncserver << END
""" + shutit.cfg['shutit.tk.vnc.vnc']['password'] + """
""" + shutit.cfg['shutit.tk.vnc.vnc']['password'] + """
END""",'/root/start_vnc.sh')
		shutit.add_line_to_file('echo "Did you expose ports 5901 and 6080?"','/root/start_vnc.sh',match_regexp='echo .Did you expose ports 5901 and 6080..')
		shutit.add_line_to_file('echo "If so, then vncviewer localhost:1 should work."','/root/start_vnc.sh',match_regexp='echo .If so, then vncviewer localhost:1 should work..')
		shutit.add_line_to_file('# stop vnc','/root/stop_vnc.sh')
		shutit.add_line_to_file("""ps -ef | grep Xvnc4 | grep -v grep | awk '{print $2}' | xargs kill""",'/root/stop_vnc.sh')
		shutit.add_line_to_file('sleep 10','/root/stop_vnc.sh')
		shutit.add_line_to_file('rm -rf /tmp/.X*-lock','/root/stop_vnc.sh')
		shutit.send_and_expect('chmod +x /root/start_vnc.sh')
		shutit.send_and_expect('chmod +x /root/stop_vnc.sh')
		return True

	def start(self,shutit):
		shutit.send_and_expect('/root/start_vnc.sh',check_exit=False)
		return True

	def stop(self,shutit):
		shutit.send_and_expect('/root/stop_vnc.sh',check_exit=False)
		return True

	def get_config(self,shutit):
		shutit.get_config('shutit.tk.vnc.vnc','password','vncpass')
		# Set build to privileged
		shutit.cfg['build']['privileged'] = True
		return True

def module():
	return vnc(
		'shutit.tk.vnc.vnc', 0.322,
		description='vnc server. contains instructions for use within /root/start_vnc.sh output.',
		depends=['shutit.tk.setup']
	)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
saferedisqueue
~~~~~~~~~~~~~~

A small wrapper around Redis that provides access to a FIFO queue and allows
upstream code to mark pop'ed items as processed successfully (ACK'ing) or
unsucessfully (FAIL'ing). Failed items are automatically requeued. Addionally
a backup is kept for items that were neither ack'ed nor fail'ed, i.e. in case
the consumer crashes. The backup items will be requeued as soon as one of the
consumer(s) comes up again.

All actions are atomic and it is safe to have multiple producers and consumers
accessing the same queue concurrently.

:copyright: (c) Ferret Go GmbH, Sean Buttinger, Fabian Neumann
:license: All Rights Reserved. For now.
"""

__version__ = '0.2.0-alpha'

import uuid

import redis

class SafeRedisQueue(object):

    AUTOCLEAN_TIMEOUT = 60

    def __init__(self, *args, **kw):
        prefix = kw.pop('prefix', 'srq')
        self.QUEUE_KEY = '%s:queue' % prefix
        self.ITEM_KEY_PREFIX = '%s:item' % prefix
        self.ACKBUF_KEY = '%s:ackbuf' % prefix
        self.BACKUP = '%s:backup' % prefix
        self.BACKUP_LOCK = '%s:lock:backup' % prefix
        self._redis = redis.StrictRedis(*args, **kw)

    def _item_key(self, uid):
        return '%s:%s' % (self.ITEM_KEY_PREFIX, uid)

    def _autoclean(self):
        if self._redis.exists(self.BACKUP_LOCK):
            return
        if self._redis.exists(self.BACKUP):
            # Get the lock
            if self._redis.setnx(self.BACKUP_LOCK, 1):
                # There's a potential deadlock here:
                # If we crash during the while loop, the BACKUP_LOCK will never
                # be set to expire, blocking the creation of a new backup queue.
                # Could be solved by setting the expire immediately, although
                # then the guard could expire before we are done with the
                # requeuing. Although that is very unlikely.
                while self._redis.rpoplpush(self.BACKUP, self.QUEUE_KEY):
                    pass
                self._redis.pipeline()\
                    .rename(self.ACKBUF_KEY, self.BACKUP)\
                    .expire(self.BACKUP_LOCK, self.AUTOCLEAN_TIMEOUT)\
                    .execute()
        else:
            with self._redis.pipeline() as pipe:
                try:
                    pipe.watch(self.BACKUP_LOCK)
                    if not pipe.exists(self.BACKUP_LOCK):
                        pipe.multi()
                        pipe.renamenx(self.ACKBUF_KEY, self.BACKUP)\
                            .setex(self.BACKUP_LOCK, self.AUTOCLEAN_TIMEOUT, 1)\
                            .execute()
                    else:
                        pipe.unwatch()
                except redis.WatchError:
                    pass

    def push_item(self, item):
        """Adds an item to the queue.

        Generates a uid for later referral.
        Enqueues the uid and stores the item.
        """
        uid = uuid.uuid4()
        self._redis.pipeline()\
                .set(self._item_key(uid), item)\
                .lpush(self.QUEUE_KEY, uid)\
                .execute()

    def get_item(self, uid):
        """Get item for uid.
        """
        item = self._redis.get(self._item_key(uid))
        return item

    def pop_item(self, timeout=0):
        """Get next item from queue. Blocks if queue is empty.

        Pops uid from queue and writes it to ackbuffer.
        Returns the corresponding item.
        """
        self._autoclean()
        uid = self._redis.brpoplpush(self.QUEUE_KEY, self.ACKBUF_KEY, timeout)
        item = self.get_item(uid)
        return uid, item

    def ack_item(self, uid):
        """Acknowledge item as successfully consumed.

        Removes uid from ackbuffer and deletes the corresponding item.
        """
        self._redis.pipeline()\
                   .lrem(self.ACKBUF_KEY, 0, uid)\
                   .lrem(self.BACKUP, 0, uid)\
                   .delete(self._item_key(uid))\
                   .execute()

    def fail_item(self, uid):
        """Report item as not successfully consumed.

        Removes uid from ackbuffer and re-enqueues it.
        """
        self._redis.pipeline()\
                   .lrem(self.ACKBUF_KEY, 0, uid)\
                   .lrem(self.BACKUP, 0, uid)\
                   .lpush(self.QUEUE_KEY, uid)\
                   .execute()

if __name__ == "__main__":
    import sys
    queue = SafeRedisQueue()

    def _usage():
        sys.stdout.writelines([
            'Commandline usage: python saferedisqueue.py <consumer | producer>\n'
            'Example:\n'
            '    $ echo "Hello World!" | python saferedisqueue.py producer\n'
            '    $ python saferedisqueue.py consumer\n'
            '    cbdabbc8-1c0f-4eb0-8733-fdb62a9c0fa6 Hello World!\n'
        ])
        sys.exit(1)

    if len(sys.argv) != 2:
        _usage()

    if sys.argv[1] == 'producer':
        for line in sys.stdin.readlines():
            queue.push_item(line.strip())
    elif sys.argv[1] == 'consumer':
        while True:
            uid, item = queue.pop_item()
            print uid, item
    else:
        _usage()
<commit_msg>Add "demo" as third option for the command line.<commit_after># -*- coding: utf-8 -*-
"""
saferedisqueue
~~~~~~~~~~~~~~

A small wrapper around Redis that provides access to a FIFO queue and allows
upstream code to mark pop'ed items as processed successfully (ACK'ing) or
unsucessfully (FAIL'ing). Failed items are automatically requeued. Addionally
a backup is kept for items that were neither ack'ed nor fail'ed, i.e. in case
the consumer crashes. The backup items will be requeued as soon as one of the
consumer(s) comes up again.

All actions are atomic and it is safe to have multiple producers and consumers
accessing the same queue concurrently.

:copyright: (c) Ferret Go GmbH, Sean Buttinger, Fabian Neumann
:license: All Rights Reserved. For now.
"""

__version__ = '0.2.0-alpha'

import uuid

import redis

class SafeRedisQueue(object):

    AUTOCLEAN_TIMEOUT = 60

    def __init__(self, *args, **kw):
        prefix = kw.pop('prefix', 'srq')
        self.QUEUE_KEY = '%s:queue' % prefix
        self.ITEM_KEY_PREFIX = '%s:item' % prefix
        self.ACKBUF_KEY = '%s:ackbuf' % prefix
        self.BACKUP = '%s:backup' % prefix
        self.BACKUP_LOCK = '%s:lock:backup' % prefix
        self._redis = redis.StrictRedis(*args, **kw)

    def _item_key(self, uid):
        return '%s:%s' % (self.ITEM_KEY_PREFIX, uid)

    def _autoclean(self):
        if self._redis.exists(self.BACKUP_LOCK):
            return
        if self._redis.exists(self.BACKUP):
            # Get the lock
            if self._redis.setnx(self.BACKUP_LOCK, 1):
                # There's a potential deadlock here:
                # If we crash during the while loop, the BACKUP_LOCK will never
                # be set to expire, blocking the creation of a new backup queue.
                # Could be solved by setting the expire immediately, although
                # then the guard could expire before we are done with the
                # requeuing. Although that is very unlikely.
                while self._redis.rpoplpush(self.BACKUP, self.QUEUE_KEY):
                    pass
                self._redis.pipeline()\
                    .rename(self.ACKBUF_KEY, self.BACKUP)\
                    .expire(self.BACKUP_LOCK, self.AUTOCLEAN_TIMEOUT)\
                    .execute()
        else:
            with self._redis.pipeline() as pipe:
                try:
                    pipe.watch(self.BACKUP_LOCK)
                    if not pipe.exists(self.BACKUP_LOCK):
                        pipe.multi()
                        pipe.renamenx(self.ACKBUF_KEY, self.BACKUP)\
                            .setex(self.BACKUP_LOCK, self.AUTOCLEAN_TIMEOUT, 1)\
                            .execute()
                    else:
                        pipe.unwatch()
                except redis.WatchError:
                    pass

    def push_item(self, item):
        """Adds an item to the queue.

        Generates a uid for later referral.
        Enqueues the uid and stores the item.
        """
        uid = uuid.uuid4()
        self._redis.pipeline()\
                .set(self._item_key(uid), item)\
                .lpush(self.QUEUE_KEY, uid)\
                .execute()

    def get_item(self, uid):
        """Get item for uid.
        """
        item = self._redis.get(self._item_key(uid))
        return item

    def pop_item(self, timeout=0):
        """Get next item from queue. Blocks if queue is empty.

        Pops uid from queue and writes it to ackbuffer.
        Returns the corresponding item.
        """
        self._autoclean()
        uid = self._redis.brpoplpush(self.QUEUE_KEY, self.ACKBUF_KEY, timeout)
        item = self.get_item(uid)
        return uid, item

    def ack_item(self, uid):
        """Acknowledge item as successfully consumed.

        Removes uid from ackbuffer and deletes the corresponding item.
        """
        self._redis.pipeline()\
                   .lrem(self.ACKBUF_KEY, 0, uid)\
                   .lrem(self.BACKUP, 0, uid)\
                   .delete(self._item_key(uid))\
                   .execute()

    def fail_item(self, uid):
        """Report item as not successfully consumed.

        Removes uid from ackbuffer and re-enqueues it.
        """
        self._redis.pipeline()\
                   .lrem(self.ACKBUF_KEY, 0, uid)\
                   .lrem(self.BACKUP, 0, uid)\
                   .lpush(self.QUEUE_KEY, uid)\
                   .execute()

if __name__ == "__main__":
    import sys

    def _usage():
        sys.stdout.writelines([
            'Commandline usage: python saferedisqueue.py <consumer | producer | demo>\n'
            'Example:\n'
            '    $ echo "Hello World" | python saferedisqueue.py producer\n'
            '    $ python saferedisqueue.py consumer\n'
            '    cbdabbc8-1c0f-4eb0-8733-fdb62a9c0fa6 Hello World\n'
        ])
        sys.exit(1)

    if len(sys.argv) != 2:
        _usage()

    queue = SafeRedisQueue()

    if sys.argv[1] == 'producer':
        for line in sys.stdin.readlines():
            queue.push_item(line.strip())
    elif sys.argv[1] == 'consumer':
        while True:
            uid, item = queue.pop_item(timeout=1)
            queue.ack_item(uid)
            print uid, item
    elif sys.argv[1] == 'demo':
        map(queue.push_item, ['Hello', 'World'])
        while True:
            uid, item = queue.pop_item(timeout=1)
            queue.ack_item(uid)
            print uid, item
    else:
        _usage()
<|endoftext|>