<commit_before>import io
from enum import Enum

from bidict import bidict
import PIL.ExifTags
import PIL.Image
import zopfli

from madam.core import operator, OperatorError
from madam.core import Asset, Processor
from madam.mime import MimeType

class ResizeMode(Enum):
    """
    Represents a behavior for image resize operations.
    """
    #: Image exactly matches the specified dimensions
    EXACT = 0
    #: Image is resized to fit completely into the specified dimensions
    FIT = 1
    #: Image is resized to completely fill the specified dimensions
    FILL = 2

class FlipOrientation(Enum):
    """
    Represents an axis for image flip operations.
    """
    #: Horizontal axis
    HORIZONTAL = 0
    #: Vertical axis
    VERTICAL = 1

def _optimized(image, pil_format, **pil_options):
    r"""
    Writes an optimized version of a Pillow image to a buffer.

    :param image: PIL image to be written
    :type image: PIL.Image.Image
    :param pil_format: PIL file format name
    :type pil_format: str
    :param \\**pil_options: pil_options
    :type \\**pil_options: dict
    :return: Buffer with image data
    :rtype: file-like object
    """
    image_buffer = io.BytesIO()

    if pil_format == 'PNG' and image.mode != 'P':
        zopfli_png = zopfli.ZopfliPNG()
        # Convert 16-bit per channel images to 8-bit per channel
        zopfli_png.lossy_8bit = False
        # Allow altering hidden colors of fully transparent pixels
        zopfli_png.lossy_transparent = True
        # Use all available optimization strategies
        zopfli_png.filter_strategies = '0me'

        pil_options.pop('optimize', False)
        essence = io.BytesIO()
        image.save(essence, 'PNG', optimize=False, **pil_options)
        essence.seek(0)
        optimized_data = zopfli_png.optimize(essence.read())
        image_buffer.write(optimized_data)
    elif pil_format == 'TIFF' and image.mode == 'P':
        pil_options.pop('compression', '')
        image.save(image_buffer, pil_format, **pil_options)
    else:
        image.save(image_buffer, pil_format, **pil_options)

    image_buffer.seek(0)
    return image_buffer

class PillowProcessor(Processor):
    """
    Represents a processor that uses Pillow as a backend.
    """
    __mime_type_to_pillow_type = bidict({
        MimeType('image/bmp'): 'BMP',
        MimeType('image/gif'): 'GIF',
        MimeType('image/jpeg'): 'JPEG',
        MimeType('image/png'): 'PNG',
        MimeType('image/tiff'): 'TIFF',
        MimeType('image/webp'): 'WEBP',
    })

    __format_options = {
        MimeType('image/gif'): dict(
            optimize=True,
        ),
        MimeType('image/jpeg'): dict(
            optimize=True,
            progressive=True,
        ),
        MimeType('image/png'): dict(
            optimize=True,
        ),
        MimeType('image/tiff'): dict(
            compression='tiff_deflate',
        ),
        MimeType('image/webp'): dict(
            method=6,
        ),
    }

    __pillow_mode_to_color_mode = bidict({
        '1': ('LUMA', 1, 'uint'),
        'L': ('LUMA', 8, 'uint'),
        'P': ('PALETTE', 8, 'uint'),
        'RGB': ('RGB', 8, 'uint'),
        'RGBA': ('RGBA', 8, 'uint'),
        'RGBX': ('RGBX', 8, 'uint'),
        'CMYK': ('CMYK', 8, 'uint'),
        'YCbCr': ('YCbCr', 8, 'uint'),
        'LAB': ('LAB', 8, 'uint'),
        'HSV': ('HSV', 8, 'uint'),
        'I;16': ('LUMA', 16, 'uint'),
        'I': ('LUMA', 32, 'uint'),
        'F': ('LUMA', 32, 'float'),
    })

    def __init__(self):
        """
        Initializes a new `PillowProcessor`.
        """
        super().__init__()

    def read(self, file):
        image = PIL.Image.open(file)
        mime_type = PillowProcessor.__mime_type_to_pillow_type.inv[image.format]
        color_space, bit_depth, data_type = PillowProcessor.__pillow_mode_to_color_mode[image.mode]
        metadata = dict(
            mime_type=str(mime_type),
            width=image.width,
            height=image.height,
            color_space=color_space,
            depth=bit_depth,
            data_type=data_type,
        )
        file.seek(0)
        asset = Asset(file, **metadata)
        return asset

    def can_read(self, file):
        try:
            PIL.Image.open(file)
            file.seek(0)
            return True
        except IOError:
            return False

    @operator
    def resize(self, asset, width, height, mode=ResizeMode.EXACT):
        """
        Creates a new Asset whose essence is resized according to the specified parameters.

        :param asset: Asset to be resized
        :type asset: Asset
        :param width: target width
        :type width: int
        :param height: target height
        :type height: int
        :param mode: resize behavior
        :type mode: ResizeMode
        :return: Asset with resized essence
        :rtype: Asset
        """
        image = PIL.Image.open(asset.essence)
        mime_type = MimeType(asset.mime_type)
        if mode == ResizeMode.EXACT:
            resized_width = width
            resized_height = height
        else:
            aspect = asset.width / asset.height
            aspect_target = width / height
            if mode == ResizeMode.FIT and aspect >= aspect_target or \
               mode == ResizeMode.FILL and aspect <= aspect_target:
                resize_factor = width / image.width
            else:
                resize_factor = height / image.height
            resized_width = max(1, round(resize_factor * image.width))
            resized_height = max(1, round(resize_factor * image.height))
        # Pillow supports resampling only for 8-bit images
        resampling_method = PIL.Image.LANCZOS if asset.depth == 8 else PIL.Image.NEAREST
        resized_image = image.resize((resized_width, resized_height),
                                     resample=resampling_method)
        resized_asset = self._image_to_asset(resized_image, mime_type=mime_type)
        return resized_asset

    def _image_to_asset(self, image, mime_type):
        """
        Converts an PIL image to a MADAM asset. The conversion can also include
        a change in file type.

        :param image: PIL image
        :type image: PIL.Image.Image
        :param mime_type: MIME type of the target asset
        :type mime_type: MimeType
        :return: MADAM asset with hte specified MIME type
        :rtype: Asset
        """
        mime_type = MimeType(mime_type)
        pil_format = PillowProcessor.__mime_type_to_pillow_type[mime_type]
        pil_options = PillowProcessor.__format_options.get(mime_type, {})
        image_buffer = _optimized(image, pil_format, **pil_options)
        asset = self.read(image_buffer)
        return asset

    def _rotate(self, asset, rotation):
        """
        Creates a new image asset from specified asset whose essence is rotated
        by the specified rotation.

        :param asset: Image asset to be rotated
        :type asset: Asset
        :param rotation: One of `PIL.Image.FLIP_LEFT_RIGHT`,
        `PIL.Image.FLIP_TOP_BOTTOM`, `PIL.Image.ROTATE_90`,
        `PIL.Image.ROTATE_180`, `PIL.Image.ROTATE_270`, or
        `PIL.Image.TRANSPOSE`
        :return: New image asset with rotated essence
        :rtype: Asset
        """
        image = PIL.Image.open(asset.essence)
        mime_type = MimeType(asset.mime_type)
        transposed_image = image.transpose(rotation)
        transposed_asset = self._image_to_asset(transposed_image, mime_type=mime_type)
        return transposed_asset

    @operator
    def transpose(self, asset):
        """
        Creates a new image asset whose essence is the transpose of the
        specified asset's essence.

        :param asset: Image asset whose essence is to be transposed
        :type asset: Asset
        :return: New image asset with transposed essence
        :rtype: Asset
        """
        return self._rotate(asset, PIL.Image.TRANSPOSE)

    @operator
    def flip(self, asset, orientation):
        """
        Creates a new asset whose essence is flipped according the specified orientation.

        :param asset: Asset whose essence is to be flipped
        :type asset: Asset
        :param orientation: axis of the flip operation
        :type orientation: FlipOrientation
        :return: Asset with flipped essence
        :rtype: Asset
        """
        if orientation == FlipOrientation.HORIZONTAL:
            flip_orientation = PIL.Image.FLIP_LEFT_RIGHT
        else:
            flip_orientation = PIL.Image.FLIP_TOP_BOTTOM
        return self._rotate(asset, flip_orientation)

    @operator
    def auto_orient(self, asset):
        """
        Creates a new asset whose essence is rotated according to the Exif
        orientation. If no orientation metadata exists or asset is not rotated,
        an identical asset object is returned.

        :param asset: Asset with orientation metadata
        :type asset: Asset
        :return: Asset with rotated essence
        :rtype: Asset
        """
        orientation = asset.metadata.get('exif', {}).get('orientation')
        if orientation is None or orientation == 1:
            return asset

        flip_horizontally = self.flip(orientation=FlipOrientation.HORIZONTAL)
        flip_vertically = self.flip(orientation=FlipOrientation.VERTICAL)

        if orientation == 2:
            oriented_asset = flip_horizontally(asset)
        elif orientation == 3:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_180)
        elif orientation == 4:
            oriented_asset = flip_vertically(asset)
        elif orientation == 5:
            oriented_asset = flip_vertically(self._rotate(asset, PIL.Image.ROTATE_90))
        elif orientation == 6:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_270)
        elif orientation == 7:
            oriented_asset = flip_horizontally(self._rotate(asset, PIL.Image.ROTATE_90))
        elif orientation == 8:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_90)
        else:
            raise OperatorError('Unable to correct image orientation with value %s' % orientation)

        return oriented_asset

    @operator
    def convert(self, asset, mime_type, color_space=None, depth=None, data_type=None):
        """
        Creates a new asset of the specified MIME type from the essence of the
        specified asset.

        :param asset: Asset whose contents will be converted
        :type asset: Asset
        :param mime_type: Target MIME type
        :type mime_type: MimeType or str
        :param color_space: Name of color space
        :type color_space: str or None
        :param depth: Bit depth per channel
        :type depth: int or None
        :param data_type: Data type of the pixels, e.g. 'uint' or 'float'
        :type data_type: str or None
        :return: New asset with converted essence
        :rtype: Asset
        """
        mime_type = MimeType(mime_type)
        try:
            image = PIL.Image.open(asset.essence)
            color_mode = color_space or asset.color_space, depth or asset.depth, data_type or asset.data_type
            pil_mode = PillowProcessor.__pillow_mode_to_color_mode.inv.get(color_mode)
            if pil_mode is not None:
                image = image.convert(pil_mode)
            converted_asset = self._image_to_asset(image, mime_type)
        except (IOError, KeyError) as pil_error:
            raise OperatorError('Could not convert image to %s: %s' %
                                (mime_type, pil_error))

        return converted_asset

    @operator
    def crop(self, asset, x, y, width, height):
        """
        Creates a new asset whose essence is cropped to the specified
        rectangular area.

        :param asset: Asset whose contents will be cropped
        :type asset: Asset
        :param x: horizontal offset of the cropping area from left
        :type x: int
        :param y: vertical offset of the cropping area from top
        :type y: int
        :param width: width of the cropping area
        :type width: int
        :param height: height of the cropping area
        :type height: int
        :return: New asset with cropped essence
        :rtype: Asset
        """
        if x == 0 and y == 0 and width == asset.width and height == asset.height:
            return asset

        max_x = max(0, min(asset.width, width + x))
        max_y = max(0, min(asset.height, height + y))
        min_x = max(0, min(asset.width, x))
        min_y = max(0, min(asset.height, y))

        if min_x == asset.width or min_y == asset.height or max_x <= min_x or max_y <= min_y:
            raise OperatorError('Invalid cropping area: <x=%r, y=%r, width=%r, height=%r>' % (x, y, width, height))

        image = PIL.Image.open(asset.essence)
        cropped_image = image.crop(box=(min_x, min_y, max_x, max_y))
        cropped_asset = self._image_to_asset(cropped_image, mime_type=asset.mime_type)

        return cropped_asset

    @operator
    def rotate(self, asset, angle, expand=False):
        """
        Creates an asset whose essence is rotated by the specified angle in
        degrees.

        .. warning:: The color model will be changed to RGB when applying
            this operation

        :param asset: Asset whose contents will be rotated
        :type asset: Asset
        :param angle: Angle in degrees, counter clockwise
        :type angle: float
        :param expand: If true, changes the dimensions of the new asset so it
            can hold the entire rotated essence, otherwise the dimensions of
            the original asset will be used.
        :type expand: bool
        :return: New asset with rotated essence
        :rtype: Asset
        """
        if angle % 360.0 == 0.0:
            return asset

        image = PIL.Image.open(asset.essence).convert('RGB')
        rotated_image = image.rotate(angle=angle, resample=PIL.Image.BICUBIC, expand=expand)
        rotated_asset = self._image_to_asset(rotated_image, mime_type=asset.mime_type)

        return rotated_asset
<commit_msg>PillowProcessor.rotate() does not implicitly change color space.<commit_after>import io
from enum import Enum

from bidict import bidict
import PIL.ExifTags
import PIL.Image
import zopfli

from madam.core import operator, OperatorError
from madam.core import Asset, Processor
from madam.mime import MimeType

class ResizeMode(Enum):
    """
    Represents a behavior for image resize operations.
    """
    #: Image exactly matches the specified dimensions
    EXACT = 0
    #: Image is resized to fit completely into the specified dimensions
    FIT = 1
    #: Image is resized to completely fill the specified dimensions
    FILL = 2

class FlipOrientation(Enum):
    """
    Represents an axis for image flip operations.
    """
    #: Horizontal axis
    HORIZONTAL = 0
    #: Vertical axis
    VERTICAL = 1

def _optimized(image, pil_format, **pil_options):
    r"""
    Writes an optimized version of a Pillow image to a buffer.

    :param image: PIL image to be written
    :type image: PIL.Image.Image
    :param pil_format: PIL file format name
    :type pil_format: str
    :param \\**pil_options: pil_options
    :type \\**pil_options: dict
    :return: Buffer with image data
    :rtype: file-like object
    """
    image_buffer = io.BytesIO()

    if pil_format == 'PNG' and image.mode != 'P':
        zopfli_png = zopfli.ZopfliPNG()
        # Convert 16-bit per channel images to 8-bit per channel
        zopfli_png.lossy_8bit = False
        # Allow altering hidden colors of fully transparent pixels
        zopfli_png.lossy_transparent = True
        # Use all available optimization strategies
        zopfli_png.filter_strategies = '0me'

        pil_options.pop('optimize', False)
        essence = io.BytesIO()
        image.save(essence, 'PNG', optimize=False, **pil_options)
        essence.seek(0)
        optimized_data = zopfli_png.optimize(essence.read())
        image_buffer.write(optimized_data)
    elif pil_format == 'TIFF' and image.mode == 'P':
        pil_options.pop('compression', '')
        image.save(image_buffer, pil_format, **pil_options)
    else:
        image.save(image_buffer, pil_format, **pil_options)

    image_buffer.seek(0)
    return image_buffer

class PillowProcessor(Processor):
    """
    Represents a processor that uses Pillow as a backend.
    """
    __mime_type_to_pillow_type = bidict({
        MimeType('image/bmp'): 'BMP',
        MimeType('image/gif'): 'GIF',
        MimeType('image/jpeg'): 'JPEG',
        MimeType('image/png'): 'PNG',
        MimeType('image/tiff'): 'TIFF',
        MimeType('image/webp'): 'WEBP',
    })

    __format_options = {
        MimeType('image/gif'): dict(
            optimize=True,
        ),
        MimeType('image/jpeg'): dict(
            optimize=True,
            progressive=True,
        ),
        MimeType('image/png'): dict(
            optimize=True,
        ),
        MimeType('image/tiff'): dict(
            compression='tiff_deflate',
        ),
        MimeType('image/webp'): dict(
            method=6,
        ),
    }

    __pillow_mode_to_color_mode = bidict({
        '1': ('LUMA', 1, 'uint'),
        'L': ('LUMA', 8, 'uint'),
        'P': ('PALETTE', 8, 'uint'),
        'RGB': ('RGB', 8, 'uint'),
        'RGBA': ('RGBA', 8, 'uint'),
        'RGBX': ('RGBX', 8, 'uint'),
        'CMYK': ('CMYK', 8, 'uint'),
        'YCbCr': ('YCbCr', 8, 'uint'),
        'LAB': ('LAB', 8, 'uint'),
        'HSV': ('HSV', 8, 'uint'),
        'I;16': ('LUMA', 16, 'uint'),
        'I': ('LUMA', 32, 'uint'),
        'F': ('LUMA', 32, 'float'),
    })

    def __init__(self):
        """
        Initializes a new `PillowProcessor`.
        """
        super().__init__()

    def read(self, file):
        image = PIL.Image.open(file)
        mime_type = PillowProcessor.__mime_type_to_pillow_type.inv[image.format]
        color_space, bit_depth, data_type = PillowProcessor.__pillow_mode_to_color_mode[image.mode]
        metadata = dict(
            mime_type=str(mime_type),
            width=image.width,
            height=image.height,
            color_space=color_space,
            depth=bit_depth,
            data_type=data_type,
        )
        file.seek(0)
        asset = Asset(file, **metadata)
        return asset

    def can_read(self, file):
        try:
            PIL.Image.open(file)
            file.seek(0)
            return True
        except IOError:
            return False

    @operator
    def resize(self, asset, width, height, mode=ResizeMode.EXACT):
        """
        Creates a new Asset whose essence is resized according to the specified parameters.

        :param asset: Asset to be resized
        :type asset: Asset
        :param width: target width
        :type width: int
        :param height: target height
        :type height: int
        :param mode: resize behavior
        :type mode: ResizeMode
        :return: Asset with resized essence
        :rtype: Asset
        """
        image = PIL.Image.open(asset.essence)
        mime_type = MimeType(asset.mime_type)
        if mode == ResizeMode.EXACT:
            resized_width = width
            resized_height = height
        else:
            aspect = asset.width / asset.height
            aspect_target = width / height
            if mode == ResizeMode.FIT and aspect >= aspect_target or \
               mode == ResizeMode.FILL and aspect <= aspect_target:
                resize_factor = width / image.width
            else:
                resize_factor = height / image.height
            resized_width = max(1, round(resize_factor * image.width))
            resized_height = max(1, round(resize_factor * image.height))
        # Pillow supports resampling only for 8-bit images
        resampling_method = PIL.Image.LANCZOS if asset.depth == 8 else PIL.Image.NEAREST
        resized_image = image.resize((resized_width, resized_height),
                                     resample=resampling_method)
        resized_asset = self._image_to_asset(resized_image, mime_type=mime_type)
        return resized_asset

    def _image_to_asset(self, image, mime_type):
        """
        Converts an PIL image to a MADAM asset. The conversion can also include
        a change in file type.

        :param image: PIL image
        :type image: PIL.Image.Image
        :param mime_type: MIME type of the target asset
        :type mime_type: MimeType
        :return: MADAM asset with hte specified MIME type
        :rtype: Asset
        """
        mime_type = MimeType(mime_type)
        pil_format = PillowProcessor.__mime_type_to_pillow_type[mime_type]
        pil_options = PillowProcessor.__format_options.get(mime_type, {})
        image_buffer = _optimized(image, pil_format, **pil_options)
        asset = self.read(image_buffer)
        return asset

    def _rotate(self, asset, rotation):
        """
        Creates a new image asset from specified asset whose essence is rotated
        by the specified rotation.

        :param asset: Image asset to be rotated
        :type asset: Asset
        :param rotation: One of `PIL.Image.FLIP_LEFT_RIGHT`,
        `PIL.Image.FLIP_TOP_BOTTOM`, `PIL.Image.ROTATE_90`,
        `PIL.Image.ROTATE_180`, `PIL.Image.ROTATE_270`, or
        `PIL.Image.TRANSPOSE`
        :return: New image asset with rotated essence
        :rtype: Asset
        """
        image = PIL.Image.open(asset.essence)
        mime_type = MimeType(asset.mime_type)
        transposed_image = image.transpose(rotation)
        transposed_asset = self._image_to_asset(transposed_image, mime_type=mime_type)
        return transposed_asset

    @operator
    def transpose(self, asset):
        """
        Creates a new image asset whose essence is the transpose of the
        specified asset's essence.

        :param asset: Image asset whose essence is to be transposed
        :type asset: Asset
        :return: New image asset with transposed essence
        :rtype: Asset
        """
        return self._rotate(asset, PIL.Image.TRANSPOSE)

    @operator
    def flip(self, asset, orientation):
        """
        Creates a new asset whose essence is flipped according the specified orientation.

        :param asset: Asset whose essence is to be flipped
        :type asset: Asset
        :param orientation: axis of the flip operation
        :type orientation: FlipOrientation
        :return: Asset with flipped essence
        :rtype: Asset
        """
        if orientation == FlipOrientation.HORIZONTAL:
            flip_orientation = PIL.Image.FLIP_LEFT_RIGHT
        else:
            flip_orientation = PIL.Image.FLIP_TOP_BOTTOM
        return self._rotate(asset, flip_orientation)

    @operator
    def auto_orient(self, asset):
        """
        Creates a new asset whose essence is rotated according to the Exif
        orientation. If no orientation metadata exists or asset is not rotated,
        an identical asset object is returned.

        :param asset: Asset with orientation metadata
        :type asset: Asset
        :return: Asset with rotated essence
        :rtype: Asset
        """
        orientation = asset.metadata.get('exif', {}).get('orientation')
        if orientation is None or orientation == 1:
            return asset

        flip_horizontally = self.flip(orientation=FlipOrientation.HORIZONTAL)
        flip_vertically = self.flip(orientation=FlipOrientation.VERTICAL)

        if orientation == 2:
            oriented_asset = flip_horizontally(asset)
        elif orientation == 3:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_180)
        elif orientation == 4:
            oriented_asset = flip_vertically(asset)
        elif orientation == 5:
            oriented_asset = flip_vertically(self._rotate(asset, PIL.Image.ROTATE_90))
        elif orientation == 6:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_270)
        elif orientation == 7:
            oriented_asset = flip_horizontally(self._rotate(asset, PIL.Image.ROTATE_90))
        elif orientation == 8:
            oriented_asset = self._rotate(asset, PIL.Image.ROTATE_90)
        else:
            raise OperatorError('Unable to correct image orientation with value %s' % orientation)

        return oriented_asset

    @operator
    def convert(self, asset, mime_type, color_space=None, depth=None, data_type=None):
        """
        Creates a new asset of the specified MIME type from the essence of the
        specified asset.

        :param asset: Asset whose contents will be converted
        :type asset: Asset
        :param mime_type: Target MIME type
        :type mime_type: MimeType or str
        :param color_space: Name of color space
        :type color_space: str or None
        :param depth: Bit depth per channel
        :type depth: int or None
        :param data_type: Data type of the pixels, e.g. 'uint' or 'float'
        :type data_type: str or None
        :return: New asset with converted essence
        :rtype: Asset
        """
        mime_type = MimeType(mime_type)
        try:
            image = PIL.Image.open(asset.essence)
            color_mode = color_space or asset.color_space, depth or asset.depth, data_type or asset.data_type
            pil_mode = PillowProcessor.__pillow_mode_to_color_mode.inv.get(color_mode)
            if pil_mode is not None:
                image = image.convert(pil_mode)
            converted_asset = self._image_to_asset(image, mime_type)
        except (IOError, KeyError) as pil_error:
            raise OperatorError('Could not convert image to %s: %s' %
                                (mime_type, pil_error))

        return converted_asset

    @operator
    def crop(self, asset, x, y, width, height):
        """
        Creates a new asset whose essence is cropped to the specified
        rectangular area.

        :param asset: Asset whose contents will be cropped
        :type asset: Asset
        :param x: horizontal offset of the cropping area from left
        :type x: int
        :param y: vertical offset of the cropping area from top
        :type y: int
        :param width: width of the cropping area
        :type width: int
        :param height: height of the cropping area
        :type height: int
        :return: New asset with cropped essence
        :rtype: Asset
        """
        if x == 0 and y == 0 and width == asset.width and height == asset.height:
            return asset

        max_x = max(0, min(asset.width, width + x))
        max_y = max(0, min(asset.height, height + y))
        min_x = max(0, min(asset.width, x))
        min_y = max(0, min(asset.height, y))

        if min_x == asset.width or min_y == asset.height or max_x <= min_x or max_y <= min_y:
            raise OperatorError('Invalid cropping area: <x=%r, y=%r, width=%r, height=%r>' % (x, y, width, height))

        image = PIL.Image.open(asset.essence)
        cropped_image = image.crop(box=(min_x, min_y, max_x, max_y))
        cropped_asset = self._image_to_asset(cropped_image, mime_type=asset.mime_type)

        return cropped_asset

    @operator
    def rotate(self, asset, angle, expand=False):
        """
        Creates an asset whose essence is rotated by the specified angle in
        degrees.

        :param asset: Asset whose contents will be rotated
        :type asset: Asset
        :param angle: Angle in degrees, counter clockwise
        :type angle: float
        :param expand: If true, changes the dimensions of the new asset so it
            can hold the entire rotated essence, otherwise the dimensions of
            the original asset will be used.
        :type expand: bool
        :return: New asset with rotated essence
        :rtype: Asset
        """
        if angle % 360.0 == 0.0:
            return asset

        image = PIL.Image.open(asset.essence)
        rotated_image = image.rotate(angle=angle, resample=PIL.Image.BICUBIC, expand=expand)
        rotated_asset = self._image_to_asset(rotated_image, mime_type=asset.mime_type)

        return rotated_asset
<|endoftext|>


<commit_before>import sys
import os
import logging
import math
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from datetime import datetime, timedelta
from collections import defaultdict, OrderedDict

import pandas as PD

"""
NOTE: This code assumes that the input and output data are 1-minute interval.
"""

HEADER_TEMPLATE = """\
 Format                 IAGA-2002                                    |
 Source of Data         SuperMAG                                     |
 IAGA CODE              {stn}                                          |
 Geodetic Latitude      XXX                                          |
 Geodetic Longitude     XXX                                          |
 Elevation              XXX                                          |
 Reported               {reported}                                         |
DATE       TIME         DOY     {stn}{C1}      {stn}{C2}      {stn}{C3}      {stn}F   |
"""

def read_sm_csv(csv_fname):
    """
    """
    df = PD.read_csv(csv_fname,
                     header=0,
                     parse_dates=[0],
                     date_parser=lambda x: datetime.strptime(x, '%Y-%m-%d %H:%M:%S'))
    return {name: group for name, group in df.groupby('IAGA')}

def get_reported(df):
    """
    """
    if all([x in df.columns for x in ['N', 'E', 'Z']]):
        return 'HDZF', {'C1': ('X', 'N'),
                        'C2': ('Y', 'E'),
                        'C3': ('Z', 'Z')}
    raise ValueError('could not determine coordinate system from columns: ' + ', '.join(df.columns))

def dataframe2iaga(fname, stn, date, df, header_template=HEADER_TEMPLATE):
    """
    ASSUMES DF SPANS ONLY A SINGLE DAY
    """
    reported, coord_map = get_reported(df)
    # build map of data records
    data_map = {}
    for _, row in df.iterrows():
        C1 = row[coord_map['C1'][1]]
        C2 = row[coord_map['C2'][1]]
        C3 = row[coord_map['C3'][1]]
        F = math.hypot(math.hypot(C1, C2), C3)
        data_map[row.Date_UTC] = (C1, C2, C3, F)
    # fill missing records
    filled_data = []
    dts = list(PD.date_range(start=date,
                             end=date + timedelta(days=1),
                             freq='min'))[:-1]
    for dt in dts:
        filled_data.append(data_map.get(dt,
                                        (88888,) * 4))
    # write data records to file
    with open(fname, 'w') as fid:
        fid.write(header_template.format(stn=stn,
                                         reported=reported,
                                         C1=coord_map['C1'][0],
                                         C2=coord_map['C2'][0],
                                         C3=coord_map['C3'][0]))
        for dt, (C1, C2, C3, F) in zip(dts, filled_data):
            fid.write('{date:%Y-%m-%d %H:%M:%S.000} {date:%j}'
                      '    {C1:>9.2f} {C2:>9.2f} {C3:>9.2f} {F:>9.2f}\n'.format(date=dt,
                                                                                C1=C1,
                                                                                C2=C2,
                                                                                C3=C3,
                                                                                F=F))
    return fname

def iaga_fname(path, stn, date, cutoff_date=datetime(2007, 1, 1)):
    """
    """
    if date >= cutoff_date.date():
        ext = 'dmin.min'
    else:
        ext = 'd.min'
    return os.path.join(path, '{stn}{date:%Y%m%d}{ext}'.format(stn=stn,
                                                               date=date,
                                                               ext=ext))

def df_map2iaga(path, df_map):
    """
    """
    output_map = defaultdict(dict)
    for stn, df in df_map.iteritems():
        for date, df_date in df.groupby(df['Date_UTC'].map(lambda x: x.date())):
            output_fname = iaga_fname(path, stn, date)
            dataframe2iaga(output_fname, stn, date, df_date)
            output_map[stn][date] = output_fname
    return output_map

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

    parser = ArgumentParser('Convert SuperMAG CSV format data to daily IAGA2002 format.',
                            formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('output_path',
                        type=str,
                        help='path to store daily IAGA2002 format files')
    parser.add_argument('csv_fname',
                        type=str,
                        help='SuperMAG CSV file')
    args = parser.parse_args(argv[1:])

    df_map = read_sm_csv(args.csv_fname)
    df_map2iaga(args.output_path,
                df_map)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    sys.exit(main())
<commit_msg>Fixed header labels.<commit_after>import sys
import os
import logging
import math
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from datetime import datetime, timedelta
from collections import defaultdict, OrderedDict

import pandas as PD

"""
NOTE: This code assumes that the input and output data are 1-minute interval.
"""

HEADER_TEMPLATE = """\
 Format                 IAGA-2002                                    |
 Source of Data         SuperMAG                                     |
 IAGA CODE              {stn}                                          |
 Geodetic Latitude      XXX                                          |
 Geodetic Longitude     XXX                                          |
 Elevation              XXX                                          |
 Reported               {reported}                                         |
DATE       TIME         DOY     {stn}{C1}      {stn}{C2}      {stn}{C3}      {stn}F   |
"""

def read_sm_csv(csv_fname):
    """
    """
    df = PD.read_csv(csv_fname,
                     header=0,
                     parse_dates=[0],
                     date_parser=lambda x: datetime.strptime(x, '%Y-%m-%d %H:%M:%S'))
    return {name: group for name, group in df.groupby('IAGA')}

def get_reported(df):
    """
    """
    if all([x in df.columns for x in ['N', 'E', 'Z']]):
        return 'HDZF', {'C1': ('H', 'N'),
                        'C2': ('D', 'E'),
                        'C3': ('Z', 'Z')}
    raise ValueError('could not determine coordinate system from columns: ' + ', '.join(df.columns))

def dataframe2iaga(fname, stn, date, df, header_template=HEADER_TEMPLATE):
    """
    ASSUMES DF SPANS ONLY A SINGLE DAY
    """
    reported, coord_map = get_reported(df)
    # build map of data records
    data_map = {}
    for _, row in df.iterrows():
        C1 = row[coord_map['C1'][1]]
        C2 = row[coord_map['C2'][1]]
        C3 = row[coord_map['C3'][1]]
        F = math.hypot(math.hypot(C1, C2), C3)
        data_map[row.Date_UTC] = (C1, C2, C3, F)
    # fill missing records
    filled_data = []
    dts = list(PD.date_range(start=date,
                             end=date + timedelta(days=1),
                             freq='min'))[:-1]
    for dt in dts:
        filled_data.append(data_map.get(dt,
                                        (88888,) * 4))
    # write data records to file
    with open(fname, 'w') as fid:
        fid.write(header_template.format(stn=stn,
                                         reported=reported,
                                         C1=coord_map['C1'][0],
                                         C2=coord_map['C2'][0],
                                         C3=coord_map['C3'][0]))
        for dt, (C1, C2, C3, F) in zip(dts, filled_data):
            fid.write('{date:%Y-%m-%d %H:%M:%S.000} {date:%j}'
                      '    {C1:>9.2f} {C2:>9.2f} {C3:>9.2f} {F:>9.2f}\n'.format(date=dt,
                                                                                C1=C1,
                                                                                C2=C2,
                                                                                C3=C3,
                                                                                F=F))
    return fname

def iaga_fname(path, stn, date, cutoff_date=datetime(2007, 1, 1)):
    """
    """
    if date >= cutoff_date.date():
        ext = 'dmin.min'
    else:
        ext = 'd.min'
    return os.path.join(path, '{stn}{date:%Y%m%d}{ext}'.format(stn=stn,
                                                               date=date,
                                                               ext=ext))

def df_map2iaga(path, df_map):
    """
    """
    output_map = defaultdict(dict)
    for stn, df in df_map.iteritems():
        for date, df_date in df.groupby(df['Date_UTC'].map(lambda x: x.date())):
            output_fname = iaga_fname(path, stn, date)
            dataframe2iaga(output_fname, stn, date, df_date)
            output_map[stn][date] = output_fname
    return output_map

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

    parser = ArgumentParser('Convert SuperMAG CSV format data to daily IAGA2002 format.',
                            formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('output_path',
                        type=str,
                        help='path to store daily IAGA2002 format files')
    parser.add_argument('csv_fname',
                        type=str,
                        help='SuperMAG CSV file')
    args = parser.parse_args(argv[1:])

    df_map = read_sm_csv(args.csv_fname)
    df_map2iaga(args.output_path,
                df_map)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    sys.exit(main())
<|endoftext|>


<commit_before>"""
 Text-Machine Lab: CliRel

 File Name :

 Creation Date : 12-10-2016

 Created By : Renan Campos

 Purpose : Module to produce tree enrichments.
           Here a tree is represented as a string consisting
           of tokens and paranthesis.

"""

import re
tree_ex = re.compile(r"([^\(\) ]+)")

def createTree(s):
  """
    Takes a string and creates a list representation
    >>> createTree(test1)
    '(S (NP (DT This)) (VP (VBZ is) (NP (DT a) (NN test) (NN sentence))) (. .))'
  """
  t = ','.join(re.sub(tree_ex, r'"\1"', s).split())
  # Make lists not tuples
  return eval(t.replace('(', '[').replace(')',']'))

def getLeaves(tree):
  """
    Returns a string of the leaves of this particular branch.
    >>> getLeaves(test1)
    'This is a test sentence'
    >>> getLeaves(test2)
    'Sentences can also have ()'
    >>> getLeaves(test3)
    'medical problem reveals medical problem'
  """
  return None

def spt(tree, a, b):
  """
    Finds the shortest path tree between tokens a and b.
  """
  return None

if __name__ == '__main__':
  test1 = '(S (NP (DT This)) (VP (VBZ is) (NP (DT a) (NN test) (NN sentence))) (. .))'
  import doctest
  doctest.testmod()
<commit_msg>Added more tree tests/functions.<commit_after>"""
 Text-Machine Lab: CliRel

 File Name : tree.py

 Creation Date : 12-10-2016

 Created By : Renan Campos

 Purpose : Module to produce tree enrichments.
           Here a tree is represented as a string consisting
           of tokens and paranthesis.

"""

import re
tree_ex = re.compile(r"([^\(\) ]+)")

def createTree(s):
  """
    Takes a string and creates a list representation
    >>> createTree('()')
    []
    >>> createTree(test0)
    ['S', ['NP', ['DT', 'This']], ['VP', ['VBZ', 'is'], ['NP', ['DT', 'a'], ['NN', 'test'], ['NN', 'sentence']]], ['.', '.']]
  """
  if s == '':
    return []
  t = ','.join(re.sub(tree_ex, r'"\1"', s).split())
  # Make lists not tuples
  return eval(t.replace('(', '[').replace(')',']'))

def getLeaves(tree):
  """
    Returns a string of the leaves of this particular branch.
    >>> getLeaves([])
    ''
    >>> getLeaves(createTree(test0))
    'This is a test sentence .'
    >>> getLeaves(createTree(test6))
    'There were many symptoms , but test reveals medical problem .'
  """
  out = ''

  if len(tree) < 2:
    return out

  for child in tree[1:]:
    if type(child) == str:
      out += child
    else:
      out += getLeaves(child)
    out += ' '

  return out[:-1]

def createString(tree):
  """
    Returns the tree as a string representation again.
    >>> createString(createTree('()'))
    '()'
    >>> createString(createTree(test0))
    '(S (NP (DT This)) (VP (VBZ is) (NP (DT a) (NN test) (NN sentence))) (. .))'
    >>> createString(createTree(test1))
    '(S (NP (DT This) (NN treatment)) (VP (VBZ improves) (NP (JJ medical) (NN problem))) (. .))'
    >>> createString(createTree(createString(createTree(test1))))
    '(S (NP (DT This) (NN treatment)) (VP (VBZ improves) (NP (JJ medical) (NN problem))) (. .))'
  """
  out = '('

  if len(tree) < 2:
    return out + ')'

  out += tree[0]
  for child in tree[1:]:
    if type(child) == str:
      out += ' ' + child
    else:
      out += ' ' + createString(child)

  return out + ')'

def spt(tree, a, b):
  """
    Finds the shortest path tree between tokens a and b.
    >>> spt(createTree(test0), 'is', 'sentence')
    ['VP', ['VBZ' 'is'], ['NP', ['DT', 'a'], ['NN', 'test'], ['NN', 'sentence']]]
  """

  return None

if __name__ == '__main__':
  test0 = '(S (NP (DT This)) (VP (VBZ is) (NP (DT a) (NN test) (NN sentence))) (. .))'
  test1 = '(S (NP (DT This) (NN treatment)) (VP (VBZ improves) (NP (JJ medical) (NN problem))) (. .))'
  test2 = '(S (NP (NNP Treatment)) (VP (VBZ worsens) (NP (JJ medical) (NN problem))) (. .))'
  test3 = '(S (NP (NNP Treatment)) (VP (VBZ causes) (NP (JJ medical) (NN problem))) (. .))'
  test4 = '(S (NP (DT The) (NN treatment)) (VP (VBZ is) (VP (VBN administered) (PP (IN for) (NP (JJ medical) (NN problem))))) (. .))'
  test5 = '(S (NP (NNP Treatment)) (VP (VBZ is) (RB not) (VP (VBN administered) (PP (IN because) (IN of) (NP (JJ medical) (NN problem))))) (. .))'
  test6 = '(S (S (NP (EX There)) (VP (VBD were) (NP (JJ many) (NNS symptoms)))) (, ,) (CC but) (S (NP (NN test)) (VP (VBZ reveals) (NP (JJ medical) (NN problem)))) (. .))'
  test7 = '(S (NP (DT The) (JJ great) (NNS tests)) (VP (VBD conducted) (S (VP (TO to) (VP (VB investigate) (NP (JJ medical) (NN problem)))))) (. .))'
  test8 = '(S (NP (JJ Medical) (NN problem)) (VP (VBD indicated) (NP (NP (JJ medical) (NN problem)) (: ;) (SBAR (WHNP (WDT which)) (S (VP (VBZ is) (ADJP (RB really) (JJ important))))))) (. .))'
  test9 = '(NP (NP (NP (DT The) (JJ great) (NNS tests)) (VP (VBD conducted) (S (VP (TO to) (VP (VB investigate) (NP (JJ medical) (NN problem)))))) (. .)) (NNP VOID))'
  test10 = '(FRAG (NP (JJ Medical) (NN problem) (VBN indicated) (JJ medical) (NN problem)) (: ;) (SBAR (WHNP (WDT which)) (S (VP (VBZ is) (ADJP (RB really) (JJ important))))) (. .) (NP (NNP VOID)))'
  test11 = '(S (NP (DT This) (NN treatment)) (VP (VBZ improves) (NP (JJ medical) (NN problem) (. .) (NNP VOID))))'
  import doctest
  doctest.testmod()
<|endoftext|>


<commit_before>import autocomplete_light
import rules
from django.contrib.auth import get_user_model
from django.contrib import admin
from django.contrib.admin.widgets import AdminRadioSelect
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.forms import (
    HiddenInput,
    NullBooleanSelect,
    ModelForm,
    ValidationError,
    ModelChoiceField,
)
from django.utils.translation import ugettext_lazy as _
from django_extensions.admin import ForeignKeyAutocompleteAdmin
from powerdns.models.powerdns import (
    CryptoKey,
    Domain,
    DomainMetadata,
    Record,
    SuperMaster,
)
from powerdns.models.authorisations import Authorisation
from rules.contrib.admin import ObjectPermissionsModelAdmin
from threadlocals.threadlocals import get_current_user

from powerdns.models.powerdns import can_delete, can_edit
from powerdns.models.templates import (
    DomainTemplate,
    RecordTemplate,
)
from powerdns.models.requests import (
    DeleteRequest,
    DomainRequest,
    RecordRequest,
)
from powerdns.utils import Owned, DomainForRecordValidator, is_owner
from powerdns.models.tsigkeys import TsigKey

class NullBooleanRadioSelect(NullBooleanSelect, AdminRadioSelect):
    pass

try:
    from django.contrib.admin import SimpleListFilter
except ImportError:
    _domain_filters = ('type', 'last_check', 'account',)
else:
    class ReverseDomainListFilter(SimpleListFilter):
        title = _('domain class')

        # Parameter for the filter that will be used in the URL query.
        parameter_name = 'domain_class'

        def lookups(self, request, model_admin):
            return (
                ('fwd', _('domain:forward')),
                ('rev', _('domain:reverse')),
            )

        def queryset(self, request, queryset):
            q = (models.Q(name__endswith='.in-addr.arpa') |
                 models.Q(name__endswith='.ip6.arpa'))
            if self.value() == 'fwd':
                return queryset.exclude(q)
            if self.value() == 'rev':
                return queryset.filter(q)
    _domain_filters = (
        ReverseDomainListFilter, 'type', 'last_check', 'account',
    )

class RecordAdminForm(ModelForm):

    def clean_type(self):
        type = self.cleaned_data['type']
        if not type:
            raise ValidationError(_("Record type is required"))
        return type

    def clean_domain(self):
        if (
            self.instance.pk and
            self.instance.domain == self.cleaned_data['domain']
        ):
            # Domain unchanged. Maybe user was assigned the record in a domain
            # She doesn't own.
            return self.cleaned_data['domain']
        validator = DomainForRecordValidator()
        validator.user = self.user
        return validator(self.cleaned_data['domain'])

class CopyingAdmin(admin.ModelAdmin):

    field_prefix = ''
    target_prefix = ''

    def get_form(self, request, obj=None, **kwargs):
        form = super().get_form(request, obj, **kwargs)
        from_pk = request.GET.get(self.from_field)
        if from_pk is not None:
            self.from_object = self.FromModel.objects.get(pk=from_pk)
            for field in self.CopyFieldsModel.copy_fields:
                form.base_fields[field[len(self.field_prefix):]].initial = \
                    getattr(self.from_object, field[len(self.target_prefix):])
        else:
            self.from_object = None
        return form

class RequestAdmin(CopyingAdmin):
    """Admin for domain/record requests"""

    def get_form(self, *args, **kwargs):
        form = super().get_form(*args, **kwargs)
        form.base_fields['target_owner'].initial =\
            form.base_fields['target_owner'].initial or get_current_user()
        return form

class OwnedAdmin(ForeignKeyAutocompleteAdmin, ObjectPermissionsModelAdmin):
    """Admin for models with owner field"""

    def save_model(self, request, object_, form, change):
        if object_.owner is None:
            object_.owner = request.user
        object_.email_owner(request.user)
        super(OwnedAdmin, self).save_model(request, object_, form, change)

    def get_related_filter(self, model, request):
        return super(OwnedAdmin, self).get_related_filter(model, request)
        user = request.user
        if not issubclass(model, Owned) or rules.is_superuser(user):
            return super(OwnedAdmin, self).get_related_filter(model, request)
        return models.Q(owner=user)

class RecordAdmin(OwnedAdmin, CopyingAdmin):
    def has_delete_permission(self, request, obj=None):
        return can_delete(request.user, obj)

    def has_change_permission(self, request, obj=None):
        return can_edit(request.user, obj)

    form = RecordAdminForm
    list_display = (
        'name',
        'type',
        'content',
        'domain',
        'owner',
        'ttl',
        'prio',
        'change_date',
        'request_change',
        'request_deletion',
    )
    list_display_links = None
    list_filter = ('type', 'ttl', 'auth', 'domain', 'created', 'modified')
    list_per_page = 250
    save_on_top = True
    search_fields = ('name', 'content',)
    readonly_fields = ('change_date', 'ordername', 'created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    fieldsets = (
        (None, {
            'fields': (
                'owner',
                'domain',
                ('type', 'name', 'content',),
                'auth',
            ),
        }),
        ('Advanced options', {
            'classes': ('collapse',),
            'fields': ('prio', 'ttl', 'ordername', 'change_date',)
        }),
        (None, {'fields': ('created', 'modified')})
    )
    formfield_overrides = {
        models.NullBooleanField: {
            'widget': NullBooleanRadioSelect(
                attrs={'class': 'radiolist inline'}
            ),
        },
    }
    FromModel = Domain
    CopyFieldsModel = Domain
    from_field = 'domain'
    field_prefix = 'record_'

    def get_form(self, request, *args, **kwargs):
        form = super().get_form(request, *args, **kwargs)
        form.user = request.user
        return form

class DomainMetadataInline(admin.TabularInline):
    model = DomainMetadata
    extra = 0

class DomainAdmin(OwnedAdmin, CopyingAdmin):
    inlines = [DomainMetadataInline]
    list_display = (
        'name',
        'type',
        'last_check',
        'account',
        'add_record_link',
        'request_change',
        'request_deletion'
    )
    list_display_links = None
    list_filter = _domain_filters + ('created', 'modified')
    list_per_page = 250
    save_on_top = True
    search_fields = ('name',)
    radio_fields = {'type': admin.HORIZONTAL}
    readonly_fields = ('notified_serial', 'created', 'modified')
    FromModel = DomainTemplate
    CopyFieldsModel = DomainTemplate
    from_field = 'template'

class SuperMasterAdmin(admin.ModelAdmin):
    list_display = ('ip', 'nameserver', 'account',)
    list_filter = ('ip', 'account', 'created', 'modified')
    search_fields = ('ip', 'nameserver',)
    readonly_fields = ('created', 'modified')

class DomainMetadataAdmin(ForeignKeyAutocompleteAdmin):
    list_display = ('domain', 'kind', 'content',)
    list_filter = ('kind', 'domain', 'created', 'modified')
    list_per_page = 250
    list_filter = ('created', 'modified')
    readonly_fields = ('created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    save_on_top = True
    search_fields = ('content',)

class CryptoKeyAdmin(ForeignKeyAutocompleteAdmin):
    list_display = ('domain', 'flags', 'active', 'content',)
    list_filter = ('active', 'domain', 'created', 'modified')
    list_per_page = 250
    readonly_fields = ('created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    save_on_top = True
    search_fields = ('content',)
    formfield_overrides = {
        models.NullBooleanField: {
            'widget': NullBooleanRadioSelect(
                attrs={'class': 'radiolist inline'}
            ),
        },
    }

class RecordTemplateInline(admin.StackedInline):
    model = RecordTemplate
    extra = 1

class DomainTemplateAdmin(ForeignKeyAutocompleteAdmin):
    inlines = [RecordTemplateInline]
    list_display = ['name', 'add_domain_link', 'is_public_domain']

RECORD_LIST_FIELDS = (
    'name',
    'type',
    'content',
    'ttl',
    'prio',
)

class RecordTemplateAdmin(ForeignKeyAutocompleteAdmin):
    form = RecordAdminForm
    list_display = RECORD_LIST_FIELDS

class DomainRequestForm(autocomplete_light.forms.ModelForm):
    class Meta:
        exclude = ['owner', 'state']

class DomainRequestAdmin(RequestAdmin):
    form = DomainRequestForm
    list_display = ['domain']
    from_field = 'domain'
    FromModel = Domain
    CopyFieldsModel = DomainRequest
    target_prefix = 'target_'
    readonly_fields = ['key']

class AuthorisationForm(autocomplete_light.forms.ModelForm):
    class Meta:
        model = Authorisation
        fields = [
            'owner',
            'authorised',
            'target',
        ]
        autocomplete_fields = ['authorised', 'target']

    owner = ModelChoiceField(
        queryset=get_user_model().objects.all(),
        widget=HiddenInput()
    )

    def clean_target(self):
        target = self.cleaned_data['target']
        user = get_current_user()
        if not (user.is_superuser or is_owner(user, target)):
            raise ValidationError(_("You cannot authorize for this"))
        return target

class AuthorisationAdmin(ObjectPermissionsModelAdmin):
    form = AuthorisationForm

    def get_changeform_initial_data(self, request, *args, **kwargs):
        data = super().get_changeform_initial_data(request, *args, **kwargs)
        data['owner'] = request.user
        return data

class DeleteRequestForm(ModelForm):
    class Meta:
        model = DeleteRequest
        fields = [
            'owner',
            'target_id',
            'content_type',
            'key',
        ]
        widgets = {
            'owner': HiddenInput(),
            'key': HiddenInput(),
            'target_id': HiddenInput(),
            'content_type': HiddenInput(),
        }

class DeleteRequestAdmin(ObjectPermissionsModelAdmin):
    form = DeleteRequestForm
    fields = ['owner', 'target_id', 'content_type']

    def add_view(self, request, extra_context=None):
        extra_context = extra_context or {}
        extra_context['deleted_item'] = str(
            ContentType.objects.get(
                pk=request.GET['content_type']
            ).get_object_for_this_type(pk=request.GET['target_id'])
        )
        return super().add_view(request, extra_context=extra_context)

class RecordRequestForm(autocomplete_light.forms.ModelForm):
    class Meta:
        exclude = ['owner', 'state']

class RecordRequestAdmin(RequestAdmin):
    form = RecordRequestForm
    list_display = ['target_' + field for field in RECORD_LIST_FIELDS]
    from_field = 'record'
    FromModel = Record
    CopyFieldsModel = RecordRequest
    target_prefix = 'target_'

    def get_form(self, request, obj=None, **kwargs):
        form = super().get_form(request, obj, **kwargs)
        if self.from_object:
            form.base_fields['domain'].initial = self.from_object.domain
        return form

admin.site.register(Domain, DomainAdmin)
admin.site.register(Record, RecordAdmin)
admin.site.register(SuperMaster, SuperMasterAdmin)
admin.site.register(DomainMetadata, DomainMetadataAdmin)
admin.site.register(CryptoKey, CryptoKeyAdmin)
admin.site.register(DomainTemplate, DomainTemplateAdmin)
admin.site.register(RecordTemplate, RecordTemplateAdmin)
admin.site.register(Authorisation, AuthorisationAdmin)
admin.site.register(DomainRequest, DomainRequestAdmin)
admin.site.register(RecordRequest, RecordRequestAdmin)
admin.site.register(DeleteRequest, DeleteRequestAdmin)
admin.site.register(TsigKey)
<commit_msg>Add search fields to RecordRequest admin (#147)<commit_after>import autocomplete_light
import rules
from django.contrib.auth import get_user_model
from django.contrib import admin
from django.contrib.admin.widgets import AdminRadioSelect
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.forms import (
    HiddenInput,
    NullBooleanSelect,
    ModelForm,
    ValidationError,
    ModelChoiceField,
)
from django.utils.translation import ugettext_lazy as _
from django_extensions.admin import ForeignKeyAutocompleteAdmin
from powerdns.models.powerdns import (
    CryptoKey,
    Domain,
    DomainMetadata,
    Record,
    SuperMaster,
)
from powerdns.models.authorisations import Authorisation
from rules.contrib.admin import ObjectPermissionsModelAdmin
from threadlocals.threadlocals import get_current_user

from powerdns.models.powerdns import can_delete, can_edit
from powerdns.models.templates import (
    DomainTemplate,
    RecordTemplate,
)
from powerdns.models.requests import (
    DeleteRequest,
    DomainRequest,
    RecordRequest,
)
from powerdns.utils import Owned, DomainForRecordValidator, is_owner
from powerdns.models.tsigkeys import TsigKey

class NullBooleanRadioSelect(NullBooleanSelect, AdminRadioSelect):
    pass

try:
    from django.contrib.admin import SimpleListFilter
except ImportError:
    _domain_filters = ('type', 'last_check', 'account',)
else:
    class ReverseDomainListFilter(SimpleListFilter):
        title = _('domain class')

        # Parameter for the filter that will be used in the URL query.
        parameter_name = 'domain_class'

        def lookups(self, request, model_admin):
            return (
                ('fwd', _('domain:forward')),
                ('rev', _('domain:reverse')),
            )

        def queryset(self, request, queryset):
            q = (models.Q(name__endswith='.in-addr.arpa') |
                 models.Q(name__endswith='.ip6.arpa'))
            if self.value() == 'fwd':
                return queryset.exclude(q)
            if self.value() == 'rev':
                return queryset.filter(q)
    _domain_filters = (
        ReverseDomainListFilter, 'type', 'last_check', 'account',
    )

class RecordAdminForm(ModelForm):

    def clean_type(self):
        type = self.cleaned_data['type']
        if not type:
            raise ValidationError(_("Record type is required"))
        return type

    def clean_domain(self):
        if (
            self.instance.pk and
            self.instance.domain == self.cleaned_data['domain']
        ):
            # Domain unchanged. Maybe user was assigned the record in a domain
            # She doesn't own.
            return self.cleaned_data['domain']
        validator = DomainForRecordValidator()
        validator.user = self.user
        return validator(self.cleaned_data['domain'])

class CopyingAdmin(admin.ModelAdmin):

    field_prefix = ''
    target_prefix = ''

    def get_form(self, request, obj=None, **kwargs):
        form = super().get_form(request, obj, **kwargs)
        from_pk = request.GET.get(self.from_field)
        if from_pk is not None:
            self.from_object = self.FromModel.objects.get(pk=from_pk)
            for field in self.CopyFieldsModel.copy_fields:
                form.base_fields[field[len(self.field_prefix):]].initial = \
                    getattr(self.from_object, field[len(self.target_prefix):])
        else:
            self.from_object = None
        return form

class RequestAdmin(CopyingAdmin):
    """Admin for domain/record requests"""

    search_fields = ['target_name', 'key', 'target_content']

    def get_form(self, *args, **kwargs):
        form = super().get_form(*args, **kwargs)
        form.base_fields['target_owner'].initial =\
            form.base_fields['target_owner'].initial or get_current_user()
        return form

class OwnedAdmin(ForeignKeyAutocompleteAdmin, ObjectPermissionsModelAdmin):
    """Admin for models with owner field"""

    def save_model(self, request, object_, form, change):
        if object_.owner is None:
            object_.owner = request.user
        object_.email_owner(request.user)
        super(OwnedAdmin, self).save_model(request, object_, form, change)

    def get_related_filter(self, model, request):
        return super(OwnedAdmin, self).get_related_filter(model, request)
        user = request.user
        if not issubclass(model, Owned) or rules.is_superuser(user):
            return super(OwnedAdmin, self).get_related_filter(model, request)
        return models.Q(owner=user)

class RecordAdmin(OwnedAdmin, CopyingAdmin):
    def has_delete_permission(self, request, obj=None):
        return can_delete(request.user, obj)

    def has_change_permission(self, request, obj=None):
        return can_edit(request.user, obj)

    form = RecordAdminForm
    list_display = (
        'name',
        'type',
        'content',
        'domain',
        'owner',
        'ttl',
        'prio',
        'change_date',
        'request_change',
        'request_deletion',
    )
    list_display_links = None
    list_filter = ('type', 'ttl', 'auth', 'domain', 'created', 'modified')
    list_per_page = 250
    save_on_top = True
    search_fields = ('name', 'content',)
    readonly_fields = ('change_date', 'ordername', 'created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    fieldsets = (
        (None, {
            'fields': (
                'owner',
                'domain',
                ('type', 'name', 'content',),
                'auth',
            ),
        }),
        ('Advanced options', {
            'classes': ('collapse',),
            'fields': ('prio', 'ttl', 'ordername', 'change_date',)
        }),
        (None, {'fields': ('created', 'modified')})
    )
    formfield_overrides = {
        models.NullBooleanField: {
            'widget': NullBooleanRadioSelect(
                attrs={'class': 'radiolist inline'}
            ),
        },
    }
    FromModel = Domain
    CopyFieldsModel = Domain
    from_field = 'domain'
    field_prefix = 'record_'

    def get_form(self, request, *args, **kwargs):
        form = super().get_form(request, *args, **kwargs)
        form.user = request.user
        return form

class DomainMetadataInline(admin.TabularInline):
    model = DomainMetadata
    extra = 0

class DomainAdmin(OwnedAdmin, CopyingAdmin):
    inlines = [DomainMetadataInline]
    list_display = (
        'name',
        'type',
        'last_check',
        'account',
        'add_record_link',
        'request_change',
        'request_deletion'
    )
    list_display_links = None
    list_filter = _domain_filters + ('created', 'modified')
    list_per_page = 250
    save_on_top = True
    search_fields = ('name',)
    radio_fields = {'type': admin.HORIZONTAL}
    readonly_fields = ('notified_serial', 'created', 'modified')
    FromModel = DomainTemplate
    CopyFieldsModel = DomainTemplate
    from_field = 'template'

class SuperMasterAdmin(admin.ModelAdmin):
    list_display = ('ip', 'nameserver', 'account',)
    list_filter = ('ip', 'account', 'created', 'modified')
    search_fields = ('ip', 'nameserver',)
    readonly_fields = ('created', 'modified')

class DomainMetadataAdmin(ForeignKeyAutocompleteAdmin):
    list_display = ('domain', 'kind', 'content',)
    list_filter = ('kind', 'domain', 'created', 'modified')
    list_per_page = 250
    list_filter = ('created', 'modified')
    readonly_fields = ('created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    save_on_top = True
    search_fields = ('content',)

class CryptoKeyAdmin(ForeignKeyAutocompleteAdmin):
    list_display = ('domain', 'flags', 'active', 'content',)
    list_filter = ('active', 'domain', 'created', 'modified')
    list_per_page = 250
    readonly_fields = ('created', 'modified')
    related_search_fields = {
        'domain': ('name',),
    }
    save_on_top = True
    search_fields = ('content',)
    formfield_overrides = {
        models.NullBooleanField: {
            'widget': NullBooleanRadioSelect(
                attrs={'class': 'radiolist inline'}
            ),
        },
    }

class RecordTemplateInline(admin.StackedInline):
    model = RecordTemplate
    extra = 1

class DomainTemplateAdmin(ForeignKeyAutocompleteAdmin):
    inlines = [RecordTemplateInline]
    list_display = ['name', 'add_domain_link', 'is_public_domain']

RECORD_LIST_FIELDS = (
    'name',
    'type',
    'content',
    'ttl',
    'prio',
)

class RecordTemplateAdmin(ForeignKeyAutocompleteAdmin):
    form = RecordAdminForm
    list_display = RECORD_LIST_FIELDS

class DomainRequestForm(autocomplete_light.forms.ModelForm):
    class Meta:
        exclude = ['owner', 'state']

class DomainRequestAdmin(RequestAdmin):
    form = DomainRequestForm
    list_display = ['domain']
    from_field = 'domain'
    FromModel = Domain
    CopyFieldsModel = DomainRequest
    target_prefix = 'target_'
    readonly_fields = ['key']

class AuthorisationForm(autocomplete_light.forms.ModelForm):
    class Meta:
        model = Authorisation
        fields = [
            'owner',
            'authorised',
            'target',
        ]
        autocomplete_fields = ['authorised', 'target']

    owner = ModelChoiceField(
        queryset=get_user_model().objects.all(),
        widget=HiddenInput()
    )

    def clean_target(self):
        target = self.cleaned_data['target']
        user = get_current_user()
        if not (user.is_superuser or is_owner(user, target)):
            raise ValidationError(_("You cannot authorize for this"))
        return target

class AuthorisationAdmin(ObjectPermissionsModelAdmin):
    form = AuthorisationForm

    def get_changeform_initial_data(self, request, *args, **kwargs):
        data = super().get_changeform_initial_data(request, *args, **kwargs)
        data['owner'] = request.user
        return data

class DeleteRequestForm(ModelForm):
    class Meta:
        model = DeleteRequest
        fields = [
            'owner',
            'target_id',
            'content_type',
            'key',
        ]
        widgets = {
            'owner': HiddenInput(),
            'key': HiddenInput(),
            'target_id': HiddenInput(),
            'content_type': HiddenInput(),
        }

class DeleteRequestAdmin(ObjectPermissionsModelAdmin):
    form = DeleteRequestForm
    fields = ['owner', 'target_id', 'content_type']

    def add_view(self, request, extra_context=None):
        extra_context = extra_context or {}
        extra_context['deleted_item'] = str(
            ContentType.objects.get(
                pk=request.GET['content_type']
            ).get_object_for_this_type(pk=request.GET['target_id'])
        )
        return super().add_view(request, extra_context=extra_context)

class RecordRequestForm(autocomplete_light.forms.ModelForm):
    class Meta:
        exclude = ['owner', 'state']

class RecordRequestAdmin(RequestAdmin):
    form = RecordRequestForm
    list_display = ['target_' + field for field in RECORD_LIST_FIELDS]
    from_field = 'record'
    FromModel = Record
    CopyFieldsModel = RecordRequest
    target_prefix = 'target_'

    def get_form(self, request, obj=None, **kwargs):
        form = super().get_form(request, obj, **kwargs)
        if self.from_object:
            form.base_fields['domain'].initial = self.from_object.domain
        return form

admin.site.register(Domain, DomainAdmin)
admin.site.register(Record, RecordAdmin)
admin.site.register(SuperMaster, SuperMasterAdmin)
admin.site.register(DomainMetadata, DomainMetadataAdmin)
admin.site.register(CryptoKey, CryptoKeyAdmin)
admin.site.register(DomainTemplate, DomainTemplateAdmin)
admin.site.register(RecordTemplate, RecordTemplateAdmin)
admin.site.register(Authorisation, AuthorisationAdmin)
admin.site.register(DomainRequest, DomainRequestAdmin)
admin.site.register(RecordRequest, RecordRequestAdmin)
admin.site.register(DeleteRequest, DeleteRequestAdmin)
admin.site.register(TsigKey)
<|endoftext|>


<commit_before>import sys
if sys.version_info[0] == 3:
    from urllib.request import urlopen
    from urllib.error import HTTPError
    from urllib.parse import urlencode
else:
    from urllib2 import urlopen, HTTPError
    from urllib import urlencode
import socket
import struct
from ctypes.util import find_library
from ctypes import Structure, pointer, c_long, CDLL
from poyonga.result import GroongaResult, GroongaSelectResult

GQTP_HEADER_SIZE = 24

class Groonga(object):

    LIBRT = CDLL(find_library("rt"))

    class CTimeSpec(Structure):
        _fields_ = [("tv_sec", c_long), ("tv_nsec", c_long)]

    def __init__(self, host='localhost', port=10041, protocol='http',
                 encoding='utf-8'):
        self.host = host
        self.port = port
        self.protocol = protocol
        self.encoding = encoding

    def _call_gqtp(self, cmd, **kwargs):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((self.host, self.port))
        _cmd = cmd
        _start = self.CTimeSpec()
        _end = self.CTimeSpec()
        _cmd_arg = "".join([" --%s '%s'" % (d, str(kwargs[d]).replace("'", r"\'")) for d in kwargs])
        _cmd = _cmd + _cmd_arg
        _cmd_str = "%08x" % len(_cmd)
        exec("_cmd_len = \"\\x%02s\\x%02s\\x%02s\\x%02s\"" % (
                _cmd_str[:2], _cmd_str[2:4], _cmd_str[4:6], _cmd_str[6:]))
        _header = "".join(["\xc7", "\x00" * 7, _cmd_len, "\x00" * 12])
        self.LIBRT.clock_gettime(0, pointer(_start))     # 0: CLOCK_REALTIME
        s.send(_header + _cmd)
        raw_data = s.recv(8192)
        proto, qtype, keylen, level, flags, status, size, opaque, cas \
                = struct.unpack("!BBHBBHIIQ", raw_data[:GQTP_HEADER_SIZE])
        while len(raw_data) < size + GQTP_HEADER_SIZE:
            raw_data += s.recv(8192)
        self.LIBRT.clock_gettime(0, pointer(_end))
        s.close()

        diff_time = (_end.tv_sec + _end.tv_nsec / 1000000000.) - \
                    (_start.tv_sec + _start.tv_nsec / 1000000000.)
        if status != 0:
            status -= 65536
            body = "\"\",[[\"\",\"\",0]]"
            _data = "[[%d,%d.%d,%lf,%s]]" % (
                    status, _start.tv_sec, _start.tv_nsec, diff_time, body)
        else:
            body = raw_data[24:]
            _data = "[[%d,%d.%d,%lf],%s]" % (
                    status, _start.tv_sec, _start.tv_nsec, diff_time, body)
        return _data

    def _call_http(self, cmd, **kwargs):
        domain = [self.protocol, "://", self.host, ":", str(self.port), "/d/"]
        if kwargs:
            url = "".join(domain) + cmd + "?" + urlencode(kwargs)
        else:
            url = "".join(domain) + cmd
        try:
            _data = urlopen(url).read()
        except HTTPError as msg:
            _data = msg.read()
        return _data

    def call(self, cmd, **kwargs):
        output_type = kwargs.get("output_type")
        if not output_type:
            output_type = "json"
        if self.protocol == "http":
            ret = self._call_http(cmd, **kwargs)
        else:
            ret = self._call_gqtp(cmd, **kwargs)
        if cmd == 'select':
            return GroongaSelectResult(ret, output_type, self.encoding)
        else:
            return GroongaResult(ret, output_type, self.encoding)
<commit_msg>support GQTP protocol for MacOSX environment<commit_after>import sys
if sys.version_info[0] == 3:
    from urllib.request import urlopen
    from urllib.error import HTTPError
    from urllib.parse import urlencode
else:
    from urllib2 import urlopen, HTTPError
    from urllib import urlencode
import socket
import struct
from ctypes.util import find_library
from ctypes import Structure, pointer, c_long, CDLL
from poyonga.result import GroongaResult, GroongaSelectResult

GQTP_HEADER_SIZE = 24

class Groonga(object):

    LIBRT = CDLL(find_library("rt"))
    LIBC = CDLL(find_library("c"))

    class CTimeSpec(Structure):
        _fields_ = [("tv_sec", c_long), ("tv_nsec", c_long)]

    class CTimeval(Structure):
        _fields_ = [("tv_sec", c_long), ("tv_usec", c_long)]

    class _TimeSpec(object):
        tv_sec = 0.
        tv_nsec = 0.

    def __init__(self, host='localhost', port=10041, protocol='http',
                 encoding='utf-8'):
        self.host = host
        self.port = port
        self.protocol = protocol
        self.encoding = encoding

    def _usec2nsec(self, nsec):
        return nsec * (1000000000 / 1000000)

    def _clock_gettime(self):
        ret = self._TimeSpec()
        if hasattr(self.LIBRT, 'clock_gettime'):
            timespec = self.CTimeSpec()
            self.LIBRT.clock_gettime(0, pointer(timespec))     # 0: CLOCK_REALTIME
            ret.tv_sec = timespec.tv_sec
            ret.tv_nsec = timespec.tv_nsec
        else:   # MacOSX and other environment
            timespec = self.CTimeval()
            self.LIBC.gettimeofday(pointer(timespec), None)
            ret.tv_sec = timespec.tv_sec
            ret.tv_nsec = self._usec2nsec(timespec.tv_usec)
        return ret

    def _call_gqtp(self, cmd, **kwargs):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((self.host, self.port))
        _cmd = cmd
        _cmd_arg = "".join([" --%s '%s'" % (d, str(kwargs[d]).replace("'", r"\'")) for d in kwargs])
        _cmd = _cmd + _cmd_arg
        _cmd_str = "%08x" % len(_cmd)
        exec("_cmd_len = \"\\x%02s\\x%02s\\x%02s\\x%02s\"" % (
                _cmd_str[:2], _cmd_str[2:4], _cmd_str[4:6], _cmd_str[6:]))
        _header = "".join(["\xc7", "\x00" * 7, _cmd_len, "\x00" * 12])
        _start = self._clock_gettime()
        s.send(_header + _cmd)
        raw_data = s.recv(8192)
        proto, qtype, keylen, level, flags, status, size, opaque, cas \
                = struct.unpack("!BBHBBHIIQ", raw_data[:GQTP_HEADER_SIZE])
        while len(raw_data) < size + GQTP_HEADER_SIZE:
            raw_data += s.recv(8192)
        _end = self._clock_gettime()
        s.close()

        diff_time = (_end.tv_sec + _end.tv_nsec / 1000000000.) - \
                    (_start.tv_sec + _start.tv_nsec / 1000000000.)
        if status != 0:
            status -= 65536
            body = "\"\",[[\"\",\"\",0]]"
            _data = "[[%d,%d.%d,%lf,%s]]" % (
                    status, _start.tv_sec, _start.tv_nsec, diff_time, body)
        else:
            body = raw_data[24:]
            _data = "[[%d,%d.%d,%lf],%s]" % (
                    status, _start.tv_sec, _start.tv_nsec, diff_time, body)
        return _data

    def _call_http(self, cmd, **kwargs):
        domain = [self.protocol, "://", self.host, ":", str(self.port), "/d/"]
        if kwargs:
            url = "".join(domain) + cmd + "?" + urlencode(kwargs)
        else:
            url = "".join(domain) + cmd
        try:
            _data = urlopen(url).read()
        except HTTPError as msg:
            _data = msg.read()
        return _data

    def call(self, cmd, **kwargs):
        output_type = kwargs.get("output_type")
        if not output_type:
            output_type = "json"
        if self.protocol == "http":
            ret = self._call_http(cmd, **kwargs)
        else:
            ret = self._call_gqtp(cmd, **kwargs)
        if cmd == 'select':
            return GroongaSelectResult(ret, output_type, self.encoding)
        else:
            return GroongaResult(ret, output_type, self.encoding)
<|endoftext|>


<commit_before>import csv
try:
    from cStringIO import StringIO
except ImportError:
    try:
        from StringIO import StringIO
    except ImportError:
        from io import StringIO
try:
    import ujson
except ImportError:
    ujson = None
import json
try:
    import msgpack
except ImportError:
    msgpack = None

class GroongaResult(object):

    def __init__(self, data, output_type="json", encoding='utf-8'):
        self.raw_result = data
        if output_type == 'tsv':
            # TODO: not implement
            c = csv.reader(StringIO(data), delimiter='\t')
        elif output_type == 'msgpack':
            if msgpack:
                unpacker = msgpack.Unpacker()
                _result = msgpack.unpackb(data)
            else:
                raise Exception("msgpack is not support")
        else:   # json or other types...
            if ujson and encoding == 'utf-8':
                _result = ujson.loads(data)
            else:
                _result = json.loads(data)
        self.status = _result[0][0]
        self.start_time = _result[0][1]
        self.elapsed = _result[0][2]
        if len(_result) == 1 and self.status != 0:
            self.body = _result[0][3]
        elif len(_result) == 1 and self.status == 0:
            self.body = ""  # handling invalid result
        else:
            self.body = _result[1]

class GroongaSelectResult(GroongaResult):

    def __init__(self, data, output_type="json", encoding='utf-8'):
        super(GroongaSelectResult, self).__init__(data, output_type, encoding)
        self.items = []
        self.hit_num = self.body[0][0][0]
        if self.status == 0:
            keys = [k[0] for k in self.body[0][1]]
            self.items = [
                    dict(list(zip(keys, item))) for item in self.body[0][2:]]
<commit_msg>fix: unable encoding value<commit_after>import csv
try:
    from cStringIO import StringIO
except ImportError:
    try:
        from StringIO import StringIO
    except ImportError:
        from io import StringIO
try:
    import ujson
except ImportError:
    ujson = None
import json
try:
    import msgpack
except ImportError:
    msgpack = None

class GroongaResult(object):

    def __init__(self, data, output_type="json", encoding='utf-8'):
        self.raw_result = data
        if output_type == 'tsv':
            # TODO: not implement
            c = csv.reader(StringIO(data), delimiter='\t')
        elif output_type == 'msgpack':
            if msgpack:
                unpacker = msgpack.Unpacker()
                _result = msgpack.unpackb(data)
            else:
                raise Exception("msgpack is not support")
        else:   # json or other types...
            if ujson and encoding == 'utf-8':
                _result = ujson.loads(data)
            else:
                _result = json.loads(data, encoding)
        self.status = _result[0][0]
        self.start_time = _result[0][1]
        self.elapsed = _result[0][2]
        if len(_result) == 1 and self.status != 0:
            self.body = _result[0][3]
        elif len(_result) == 1 and self.status == 0:
            self.body = ""  # handling invalid result
        else:
            self.body = _result[1]

class GroongaSelectResult(GroongaResult):

    def __init__(self, data, output_type="json", encoding='utf-8'):
        super(GroongaSelectResult, self).__init__(data, output_type, encoding)
        self.items = []
        self.hit_num = self.body[0][0][0]
        if self.status == 0:
            keys = [k[0] for k in self.body[0][1]]
            self.items = [
                    dict(list(zip(keys, item))) for item in self.body[0][2:]]
<|endoftext|>


<commit_before>#!/usr/bin/env python
import argparse
from cStringIO import StringIO
import logging
from math import ceil
from multiprocessing import Pool
import time
import urlparse

import boto
from boto.s3.connection import OrdinaryCallingFormat

parser = argparse.ArgumentParser(description="Transfer large files to S3",
        prog="s3-mp-upload")
parser.add_argument("src", type=file, help="The file to transfer")
parser.add_argument("dest", help="The S3 destination object")
parser.add_argument("-np", "--num-processes", help="Number of processors to use",
        type=int, default=2)
parser.add_argument("-f", "--force", help="Overwrite an existing S3 key",
        action="store_true")
parser.add_argument("-s", "--split", help="Split size, in Mb", type=int, default=50)
parser.add_argument("-rrs", "--reduced-redundancy", help="Use reduced redundancy storage. Default is standard.", default=False,  action="store_true")
parser.add_argument("--insecure", dest='secure', help="Use HTTP for connection",
        default=True, action="store_false")
parser.add_argument("-t", "--max-tries", help="Max allowed retries for http timeout", type=int, default=5)
parser.add_argument("-v", "--verbose", help="Be more verbose", default=False, action="store_true")
parser.add_argument("-q", "--quiet", help="Be less verbose (for use in cron jobs)", default=False, action="store_true")

logger = logging.getLogger("s3-mp-upload")

def do_part_upload(args):
    """
    Upload a part of a MultiPartUpload

    Open the target file and read in a chunk. Since we can't pickle
    S3Connection or MultiPartUpload objects, we have to reconnect and lookup
    the MPU object with each part upload.

    :type args: tuple of (string, string, string, int, int, int)
    :param args: The actual arguments of this method. Due to lameness of
                 multiprocessing, we have to extract these outside of the
                 function definition.

                 The arguments are: S3 Bucket name, MultiPartUpload id, file
                 name, the part number, part offset, part size
    """
    # Multiprocessing args lameness
    bucket_name, mpu_id, fname, i, start, size, secure, max_tries, current_tries = args
    logger.debug("do_part_upload got args: %s" % (args,))

    # Connect to S3, get the MultiPartUpload
    s3 = boto.connect_s3(calling_format=OrdinaryCallingFormat())
    s3.is_secure = secure
    bucket = s3.lookup(bucket_name)
    mpu = None
    for mp in bucket.list_multipart_uploads():
        if mp.id == mpu_id:
            mpu = mp
            break
    if mpu is None:
        raise Exception("Could not find MultiPartUpload %s" % mpu_id)

    # Read the chunk from the file
    fp = open(fname, 'rb')
    fp.seek(start)
    data = fp.read(size)
    fp.close()
    if not data:
        raise Exception("Unexpectedly tried to read an empty chunk")

    def progress(x,y):
        logger.debug("Part %d: %0.2f%%" % (i+1, 100.*x/y))

    try:
        # Do the upload
        t1 = time.time()
        mpu.upload_part_from_file(StringIO(data), i+1, cb=progress)

        # Print some timings
        t2 = time.time() - t1
        s = len(data)/1024./1024.
        logger.info("Uploaded part %s (%0.2fM) in %0.2fs at %0.2fMBps" % (i+1, s, t2, s/t2))
    except Exception, err:
        logger.debug("Retry request %d of max %d times" % (current_tries, max_tries))
        if (current_tries > max_tries):
            logger.error(err)
        else:
            time.sleep(3)
            current_tries += 1
            do_part_upload((bucket_name, mpu_id, fname, i, start, size, secure, max_tries, current_tries))

def main(src, dest, num_processes=2, split=50, force=False, reduced_redundancy=False, verbose=False, quiet=False, secure=True, max_tries=5):
    # Check that dest is a valid S3 url
    split_rs = urlparse.urlsplit(dest)
    if split_rs.scheme != "s3":
        raise ValueError("'%s' is not an S3 url" % dest)

    s3 = boto.connect_s3(calling_format=OrdinaryCallingFormat())
    s3.is_secure = secure
    bucket = s3.lookup(split_rs.netloc)
    if bucket == None:
        raise ValueError("'%s' is not a valid bucket" % split_rs.netloc)
    key = bucket.get_key(split_rs.path)
    # See if we're overwriting an existing key
    if key is not None:
        if not force:
            raise ValueError("'%s' already exists. Specify -f to overwrite it" % dest)

    # Determine the splits
    part_size = max(5*1024*1024, 1024*1024*split)
    src.seek(0,2)
    size = src.tell()
    num_parts = int(ceil(float(size) / part_size))

    # If file is less than 5M, just upload it directly
    if size < 5*1024*1024:
        src.seek(0)
        t1 = time.time()
        k = boto.s3.key.Key(bucket,split_rs.path)
        k.set_contents_from_file(src)
        t2 = time.time() - t1
        s = size/1024./1024.
        logger.info("Finished uploading %0.2fM in %0.2fs (%0.2fMBps)" % (s, t2, s/t2))
        return

    # Create the multi-part upload object
    mpu = bucket.initiate_multipart_upload(split_rs.path, reduced_redundancy=reduced_redundancy)
    logger.info("Initialized upload: %s" % mpu.id)

    # Generate arguments for invocations of do_part_upload
    def gen_args(num_parts, fold_last):
        for i in range(num_parts+1):
            part_start = part_size*i
            if i == (num_parts-1) and fold_last is True:
                yield (bucket.name, mpu.id, src.name, i, part_start, part_size*2, secure, max_tries, 0)
                break
            else:
                yield (bucket.name, mpu.id, src.name, i, part_start, part_size, secure, max_tries, 0)

    # If the last part is less than 5M, just fold it into the previous part
    fold_last = ((size % part_size) < 5*1024*1024)

    # Do the thing
    try:
        # Create a pool of workers
        pool = Pool(processes=num_processes)
        t1 = time.time()
        pool.map_async(do_part_upload, gen_args(num_parts, fold_last)).get(9999999)
        # Print out some timings
        t2 = time.time() - t1
        s = size/1024./1024.
        # Finalize
        src.close()
        mpu.complete_upload()
        logger.info("Finished uploading %0.2fM in %0.2fs (%0.2fMBps)" % (s, t2, s/t2))
    except KeyboardInterrupt:
        logger.warn("Received KeyboardInterrupt, canceling upload")
        pool.terminate()
        mpu.cancel_upload()
    except Exception, err:
        logger.error("Encountered an error, canceling upload")
        logger.error(err)
        mpu.cancel_upload()

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    args = parser.parse_args()
    arg_dict = vars(args)
    if arg_dict['quiet'] == True:
        logger.setLevel(logging.WARNING)
    if arg_dict['verbose'] == True:
        logger.setLevel(logging.DEBUG)
    logger.debug("CLI args: %s" % args)
    main(**arg_dict)
<commit_msg>Adding a bit of information for debug purposes<commit_after>#!/usr/bin/env python
import argparse
from cStringIO import StringIO
import logging
from math import ceil
from multiprocessing import Pool
import time
import urlparse

import boto
from boto.s3.connection import OrdinaryCallingFormat

parser = argparse.ArgumentParser(description="Transfer large files to S3",
        prog="s3-mp-upload")
parser.add_argument("src", type=file, help="The file to transfer")
parser.add_argument("dest", help="The S3 destination object")
parser.add_argument("-np", "--num-processes", help="Number of processors to use",
        type=int, default=2)
parser.add_argument("-f", "--force", help="Overwrite an existing S3 key",
        action="store_true")
parser.add_argument("-s", "--split", help="Split size, in Mb", type=int, default=50)
parser.add_argument("-rrs", "--reduced-redundancy", help="Use reduced redundancy storage. Default is standard.", default=False,  action="store_true")
parser.add_argument("--insecure", dest='secure', help="Use HTTP for connection",
        default=True, action="store_false")
parser.add_argument("-t", "--max-tries", help="Max allowed retries for http timeout", type=int, default=5)
parser.add_argument("-v", "--verbose", help="Be more verbose", default=False, action="store_true")
parser.add_argument("-q", "--quiet", help="Be less verbose (for use in cron jobs)", default=False, action="store_true")

logger = logging.getLogger("s3-mp-upload")

def do_part_upload(args):
    """
    Upload a part of a MultiPartUpload

    Open the target file and read in a chunk. Since we can't pickle
    S3Connection or MultiPartUpload objects, we have to reconnect and lookup
    the MPU object with each part upload.

    :type args: tuple of (string, string, string, int, int, int)
    :param args: The actual arguments of this method. Due to lameness of
                 multiprocessing, we have to extract these outside of the
                 function definition.

                 The arguments are: S3 Bucket name, MultiPartUpload id, file
                 name, the part number, part offset, part size
    """
    # Multiprocessing args lameness
    bucket_name, mpu_id, fname, i, start, size, secure, max_tries, current_tries = args
    logger.debug("do_part_upload got args: %s" % (args,))

    # Connect to S3, get the MultiPartUpload
    s3 = boto.connect_s3(calling_format=OrdinaryCallingFormat())
    s3.is_secure = secure
    bucket = s3.lookup(bucket_name)
    mpu = None
    for mp in bucket.list_multipart_uploads():
        if mp.id == mpu_id:
            mpu = mp
            break
    if mpu is None:
        raise Exception("Could not find MultiPartUpload %s" % mpu_id)

    # Read the chunk from the file
    fp = open(fname, 'rb')
    fp.seek(start)
    data = fp.read(size)
    fp.close()
    if not data:
        raise Exception("Unexpectedly tried to read an empty chunk")

    def progress(x,y):
        logger.debug("Part %d: %0.2f%%" % (i+1, 100.*x/y))

    try:
        # Do the upload
        t1 = time.time()
        mpu.upload_part_from_file(StringIO(data), i+1, cb=progress)

        # Print some timings
        t2 = time.time() - t1
        s = len(data)/1024./1024.
        logger.info("Uploaded part %s (%0.2fM) in %0.2fs at %0.2fMBps" % (i+1, s, t2, s/t2))
    except Exception, err:
        logger.debug("Retry request %d of max %d times" % (current_tries, max_tries))
        if (current_tries > max_tries):
            logger.error(err)
        else:
            time.sleep(3)
            current_tries += 1
            do_part_upload((bucket_name, mpu_id, fname, i, start, size, secure, max_tries, current_tries))

def main(src, dest, num_processes=2, split=50, force=False, reduced_redundancy=False, verbose=False, quiet=False, secure=True, max_tries=5):
    # Check that dest is a valid S3 url
    split_rs = urlparse.urlsplit(dest)
    if split_rs.scheme != "s3":
        raise ValueError("'%s' is not an S3 url" % dest)

    s3 = boto.connect_s3(calling_format=OrdinaryCallingFormat())
    s3.is_secure = secure
    bucket = s3.lookup(split_rs.netloc)
    if bucket == None:
        raise ValueError("'%s' is not a valid bucket" % split_rs.netloc)
    key = bucket.get_key(split_rs.path)
    # See if we're overwriting an existing key
    if key is not None:
        if not force:
            raise ValueError("'%s' already exists. Specify -f to overwrite it" % dest)

    # Determine the splits
    part_size = max(5*1024*1024, 1024*1024*split)
    src.seek(0,2)
    size = src.tell()
    num_parts = int(ceil(float(size) / part_size))

    logger.info("Part Size: {}".format(part_size)
    logger.info("File Size: {}".format(size)
    logger.info("Num Parts: {}".format(num_parts)

    # If file is less than 5M, just upload it directly
    if size < 5*1024*1024:
        src.seek(0)
        t1 = time.time()
        k = boto.s3.key.Key(bucket,split_rs.path)
        k.set_contents_from_file(src)
        t2 = time.time() - t1
        s = size/1024./1024.
        logger.info("Finished uploading %0.2fM in %0.2fs (%0.2fMBps)" % (s, t2, s/t2))
        return

    # Create the multi-part upload object
    mpu = bucket.initiate_multipart_upload(split_rs.path, reduced_redundancy=reduced_redundancy)
    logger.info("Initialized upload: %s" % mpu.id)

    # Generate arguments for invocations of do_part_upload
    def gen_args(num_parts, fold_last):
        for i in range(num_parts+1):
            part_start = part_size*i
            if i == (num_parts-1) and fold_last is True:
                yield (bucket.name, mpu.id, src.name, i, part_start, part_size*2, secure, max_tries, 0)
                break
            else:
                yield (bucket.name, mpu.id, src.name, i, part_start, part_size, secure, max_tries, 0)

    # If the last part is less than 5M, just fold it into the previous part
    fold_last = ((size % part_size) < 5*1024*1024)

    # Do the thing
    try:
        # Create a pool of workers
        pool = Pool(processes=num_processes)
        t1 = time.time()
        pool.map_async(do_part_upload, gen_args(num_parts, fold_last)).get(9999999)
        # Print out some timings
        t2 = time.time() - t1
        s = size/1024./1024.
        # Finalize
        src.close()
        mpu.complete_upload()
        logger.info("Finished uploading %0.2fM in %0.2fs (%0.2fMBps)" % (s, t2, s/t2))
    except KeyboardInterrupt:
        logger.warn("Received KeyboardInterrupt, canceling upload")
        pool.terminate()
        mpu.cancel_upload()
    except Exception, err:
        logger.error("Encountered an error, canceling upload")
        logger.error(err)
        mpu.cancel_upload()

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    args = parser.parse_args()
    arg_dict = vars(args)
    if arg_dict['quiet'] == True:
        logger.setLevel(logging.WARNING)
    if arg_dict['verbose'] == True:
        logger.setLevel(logging.DEBUG)
    logger.debug("CLI args: %s" % args)
    main(**arg_dict)
<|endoftext|>