<commit_before># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# 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.

"""Converts tracks or albums to external directory
"""
import os
import threading
import subprocess
import tempfile
import shlex
from string import Template

from beets import ui, util, plugins, config
from beets.plugins import BeetsPlugin
from beetsplug.embedart import embed_item
from beets.util.confit import ConfigTypeError

_fs_lock = threading.Lock()
_temp_files = []  # Keep track of temporary transcoded files for deletion.

# Some convenient alternate names for formats.
ALIASES = {
    u'wma': u'windows media',
    u'vorbis': u'ogg',
}

LOSSLESS_FORMATS = ['ape', 'flac', 'alac', 'wav']

def replace_ext(path, ext):
    """Return the path with its extension replaced by `ext`.

    The new extension must not contain a leading dot.
    """
    return os.path.splitext(path)[0] + '.' + ext

def get_format(format=None):
    """Return the command tempate and the extension from the config.
    """
    if not format:
        format = config['convert']['format'].get(unicode).lower()
    format = ALIASES.get(format, format)

    try:
        format_info = config['convert']['formats'][format].get(dict)
        command = format_info['command']
        extension = format_info['extension']
    except KeyError:
        raise ui.UserError(
            u'convert: format {0} needs "command" and "extension" fields'
            .format(format)
        )
    except ConfigTypeError:
        command = config['convert']['formats'][format].get(str)
        extension = format

    # Convenience and backwards-compatibility shortcuts.
    keys = config['convert'].keys()
    if 'command' in keys:
        command = config['convert']['command'].get(unicode)
    elif 'opts' in keys:
        # Undocumented option for backwards compatibility with < 1.3.1.
        command = u'ffmpeg -i $source -y {0} $dest'.format(
            config['convert']['opts'].get(unicode)
        )
    if 'extension' in keys:
        extension = config['convert']['extension'].get(unicode)

    return (command.encode('utf8'), extension.encode('utf8'))

def should_transcode(item, format):
    """Determine whether the item should be transcoded as part of
    conversion (i.e., its bitrate is high or it has the wrong format).
    """
    if config['convert']['never_convert_lossy_files'] and \
            not (item.format.lower() in LOSSLESS_FORMATS):
        return False
    maxbr = config['convert']['max_bitrate'].get(int)
    return format.lower() != item.format.lower() or \
        item.bitrate >= 1000 * maxbr

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'pretend': False,
            u'threads': util.cpu_count(),
            u'format': u'mp3',
            u'formats': {
                u'aac': {
                    u'command': u'ffmpeg -i $source -y -vn -acodec libfaac '
                                u'-aq 100 $dest',
                    u'extension': u'm4a',
                },
                u'alac': {
                    u'command': u'ffmpeg -i $source -y -vn -acodec alac $dest',
                    u'extension': u'm4a',
                },
                u'flac': u'ffmpeg -i $source -y -vn -acodec flac $dest',
                u'mp3': u'ffmpeg -i $source -y -vn -aq 2 $dest',
                u'opus':
                    u'ffmpeg -i $source -y -vn -acodec libopus -ab 96k $dest',
                u'ogg':
                    u'ffmpeg -i $source -y -vn -acodec libvorbis -aq 2 $dest',
                u'wma':
                    u'ffmpeg -i $source -y -vn -acodec wmav2 -vn $dest',
            },
            u'max_bitrate': 500,
            u'auto': False,
            u'quiet': False,
            u'embed': True,
            u'paths': {},
            u'never_convert_lossy_files': False,
            u'copy_album_art': False,
        })
        self._import_stages = [self.auto_convert]

        self.register_listener('import_task_files', self._cleanup)

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-p', '--pretend', action='store_true',
                              help='show actions but do nothing')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum available processors')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.parser.add_option('-f', '--format', action='store', dest='format',
                              help='set the destination directory')
        cmd.parser.add_option('-y', '--yes', action='store_true', dest='yes',
                              help='do not ask for confirmation')
        cmd.func = self.convert_func
        return [cmd]

    def auto_convert(self, config, task):
        if self.config['auto']:
            for item in task.imported_items():
                self.convert_on_import(config.lib, item)

    # Utilities converted from functions to methods on logging overhaul

    def encode(self, command, source, dest, pretend=False):
        """Encode `source` to `dest` using command template `command`.

        Raises `subprocess.CalledProcessError` if the command exited with a
        non-zero status code.
        """
        quiet = config['convert']['quiet'].get()

        if not quiet and not pretend:
            self._log.info(u'Encoding {0}', util.displayable_path(source))

        # Substitute $source and $dest in the argument list.
        args = shlex.split(command)
        for i, arg in enumerate(args):
            args[i] = Template(arg).safe_substitute({
                'source': source,
                'dest': dest,
            })

        if pretend:
            self._log.info(' '.join(args))
            return

        try:
            util.command_output(args)
        except subprocess.CalledProcessError as exc:
            # Something went wrong (probably Ctrl+C), remove temporary files
            self._log.info(u'Encoding {0} failed. Cleaning up...',
                           util.displayable_path(source))
            self._log.debug(u'Command {0} exited with status {1}',
                            exc.cmd.decode('utf8', 'ignore'),
                            exc.returncode)
            util.remove(dest)
            util.prune_dirs(os.path.dirname(dest))
            raise
        except OSError as exc:
            raise ui.UserError(
                u"convert: could invoke '{0}': {1}".format(
                    ' '.join(args), exc
                )
            )

        if not quiet and not pretend:
            self._log.info(u'Finished encoding {0}',
                           util.displayable_path(source))

    def convert_item(self, dest_dir, keep_new, path_formats, format,
                     pretend=False):
        command, ext = get_format(format)
        item, original, converted = None, None, None
        while True:
            item = yield (item, original, converted)
            dest = item.destination(basedir=dest_dir,
                                    path_formats=path_formats)

            # When keeping the new file in the library, we first move the
            # current (pristine) file to the destination. We'll then copy it
            # back to its old path or transcode it to a new path.
            if keep_new:
                original = dest
                converted = item.path
                if should_transcode(item, format):
                    converted = replace_ext(converted, ext)
            else:
                original = item.path
                if should_transcode(item, format):
                    dest = replace_ext(dest, ext)
                converted = dest

            # Ensure that only one thread tries to create directories at a
            # time. (The existence check is not atomic with the directory
            # creation inside this function.)
            if not pretend:
                with _fs_lock:
                    util.mkdirall(dest)

            if os.path.exists(util.syspath(dest)):
                self._log.info(u'Skipping {0} (target file exists)',
                               util.displayable_path(item.path))
                continue

            if keep_new:
                if pretend:
                    self._log.info(u'mv {0} {1}',
                                   util.displayable_path(item.path),
                                   util.displayable_path(original))
                else:
                    self._log.info(u'Moving to {0}',
                                   util.displayable_path(original))
                    util.move(item.path, original)

            if should_transcode(item, format):
                try:
                    self.encode(command, original, converted, pretend)
                except subprocess.CalledProcessError:
                    continue
            else:
                if pretend:
                    self._log.info(u'cp {0} {1}',
                                   util.displayable_path(original),
                                   util.displayable_path(converted))
                else:
                    # No transcoding necessary.
                    self._log.info(u'Copying {0}',
                                   util.displayable_path(item.path))
                    util.copy(original, converted)

            if pretend:
                continue

            # Write tags from the database to the converted file.
            item.try_write(path=converted)

            if keep_new:
                # If we're keeping the transcoded file, read it again (after
                # writing) to get new bitrate, duration, etc.
                item.path = converted
                item.read()
                item.store()  # Store new path and audio data.

            if config['convert']['embed']:
                album = item.get_album()
                if album and album.artpath:
                    embed_item(item, album.artpath, itempath=converted)

            if keep_new:
                plugins.send('after_convert', item=item,
                             dest=dest, keepnew=True)
            else:
                plugins.send('after_convert', item=item,
                             dest=converted, keepnew=False)

    def copy_album_art(self, album, dest_dir, path_formats, pretend=False):
        """Copies the associated cover art of the album. Album must have at
        least one track.
        """
        if not album or not album.artpath:
            return

        album_item = album.items().get()
        # Album shouldn't be empty.
        if not album_item:
            return

        # Get the destination of the first item (track) of the album, we use
        # this function to format the path accordingly to path_formats.
        dest = album_item.destination(basedir=dest_dir,
                                      path_formats=path_formats)

        # Remove item from the path.
        dest = os.path.join(*util.components(dest)[:-1])

        dest = album.art_destination(album.artpath, item_dir=dest)
        if album.artpath == dest:
            return

        if not pretend:
            util.mkdirall(dest)

        if os.path.exists(util.syspath(dest)):
            self._log.info(u'Skipping {0} (target file exists)',
                           util.displayable_path(album.artpath))
            return

        if pretend:
            self._log.info(u'cp {0} {1}',
                           util.displayable_path(album.artpath),
                           util.displayable_path(dest))
        else:
            self._log.info(u'Copying cover art to {0}',
                           util.displayable_path(dest))
            util.copy(album.artpath, dest)

    def convert_func(self, lib, opts, args):
        if not opts.dest:
            opts.dest = config['convert']['dest'].get()
        if not opts.dest:
            raise ui.UserError('no convert destination set')
        opts.dest = util.bytestring_path(opts.dest)

        if not opts.threads:
            opts.threads = config['convert']['threads'].get(int)

        if config['convert']['paths']:
            path_formats = ui.get_path_formats(config['convert']['paths'])
        else:
            path_formats = ui.get_path_formats()

        if not opts.format:
            opts.format = config['convert']['format'].get(unicode).lower()

        pretend = opts.pretend if opts.pretend is not None else \
            config['convert']['pretend'].get(bool)

        if not pretend:
            ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

            if not (opts.yes or ui.input_yn("Convert? (Y/n)")):
                return

        if opts.album:
            albums = lib.albums(ui.decargs(args))
            items = (i for a in albums for i in a.items())
            if config['convert']['copy_album_art']:
                for album in albums:
                    self.copy_album_art(album, opts.dest, path_formats,
                                        pretend)
        else:
            items = iter(lib.items(ui.decargs(args)))
        convert = [self.convert_item(opts.dest,
                                     opts.keep_new,
                                     path_formats,
                                     opts.format,
                                     pretend)
                   for _ in range(opts.threads)]
        pipe = util.pipeline.Pipeline([items, convert])
        pipe.run_parallel()

    def convert_on_import(self, lib, item):
        """Transcode a file automatically after it is imported into the
        library.
        """
        format = config['convert']['format'].get(unicode).lower()
        if should_transcode(item, format):
            command, ext = get_format()
            fd, dest = tempfile.mkstemp('.' + ext)
            os.close(fd)
            _temp_files.append(dest)  # Delete the transcode later.
            try:
                self.encode(command, item.path, dest)
            except subprocess.CalledProcessError:
                return
            item.path = dest
            item.write()
            item.read()  # Load new audio information data.
            item.store()

    def _cleanup(self, task, session):
        for path in task.old_paths:
            if path in _temp_files:
                if os.path.isfile(path):
                    util.remove(path)
                _temp_files.remove(path)
<commit_msg>Fix usage of embedart by convert plugin<commit_after># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# 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.

"""Converts tracks or albums to external directory
"""
import os
import threading
import subprocess
import tempfile
import shlex
from string import Template

from beets import ui, util, plugins, config
from beets.plugins import BeetsPlugin
from beetsplug.embedart import EmbedCoverArtPlugin
from beets.util.confit import ConfigTypeError

_fs_lock = threading.Lock()
_temp_files = []  # Keep track of temporary transcoded files for deletion.

# Some convenient alternate names for formats.
ALIASES = {
    u'wma': u'windows media',
    u'vorbis': u'ogg',
}

LOSSLESS_FORMATS = ['ape', 'flac', 'alac', 'wav']

def replace_ext(path, ext):
    """Return the path with its extension replaced by `ext`.

    The new extension must not contain a leading dot.
    """
    return os.path.splitext(path)[0] + '.' + ext

def get_format(format=None):
    """Return the command tempate and the extension from the config.
    """
    if not format:
        format = config['convert']['format'].get(unicode).lower()
    format = ALIASES.get(format, format)

    try:
        format_info = config['convert']['formats'][format].get(dict)
        command = format_info['command']
        extension = format_info['extension']
    except KeyError:
        raise ui.UserError(
            u'convert: format {0} needs "command" and "extension" fields'
            .format(format)
        )
    except ConfigTypeError:
        command = config['convert']['formats'][format].get(str)
        extension = format

    # Convenience and backwards-compatibility shortcuts.
    keys = config['convert'].keys()
    if 'command' in keys:
        command = config['convert']['command'].get(unicode)
    elif 'opts' in keys:
        # Undocumented option for backwards compatibility with < 1.3.1.
        command = u'ffmpeg -i $source -y {0} $dest'.format(
            config['convert']['opts'].get(unicode)
        )
    if 'extension' in keys:
        extension = config['convert']['extension'].get(unicode)

    return (command.encode('utf8'), extension.encode('utf8'))

def should_transcode(item, format):
    """Determine whether the item should be transcoded as part of
    conversion (i.e., its bitrate is high or it has the wrong format).
    """
    if config['convert']['never_convert_lossy_files'] and \
            not (item.format.lower() in LOSSLESS_FORMATS):
        return False
    maxbr = config['convert']['max_bitrate'].get(int)
    return format.lower() != item.format.lower() or \
        item.bitrate >= 1000 * maxbr

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'pretend': False,
            u'threads': util.cpu_count(),
            u'format': u'mp3',
            u'formats': {
                u'aac': {
                    u'command': u'ffmpeg -i $source -y -vn -acodec libfaac '
                                u'-aq 100 $dest',
                    u'extension': u'm4a',
                },
                u'alac': {
                    u'command': u'ffmpeg -i $source -y -vn -acodec alac $dest',
                    u'extension': u'm4a',
                },
                u'flac': u'ffmpeg -i $source -y -vn -acodec flac $dest',
                u'mp3': u'ffmpeg -i $source -y -vn -aq 2 $dest',
                u'opus':
                    u'ffmpeg -i $source -y -vn -acodec libopus -ab 96k $dest',
                u'ogg':
                    u'ffmpeg -i $source -y -vn -acodec libvorbis -aq 2 $dest',
                u'wma':
                    u'ffmpeg -i $source -y -vn -acodec wmav2 -vn $dest',
            },
            u'max_bitrate': 500,
            u'auto': False,
            u'quiet': False,
            u'embed': True,
            u'paths': {},
            u'never_convert_lossy_files': False,
            u'copy_album_art': False,
        })
        self._import_stages = [self.auto_convert]

        self.register_listener('import_task_files', self._cleanup)

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-p', '--pretend', action='store_true',
                              help='show actions but do nothing')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum available processors')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.parser.add_option('-f', '--format', action='store', dest='format',
                              help='set the destination directory')
        cmd.parser.add_option('-y', '--yes', action='store_true', dest='yes',
                              help='do not ask for confirmation')
        cmd.func = self.convert_func
        return [cmd]

    def auto_convert(self, config, task):
        if self.config['auto']:
            for item in task.imported_items():
                self.convert_on_import(config.lib, item)

    # Utilities converted from functions to methods on logging overhaul

    def encode(self, command, source, dest, pretend=False):
        """Encode `source` to `dest` using command template `command`.

        Raises `subprocess.CalledProcessError` if the command exited with a
        non-zero status code.
        """
        quiet = config['convert']['quiet'].get()

        if not quiet and not pretend:
            self._log.info(u'Encoding {0}', util.displayable_path(source))

        # Substitute $source and $dest in the argument list.
        args = shlex.split(command)
        for i, arg in enumerate(args):
            args[i] = Template(arg).safe_substitute({
                'source': source,
                'dest': dest,
            })

        if pretend:
            self._log.info(' '.join(args))
            return

        try:
            util.command_output(args)
        except subprocess.CalledProcessError as exc:
            # Something went wrong (probably Ctrl+C), remove temporary files
            self._log.info(u'Encoding {0} failed. Cleaning up...',
                           util.displayable_path(source))
            self._log.debug(u'Command {0} exited with status {1}',
                            exc.cmd.decode('utf8', 'ignore'),
                            exc.returncode)
            util.remove(dest)
            util.prune_dirs(os.path.dirname(dest))
            raise
        except OSError as exc:
            raise ui.UserError(
                u"convert: could invoke '{0}': {1}".format(
                    ' '.join(args), exc
                )
            )

        if not quiet and not pretend:
            self._log.info(u'Finished encoding {0}',
                           util.displayable_path(source))

    def convert_item(self, dest_dir, keep_new, path_formats, format,
                     pretend=False):
        command, ext = get_format(format)
        item, original, converted = None, None, None
        while True:
            item = yield (item, original, converted)
            dest = item.destination(basedir=dest_dir,
                                    path_formats=path_formats)

            # When keeping the new file in the library, we first move the
            # current (pristine) file to the destination. We'll then copy it
            # back to its old path or transcode it to a new path.
            if keep_new:
                original = dest
                converted = item.path
                if should_transcode(item, format):
                    converted = replace_ext(converted, ext)
            else:
                original = item.path
                if should_transcode(item, format):
                    dest = replace_ext(dest, ext)
                converted = dest

            # Ensure that only one thread tries to create directories at a
            # time. (The existence check is not atomic with the directory
            # creation inside this function.)
            if not pretend:
                with _fs_lock:
                    util.mkdirall(dest)

            if os.path.exists(util.syspath(dest)):
                self._log.info(u'Skipping {0} (target file exists)',
                               util.displayable_path(item.path))
                continue

            if keep_new:
                if pretend:
                    self._log.info(u'mv {0} {1}',
                                   util.displayable_path(item.path),
                                   util.displayable_path(original))
                else:
                    self._log.info(u'Moving to {0}',
                                   util.displayable_path(original))
                    util.move(item.path, original)

            if should_transcode(item, format):
                try:
                    self.encode(command, original, converted, pretend)
                except subprocess.CalledProcessError:
                    continue
            else:
                if pretend:
                    self._log.info(u'cp {0} {1}',
                                   util.displayable_path(original),
                                   util.displayable_path(converted))
                else:
                    # No transcoding necessary.
                    self._log.info(u'Copying {0}',
                                   util.displayable_path(item.path))
                    util.copy(original, converted)

            if pretend:
                continue

            # Write tags from the database to the converted file.
            item.try_write(path=converted)

            if keep_new:
                # If we're keeping the transcoded file, read it again (after
                # writing) to get new bitrate, duration, etc.
                item.path = converted
                item.read()
                item.store()  # Store new path and audio data.

            if config['convert']['embed']:
                album = item.get_album()
                if album and album.artpath:
                    EmbedCoverArtPlugin().embed_item(item, album.artpath,
                                                     itempath=converted)

            if keep_new:
                plugins.send('after_convert', item=item,
                             dest=dest, keepnew=True)
            else:
                plugins.send('after_convert', item=item,
                             dest=converted, keepnew=False)

    def copy_album_art(self, album, dest_dir, path_formats, pretend=False):
        """Copies the associated cover art of the album. Album must have at
        least one track.
        """
        if not album or not album.artpath:
            return

        album_item = album.items().get()
        # Album shouldn't be empty.
        if not album_item:
            return

        # Get the destination of the first item (track) of the album, we use
        # this function to format the path accordingly to path_formats.
        dest = album_item.destination(basedir=dest_dir,
                                      path_formats=path_formats)

        # Remove item from the path.
        dest = os.path.join(*util.components(dest)[:-1])

        dest = album.art_destination(album.artpath, item_dir=dest)
        if album.artpath == dest:
            return

        if not pretend:
            util.mkdirall(dest)

        if os.path.exists(util.syspath(dest)):
            self._log.info(u'Skipping {0} (target file exists)',
                           util.displayable_path(album.artpath))
            return

        if pretend:
            self._log.info(u'cp {0} {1}',
                           util.displayable_path(album.artpath),
                           util.displayable_path(dest))
        else:
            self._log.info(u'Copying cover art to {0}',
                           util.displayable_path(dest))
            util.copy(album.artpath, dest)

    def convert_func(self, lib, opts, args):
        if not opts.dest:
            opts.dest = config['convert']['dest'].get()
        if not opts.dest:
            raise ui.UserError('no convert destination set')
        opts.dest = util.bytestring_path(opts.dest)

        if not opts.threads:
            opts.threads = config['convert']['threads'].get(int)

        if config['convert']['paths']:
            path_formats = ui.get_path_formats(config['convert']['paths'])
        else:
            path_formats = ui.get_path_formats()

        if not opts.format:
            opts.format = config['convert']['format'].get(unicode).lower()

        pretend = opts.pretend if opts.pretend is not None else \
            config['convert']['pretend'].get(bool)

        if not pretend:
            ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

            if not (opts.yes or ui.input_yn("Convert? (Y/n)")):
                return

        if opts.album:
            albums = lib.albums(ui.decargs(args))
            items = (i for a in albums for i in a.items())
            if config['convert']['copy_album_art']:
                for album in albums:
                    self.copy_album_art(album, opts.dest, path_formats,
                                        pretend)
        else:
            items = iter(lib.items(ui.decargs(args)))
        convert = [self.convert_item(opts.dest,
                                     opts.keep_new,
                                     path_formats,
                                     opts.format,
                                     pretend)
                   for _ in range(opts.threads)]
        pipe = util.pipeline.Pipeline([items, convert])
        pipe.run_parallel()

    def convert_on_import(self, lib, item):
        """Transcode a file automatically after it is imported into the
        library.
        """
        format = config['convert']['format'].get(unicode).lower()
        if should_transcode(item, format):
            command, ext = get_format()
            fd, dest = tempfile.mkstemp('.' + ext)
            os.close(fd)
            _temp_files.append(dest)  # Delete the transcode later.
            try:
                self.encode(command, item.path, dest)
            except subprocess.CalledProcessError:
                return
            item.path = dest
            item.write()
            item.read()  # Load new audio information data.
            item.store()

    def _cleanup(self, task, session):
        for path in task.old_paths:
            if path in _temp_files:
                if os.path.isfile(path):
                    util.remove(path)
                _temp_files.remove(path)
<|endoftext|>


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

"""
image_scraper is used to find all images from a web page
"""
# @author chengdujin
# @contact chengdujin@gmail.com
# @created Jul. 23, 2013

import sys
reload(sys)
sys.setdefaultencoding('UTF-8')
sys.path.append('..')

from BeautifulSoup import BeautifulSoup
from cStringIO import StringIO
from data_processor import thumbnail
import Image
import os
import re
import urllib2
import urlparse

# CONSTANTS
from config import IMAGES_LOCAL_DIR
from config import IMAGES_PUBLIC_DIR
from config import MIN_IMAGE_SIZE
from config import TRANSCODED_LOCAL_DIR

if not os.path.exists(IMAGES_LOCAL_DIR):
    os.mkdir(IMAGES_LOCAL_DIR)

def _link_process(link):
    """
    get rid of cdn prefix
    """
    link = link.replace("\/", "/")
    image_url_complex = urllib2.unquote(link.strip())
    if image_url_complex:
        # as the name could be http://xxx.com/yyy--http://zzz.jpg
        # or http://xxx.com/yyy--https://zzz.jpg
        last_http_index = image_url_complex.rfind('http')
        image_url = image_url_complex[last_http_index:]
        # response is the signal of a valid image
        response = None
        try:
            response = urllib2.urlopen(image_url)
        except urllib2.URLError as k:
            path = re.split('https?://?', image_url)[-1]
            scheme = urlparse.urlparse(image_url).scheme
            image_url = '%s://%s' % (scheme, path)
            try:
                response = urllib2.urlopen(image_url)
            except urllib2.URLError as k:
                pass
            except Exception as k:
                print k
        if response:
            return image_url

def find_image(link=None):
    """
    find an image from the link
    """
    if not link:
        return None

    link = _link_process(link)
    image_normalized = normalize(link)
    return image_normalized[0] if image_normalized else None

def find_images(content=None):
    """
    find out all images from content and its size info
    """
    if not content:
        return None

    # determine the type of content
    if isinstance(content, str) and content.startswith(TRANSCODED_LOCAL_DIR):
        # then its a file
        f = open(content, 'r')
        content = f.read()

    soup = BeautifulSoup(content.decode('utf-8', 'ignore'))
    images_normalized = []
    images = soup.findAll('img')
    for image in images:
        if image.get('src'):
            image_normalized = find_image(image.get('src'))
            if image_normalized:
                images_normalized.append(image_normalized)
    return images_normalized

def find_biggest_image(images=None):
    """
    find the biggest in resolution from a pile of images
    """
    if not images:
        return None

    biggest = None
    for image in images:
        resolution_max = MIN_IMAGE_SIZE[0] * MIN_IMAGE_SIZE[1]
        resolution_image = int(image['width']) * int(image['height'])
        if resolution_image > resolution_max:
            biggest = image
            resolution_max = resolution_image
    return biggest

def dedupe_images(images):
    """
    remove same images
    image: {'url':xxx, 'width':yyy, 'height':zzz}
    images = [image, image, image]
    """
    if not images:
        return None

    image_urls = []

    def _exists(image):
        """
        return boolean if image exists in list image_urls
        """
        exists = image['url'] in image_urls
        if not exists:
            image_urls.append(image['url'])
            return False
        else:
            return True

    return filter(lambda x: not _exists(x), images)

# TODO: boundary checker
def scale_image(image=None, size_expected=MIN_IMAGE_SIZE,
                resize_by_width=True, crop_by_center=True, relative_path=None):
    """
    resize an image as requested
    resize_by_width: resize image according to its width(True)/height(False)
    crop_by_center: crop image from its center(True) or by point(0, 0)(False)
    """
    if not image or not size_expected or not relative_path:
        return None, None

    width = int(image['width'])
    height = int(image['height'])
    width_expected = size_expected[0]
    height_expected = size_expected[1]

    if width >= width_expected and height >= height_expected:
        if resize_by_width:
            height_new = width_expected * height / width
            width_new = width_expected
        else:
            width_new = height_expected * width / height
            height_new = height_expected

        # larger and equal than is important here
        if width_new >= width_expected and height_new >= height_expected:
            # resize
            size_new = width_new, height_new
            image_downloaded = StringIO(urllib2.urlopen(image['url']).read())
            image_data = Image.open(image_downloaded)
            image_data.thumbnail(size_new, Image.ANTIALIAS)
            # crop
            if crop_by_center:
                left = (width_new - width_expected) / 2
                top = (height_new - height_expected) / 2
                right = (width_new + width_expected) / 2
                bottom = (height_new + height_expected) / 2
                image_cropped = image_data.crop((left, top, right, bottom))
            else:
                left = 0
                top = 0
                right = width_expected
                bottom = height_expected
                image_cropped = image_data.crop((left, top, right, bottom))
            # storing
            if image_cropped:
                image_web_path = '%s%s.jpg' % (
                    IMAGES_PUBLIC_DIR, relative_path)
                image_local_path = '%s%s.jpg' % (
                    IMAGES_LOCAL_DIR, relative_path)
                image_cropped = image_cropped.convert('RGB')
                image_cropped.save(image_local_path, 'JPEG')
                return {'url': image_web_path, 'width': width_expected, 'height': height_expected}, {'url': image_local_path, 'width': width_expected, 'height': height_expected}
            else:
                return None, None
        else:
            return scale_image(image, size_expected, not resize_by_width, crop_by_center, relative_path)
    else:
        return None, None

def normalize(images):
    """
    for list of images, remove images that don't match with MIN_IMAGE_SIZE;
    for an image, return None if it doesn't matches with MIN_IMAGE_SIZE
    """
    def _check_image(image):
        """
        check an image if it matches with MIN_IMAGE_SIZE
        """
        if not image:
            raise Exception('ERROR: Method not well formed!')

        try:
            if 'url' in image:
                image_url = image['url']
                if thumbnail.is_valid_image(image_url):
                    return image
            else:
                if thumbnail.is_valid_image(image):
                    width, height = thumbnail.get_image_size(image)
                    return {'url': image, 'width': width, 'height': height}
            return None
        except IOError as k:
            print k
            return None

    if isinstance(images, str) or isinstance(images, unicode):
        image = _check_image(images)
        return [image] if image else None
    elif isinstance(images, list):
        images_new = []
        for image in images:
            image_new = _check_image(image)
            if image_new:
                images_new.append(image_new)
        return images_new if images_new else None
<commit_msg>added timeout to urlopen<commit_after>#!/usr/bin/env python
#-*- coding: utf-8 -*-

"""
image_scraper is used to find all images from a web page
"""
# @author chengdujin
# @contact chengdujin@gmail.com
# @created Jul. 23, 2013

import sys
reload(sys)
sys.setdefaultencoding('UTF-8')
sys.path.append('..')

from BeautifulSoup import BeautifulSoup
from cStringIO import StringIO
from data_processor import thumbnail
import Image
import os
import re
import urllib2
import urlparse

# CONSTANTS
from config import IMAGES_LOCAL_DIR
from config import IMAGES_PUBLIC_DIR
from config import MIN_IMAGE_SIZE
from config import TRANSCODED_LOCAL_DIR
from config import UCK_TIMEOUT

if not os.path.exists(IMAGES_LOCAL_DIR):
    os.mkdir(IMAGES_LOCAL_DIR)

def _link_process(link):
    """
    get rid of cdn prefix
    """
    link = link.replace("\/", "/")
    image_url_complex = urllib2.unquote(link.strip())
    if image_url_complex:
        # as the name could be http://xxx.com/yyy--http://zzz.jpg
        # or http://xxx.com/yyy--https://zzz.jpg
        last_http_index = image_url_complex.rfind('http')
        image_url = image_url_complex[last_http_index:]
        # response is the signal of a valid image
        response = None
        try:
            response = urllib2.urlopen(image_url, timeout=UCK_TIMEOUT)
        except urllib2.URLError as k:
            path = re.split('https?://?', image_url)[-1]
            scheme = urlparse.urlparse(image_url).scheme
            image_url = '%s://%s' % (scheme, path)
            try:
                response = urllib2.urlopen(image_url, timeout=UCK_TIMEOUT)
            except urllib2.URLError as k:
                pass
            except Exception as k:
                print k
        if response:
            return image_url

def find_image(link=None):
    """
    find an image from the link
    """
    if not link:
        return None

    link = _link_process(link)
    image_normalized = normalize(link)
    return image_normalized[0] if image_normalized else None

def find_images(content=None):
    """
    find out all images from content and its size info
    """
    if not content:
        return None

    # determine the type of content
    if isinstance(content, str) and content.startswith(TRANSCODED_LOCAL_DIR):
        # then its a file
        f = open(content, 'r')
        content = f.read()

    soup = BeautifulSoup(content.decode('utf-8', 'ignore'))
    images_normalized = []
    images = soup.findAll('img')
    for image in images:
        if image.get('src'):
            image_normalized = find_image(image.get('src'))
            if image_normalized:
                images_normalized.append(image_normalized)
    return images_normalized

def find_biggest_image(images=None):
    """
    find the biggest in resolution from a pile of images
    """
    if not images:
        return None

    biggest = None
    for image in images:
        resolution_max = MIN_IMAGE_SIZE[0] * MIN_IMAGE_SIZE[1]
        resolution_image = int(image['width']) * int(image['height'])
        if resolution_image > resolution_max:
            biggest = image
            resolution_max = resolution_image
    return biggest

def dedupe_images(images):
    """
    remove same images
    image: {'url':xxx, 'width':yyy, 'height':zzz}
    images = [image, image, image]
    """
    if not images:
        return None

    image_urls = []

    def _exists(image):
        """
        return boolean if image exists in list image_urls
        """
        exists = image['url'] in image_urls
        if not exists:
            image_urls.append(image['url'])
            return False
        else:
            return True

    return filter(lambda x: not _exists(x), images)

# TODO: boundary checker
def scale_image(image=None, size_expected=MIN_IMAGE_SIZE,
                resize_by_width=True, crop_by_center=True, relative_path=None):
    """
    resize an image as requested
    resize_by_width: resize image according to its width(True)/height(False)
    crop_by_center: crop image from its center(True) or by point(0, 0)(False)
    """
    if not image or not size_expected or not relative_path:
        return None, None

    width = int(image['width'])
    height = int(image['height'])
    width_expected = size_expected[0]
    height_expected = size_expected[1]

    if width >= width_expected and height >= height_expected:
        if resize_by_width:
            height_new = width_expected * height / width
            width_new = width_expected
        else:
            width_new = height_expected * width / height
            height_new = height_expected

        # larger and equal than is important here
        if width_new >= width_expected and height_new >= height_expected:
            # resize
            size_new = width_new, height_new
            image_downloaded = StringIO(urllib2.urlopen(image['url'], timeout=UCK_TIMEOUT).read())
            image_data = Image.open(image_downloaded)
            image_data.thumbnail(size_new, Image.ANTIALIAS)
            # crop
            if crop_by_center:
                left = (width_new - width_expected) / 2
                top = (height_new - height_expected) / 2
                right = (width_new + width_expected) / 2
                bottom = (height_new + height_expected) / 2
                image_cropped = image_data.crop((left, top, right, bottom))
            else:
                left = 0
                top = 0
                right = width_expected
                bottom = height_expected
                image_cropped = image_data.crop((left, top, right, bottom))
            # storing
            if image_cropped:
                image_web_path = '%s%s.jpg' % (
                    IMAGES_PUBLIC_DIR, relative_path)
                image_local_path = '%s%s.jpg' % (
                    IMAGES_LOCAL_DIR, relative_path)
                image_cropped = image_cropped.convert('RGB')
                image_cropped.save(image_local_path, 'JPEG')
                return {'url': image_web_path, 'width': width_expected, 'height': height_expected}, {'url': image_local_path, 'width': width_expected, 'height': height_expected}
            else:
                return None, None
        else:
            return scale_image(image, size_expected, not resize_by_width, crop_by_center, relative_path)
    else:
        return None, None

def normalize(images):
    """
    for list of images, remove images that don't match with MIN_IMAGE_SIZE;
    for an image, return None if it doesn't matches with MIN_IMAGE_SIZE
    """
    def _check_image(image):
        """
        check an image if it matches with MIN_IMAGE_SIZE
        """
        if not image:
            raise Exception('ERROR: Method not well formed!')

        try:
            if 'url' in image:
                image_url = image['url']
                if thumbnail.is_valid_image(image_url):
                    return image
            else:
                if thumbnail.is_valid_image(image):
                    width, height = thumbnail.get_image_size(image)
                    return {'url': image, 'width': width, 'height': height}
            return None
        except IOError as k:
            print k
            return None

    if isinstance(images, str) or isinstance(images, unicode):
        image = _check_image(images)
        return [image] if image else None
    elif isinstance(images, list):
        images_new = []
        for image in images:
            image_new = _check_image(image)
            if image_new:
                images_new.append(image_new)
        return images_new if images_new else None
<|endoftext|>


<commit_before>from uuid import uuid4
from cqlengine.exceptions import ValidationError
from cqlengine.query import QueryException

from cqlengine.tests.base import BaseCassEngTestCase
from cqlengine.models import Model
from cqlengine.management import sync_table, drop_table
from cqlengine import columns

class TestQueryUpdateModel(Model):
    partition   = columns.UUID(primary_key=True, default=uuid4)
    cluster     = columns.Integer(primary_key=True)
    count       = columns.Integer(required=False)
    text        = columns.Text(required=False, index=True)
    text_set    = columns.Set(columns.Text, required=False)
    text_list   = columns.List(columns.Text, required=False)
    text_map    = columns.Map(columns.Text, columns.Text, required=False)

class QueryUpdateTests(BaseCassEngTestCase):

    @classmethod
    def setUpClass(cls):
        super(QueryUpdateTests, cls).setUpClass()
        sync_table(TestQueryUpdateModel)

    @classmethod
    def tearDownClass(cls):
        super(QueryUpdateTests, cls).tearDownClass()
        drop_table(TestQueryUpdateModel)

    def test_update_values(self):
        """ tests calling udpate on a queryset """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count=6)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == (6 if i == 3 else i)
            assert row.text == str(i)

    def test_update_values_validation(self):
        """ tests calling udpate on models with values passed in """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count='asdf')

    def test_invalid_update_kwarg(self):
        """ tests that passing in a kwarg to the update method that isn't a column will fail """
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=uuid4(), cluster=3).update(bacon=5000)

    def test_primary_key_update_failure(self):
        """ tests that attempting to update the value of a primary key will fail """
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=uuid4(), cluster=3).update(cluster=5000)

    def test_null_update_deletes_column(self):
        """ setting a field to null in the update should issue a delete statement """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(text=None)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == (None if i == 3 else str(i))

    def test_mixed_value_and_null_update(self):
        """ tests that updating a columns value, and removing another works properly """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count=6, text=None)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == (6 if i == 3 else i)
            assert row.text == (None if i == 3 else str(i))

    def test_counter_updates(self):
        pass

    def test_set_add_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(text_set__add={'bar'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"foo", "bar"})

    def test_set_add_updates_new_record(self):
        """ If the key doesn't exist yet, an update creates the record
        """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(text_set__add={'bar'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"bar"})

    def test_set_remove_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo", "baz"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_set__remove={'foo'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"baz"})

    def test_set_remove_new_record(self):
        """ Removing something not in the set should silently do nothing
        """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_set__remove={'afsd'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"foo"})

    def test_list_append_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_list=["foo"])
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_list__append=['bar'])
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_list, ["foo", "bar"])

    def test_list_prepend_updates(self):
        """ Prepend two things since order is reversed by default by CQL """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_list=["foo"])
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_list__prepend=['bar', 'baz'])
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_list, ["bar", "baz", "foo"])

    def test_map_merge_updates(self):
        """ Merge a dictionary into existing value """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster,
                text_map={"foo": '1', "bar": '2'})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_map__merge={"bar": '3', "baz": '4'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_map, {"foo": '1', "bar": '3', "baz": '4'})
<commit_msg>Tries to add a test that doesn't work because of a bug in the underlying cql library.  So doing an update with a dict __merge with a None value will not be possible for the time being.<commit_after>from uuid import uuid4
from cqlengine.exceptions import ValidationError
from cqlengine.query import QueryException

from cqlengine.tests.base import BaseCassEngTestCase
from cqlengine.models import Model
from cqlengine.management import sync_table, drop_table
from cqlengine import columns

class TestQueryUpdateModel(Model):
    partition   = columns.UUID(primary_key=True, default=uuid4)
    cluster     = columns.Integer(primary_key=True)
    count       = columns.Integer(required=False)
    text        = columns.Text(required=False, index=True)
    text_set    = columns.Set(columns.Text, required=False)
    text_list   = columns.List(columns.Text, required=False)
    text_map    = columns.Map(columns.Text, columns.Text, required=False)

class QueryUpdateTests(BaseCassEngTestCase):

    @classmethod
    def setUpClass(cls):
        super(QueryUpdateTests, cls).setUpClass()
        sync_table(TestQueryUpdateModel)

    @classmethod
    def tearDownClass(cls):
        super(QueryUpdateTests, cls).tearDownClass()
        drop_table(TestQueryUpdateModel)

    def test_update_values(self):
        """ tests calling udpate on a queryset """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count=6)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == (6 if i == 3 else i)
            assert row.text == str(i)

    def test_update_values_validation(self):
        """ tests calling udpate on models with values passed in """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count='asdf')

    def test_invalid_update_kwarg(self):
        """ tests that passing in a kwarg to the update method that isn't a column will fail """
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=uuid4(), cluster=3).update(bacon=5000)

    def test_primary_key_update_failure(self):
        """ tests that attempting to update the value of a primary key will fail """
        with self.assertRaises(ValidationError):
            TestQueryUpdateModel.objects(partition=uuid4(), cluster=3).update(cluster=5000)

    def test_null_update_deletes_column(self):
        """ setting a field to null in the update should issue a delete statement """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(text=None)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == (None if i == 3 else str(i))

    def test_mixed_value_and_null_update(self):
        """ tests that updating a columns value, and removing another works properly """
        partition = uuid4()
        for i in range(5):
            TestQueryUpdateModel.create(partition=partition, cluster=i, count=i, text=str(i))

        # sanity check
        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == i
            assert row.text == str(i)

        # perform update
        TestQueryUpdateModel.objects(partition=partition, cluster=3).update(count=6, text=None)

        for i, row in enumerate(TestQueryUpdateModel.objects(partition=partition)):
            assert row.cluster == i
            assert row.count == (6 if i == 3 else i)
            assert row.text == (None if i == 3 else str(i))

    def test_counter_updates(self):
        pass

    def test_set_add_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(text_set__add={'bar'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"foo", "bar"})

    def test_set_add_updates_new_record(self):
        """ If the key doesn't exist yet, an update creates the record
        """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(text_set__add={'bar'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"bar"})

    def test_set_remove_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo", "baz"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_set__remove={'foo'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"baz"})

    def test_set_remove_new_record(self):
        """ Removing something not in the set should silently do nothing
        """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_set={"foo"})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_set__remove={'afsd'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_set, {"foo"})

    def test_list_append_updates(self):
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_list=["foo"])
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_list__append=['bar'])
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_list, ["foo", "bar"])

    def test_list_prepend_updates(self):
        """ Prepend two things since order is reversed by default by CQL """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster, text_list=["foo"])
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_list__prepend=['bar', 'baz'])
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_list, ["bar", "baz", "foo"])

    def test_map_merge_updates(self):
        """ Merge a dictionary into existing value """
        partition = uuid4()
        cluster = 1
        TestQueryUpdateModel.objects.create(
                partition=partition, cluster=cluster,
                text_map={"foo": '1', "bar": '2'})
        TestQueryUpdateModel.objects(
                partition=partition, cluster=cluster).update(
                text_map__merge={"bar": '3', "baz": '4'})
        obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        self.assertEqual(obj.text_map, {"foo": '1', "bar": '3', "baz": '4'})

    def test_map_merge_none_deletes_key(self):
        """ The CQL behavior is if you set a key in a map to null it deletes
        that key from the map.  Test that this works with __merge.

        This test fails because of a bug in the cql python library not
        converting None to null (and the cql library is no longer in active
        developement).
        """
        # partition = uuid4()
        # cluster = 1
        # TestQueryUpdateModel.objects.create(
        #         partition=partition, cluster=cluster,
        #         text_map={"foo": '1', "bar": '2'})
        # TestQueryUpdateModel.objects(
        #         partition=partition, cluster=cluster).update(
        #         text_map__merge={"bar": None})
        # obj = TestQueryUpdateModel.objects.get(partition=partition, cluster=cluster)
        # self.assertEqual(obj.text_map, {"foo": '1'})
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import json
import paho.mqtt.client as mqtt
import time
from influxdb import InfluxDBClient

# The callback for when the client receives a CONNACK response from the server.
def on_connect(client, userdata, flags, rc):
    #print("Connected with result code "+str(rc))

    # Subscribing in on_connect() means that if we lose the connection and
    # reconnect then subscriptions will be renewed.
    client.subscribe("k34/#")
    client.subscribe("sensors/#")

# The callback for when a PUBLISH message is received from the server.
def on_message(client, userdata, msg):
    #print(msg.topic)
    #print(str(msg.payload))
    try:
      pl = json.loads(msg.payload)
      if msg.topic == "k34/bmp085":
        json_body = [
          {
            "measurement": "pressure",
            "time": pl["time"] * 1000000,
            "fields": {
              "pressure": pl["pressure"],
              "temperature": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/heat/water":
        json_body = [
          {
            "measurement": "heatWater",
            "time": pl["timestamp"] * 1000000,
            "fields": {
              "feedTemp": pl["temperature"]["feed"],
              "returnTemp": pl["temperature"]["return"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "sensors/power/p1meter/actual_consumption":
        json_body = [
          {
            "measurement": "power",
            "time": long(time.time() * 1000000000),
            "fields": {
              "power": pl
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/tempout":
        json_body = [
          {
            "measurement": "tempout",
            "time": long(time.time() * 1000000000),
            "fields": {
              "tempout": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/tempgarage":
        json_body = [
          {
            "measurement": "tempgarage",
            "time": long(time.time() * 1000000000),
            "fields": {
              "tempout": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/heat/burner":
        json_body = [
          {
            "measurement": "oilBurner",
            "time": pl["timestamp"] * 1000000,
            "fields": {
              "burner": pl["burner"]
            }
          }
        ]
        db.write_points(json_body)
      else:
        pass
    except:
      print("something went wrong, probably unintelligible payload")

db = InfluxDBClient("localhost", 8086, "root", "root", "k34db")

client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message

client.connect("localhost", 1883, 60)

# Blocking call that processes network traffic, dispatches callbacks and
# handles reconnecting.
client.loop_forever()
<commit_msg>casted power value to int<commit_after># -*- coding: utf-8 -*-

import json
import paho.mqtt.client as mqtt
import time
from influxdb import InfluxDBClient

# The callback for when the client receives a CONNACK response from the server.
def on_connect(client, userdata, flags, rc):
    #print("Connected with result code "+str(rc))

    # Subscribing in on_connect() means that if we lose the connection and
    # reconnect then subscriptions will be renewed.
    client.subscribe("k34/#")
    client.subscribe("sensors/#")

# The callback for when a PUBLISH message is received from the server.
def on_message(client, userdata, msg):
    #print(msg.topic)
    #print(str(msg.payload))
    try:
      pl = json.loads(msg.payload)
      if msg.topic == "k34/bmp085":
        json_body = [
          {
            "measurement": "pressure",
            "time": pl["time"] * 1000000,
            "fields": {
              "pressure": pl["pressure"],
              "temperature": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/heat/water":
        json_body = [
          {
            "measurement": "heatWater",
            "time": pl["timestamp"] * 1000000,
            "fields": {
              "feedTemp": pl["temperature"]["feed"],
              "returnTemp": pl["temperature"]["return"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "sensors/power/p1meter/actual_consumption":
        json_body = [
          {
            "measurement": "power",
            "time": long(time.time() * 1000000000),
            "fields": {
              "power": int(pl)
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/tempout":
        json_body = [
          {
            "measurement": "tempout",
            "time": long(time.time() * 1000000000),
            "fields": {
              "tempout": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/tempgarage":
        json_body = [
          {
            "measurement": "tempgarage",
            "time": long(time.time() * 1000000000),
            "fields": {
              "tempout": pl["temperature"]
            }
          }
        ]
        db.write_points(json_body)
      elif msg.topic == "k34/heat/burner":
        json_body = [
          {
            "measurement": "oilBurner",
            "time": pl["timestamp"] * 1000000,
            "fields": {
              "burner": pl["burner"]
            }
          }
        ]
        db.write_points(json_body)
      else:
        pass
    except:
      print("something went wrong, probably unintelligible payload")

db = InfluxDBClient("localhost", 8086, "root", "root", "k34db")

client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message

client.connect("localhost", 1883, 60)

# Blocking call that processes network traffic, dispatches callbacks and
# handles reconnecting.
client.loop_forever()
<|endoftext|>


<commit_before>import os
import mmap
import cPickle as pickle
import operator
from functools32 import lru_cache

from ..data import Store, ObjectStore
from nel import logging
log = logging.getLogger()

class FileObjectStore(ObjectStore):
    def __init__(self, path):
        self.path = path
        self.store = mmdict(path)

    @classmethod
    def get_protocol(cls):
        return 'file'

    def iter_ids(self):
        return self.store.iterkeys()

    def exists(self, oid):
        return oid in self.store

    def fetch(self, oid):
        return self.store[oid]

    def fetch_many(self, oids):
        return [self.fetch(oid) for oid in oids]

    def fetch_all(self):
        return self.store.iteritems()

    def save_many(self, obj_iter):
        self.store.close()
        mmdict.write(self.path, ((o['_id'], o) for o in obj_iter))
        self.store = mmdict(self.path)

    @classmethod
    def GetPath(cls, store_id, uri):
        path = store_id.replace(':', '/')
        if uri and uri.startswith('file://'):
            path = os.path.join(uri[7:], path)
        return path

    @classmethod
    def Get(cls, store_id, uri='file://', **kwargs):
        return cls(cls.GetPath(store_id, uri))

class mmdict(object):
    def __init__(self, path):
        self.path = path
        self.index = {}

        index_path = self.path + '.index'
        if os.path.exists(index_path):
            log.debug('Loading mmap store: %s ...' % index_path)
            with open(index_path, 'rb') as f:
                self.index = dict(self.deserialise(f))

            self.data_file = open(path + '.data', 'rb')
            self.data_mmap = mmap.mmap(self.data_file.fileno(), 0, prot=mmap.PROT_READ)
        else:
            log.warn('No existing mmap store found: %s ...' % index_path)

    @staticmethod
    def serialise(obj, f):
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

    @staticmethod
    def deserialise(f):
        return pickle.load(f)

    @staticmethod
    def static_itervalues(path):
        with open(path + '.data', 'rb') as f:
            while True:
                try:
                    yield mmdict.deserialise(f)
                except EOFError: break

    def iteritems(self):
        sorted_idx = sorted(self.index.iteritems(), key=operator.itemgetter(1))

        for i, v in enumerate(self.itervalues()):
            yield (sorted_idx[i][0], v)

    def iterkeys(self):
        return self.index.iterkeys()

    def itervalues(self):
        self.data_mmap.seek(0)
        while True:
            try:
                yield self.deserialise(self.data_mmap)
            except EOFError: break

    def __len__(self):
        return len(self.index)

    def __contains__(self, key):
        return key in self.index

    @lru_cache(maxsize=20000)
    def __getitem__(self, key):
        if key not in self:
            return None

        self.data_mmap.seek(self.index[key])
        return self.deserialise(self.data_mmap)

    def __enter__(self):
        return self

    def close(self):
        if hasattr(self, 'data_mmap') and self.data_mmap != None:
            self.data_mmap.close()
        if hasattr(self, 'data_file') and self.data_file != None:
            self.data_file.close()

    def __exit__(self, type, value, traceback):
        self.close()

    def __del__(self):
        self.close()

    @staticmethod
    def write(path, iter_kvs):
        index = []
        with open(path + '.data', 'wb') as f:
            for key, value in iter_kvs:
                index.append((key, f.tell()))
                mmdict.serialise(value, f)
        with open(path + '.index','wb') as f:
            mmdict.serialise(index, f)
<commit_msg>Update file.py<commit_after>import os
import mmap
import cPickle as pickle
import operator
from functools32 import lru_cache

from ..data import Store, ObjectStore
from nel import logging
log = logging.getLogger()

class FileObjectStore(ObjectStore):
    def __init__(self, path):
        self.path = path
        self.store = mmdict(path)

    @classmethod
    def get_protocol(cls):
        return 'file'

    def iter_ids(self):
        return self.store.iterkeys()

    def exists(self, oid):
        return oid in self.store

    def fetch(self, oid):
        return self.store[oid]

    def fetch_many(self, oids):
        return [self.fetch(oid) for oid in oids]

    def fetch_all(self):
        return self.store.iteritems()

    def save_many(self, obj_iter):
        self.store.close()
        mmdict.write(self.path, ((o['_id'], o) for o in obj_iter))
        self.store = mmdict(self.path)

    def save(self,obj):
        self.save_man([obj])

    @classmethod
    def GetPath(cls, store_id, uri):
        path = store_id.replace(':', '/')
        if uri and uri.startswith('file://'):
            path = os.path.join(uri[7:], path)
        return path

    @classmethod
    def Get(cls, store_id, uri='file://', **kwargs):
        return cls(cls.GetPath(store_id, uri))

class mmdict(object):
    def __init__(self, path):
        self.path = path
        self.index = {}

        index_path = self.path + '.index'
        if os.path.exists(index_path):
            log.debug('Loading mmap store: %s ...' % index_path)
            with open(index_path, 'rb') as f:
                self.index = dict(self.deserialise(f))

            self.data_file = open(path + '.data', 'rb')
            self.data_mmap = mmap.mmap(self.data_file.fileno(), 0, prot=mmap.PROT_READ)
        else:
            log.warn('No existing mmap store found: %s ...' % index_path)

    @staticmethod
    def serialise(obj, f):
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

    @staticmethod
    def deserialise(f):
        return pickle.load(f)

    @staticmethod
    def static_itervalues(path):
        with open(path + '.data', 'rb') as f:
            while True:
                try:
                    yield mmdict.deserialise(f)
                except EOFError: break

    def iteritems(self):
        sorted_idx = sorted(self.index.iteritems(), key=operator.itemgetter(1))

        for i, v in enumerate(self.itervalues()):
            yield (sorted_idx[i][0], v)

    def iterkeys(self):
        return self.index.iterkeys()

    def itervalues(self):
        self.data_mmap.seek(0)
        while True:
            try:
                yield self.deserialise(self.data_mmap)
            except EOFError: break

    def __len__(self):
        return len(self.index)

    def __contains__(self, key):
        return key in self.index

    @lru_cache(maxsize=20000)
    def __getitem__(self, key):
        if key not in self:
            return None

        self.data_mmap.seek(self.index[key])
        return self.deserialise(self.data_mmap)

    def __enter__(self):
        return self

    def close(self):
        if hasattr(self, 'data_mmap') and self.data_mmap != None:
            self.data_mmap.close()
        if hasattr(self, 'data_file') and self.data_file != None:
            self.data_file.close()

    def __exit__(self, type, value, traceback):
        self.close()

    def __del__(self):
        self.close()

    @staticmethod
    def write(path, iter_kvs):
        index = []
        with open(path + '.data', 'wb') as f:
            for key, value in iter_kvs:
                index.append((key, f.tell()))
                mmdict.serialise(value, f)
        with open(path + '.index','wb') as f:
            mmdict.serialise(index, f)
<|endoftext|>


<commit_before>#
# Copyright (C) 2011 - 2017 Satoru SATOH <ssato @ redhat.com>
# License: MIT
#
# Some XML modules may be missing and Base.{load,dumps}_impl are not overriden:
# pylint: disable=import-error
"""XML files parser backend, should be available always.

.. versionchanged:: 0.1.0
   Added XML dump support.

- Format to support: XML, e.g. http://www.w3.org/TR/xml11/
- Requirements: one of the followings

  - lxml2.etree if available
  - xml.etree.ElementTree in standard lib if python >= 2.5
  - elementtree.ElementTree (otherwise)

- Limitations:

  - '<prefix>attrs', '<prefix>text' and '<prefix>children' are used as special
    parameter to keep XML structure of original data. You have to cusomize
    <prefix> (default: '@') if any config parameters conflict with some of
    them.

  - Some data or structures of original XML file may be lost if make it backed
    to XML file; XML file - (anyconfig.load) -> config - (anyconfig.dump) ->
    XML file

  - XML specific features (namespace, etc.) may not be processed correctly.

- Special Options: None supported
"""
from __future__ import absolute_import
from io import BytesIO

import re
import sys
try:
    # First, try lxml which is compatible with elementtree and looks faster a
    # lot. See also: http://getpython3.com/diveintopython3/xml.html
    from lxml2 import etree as ET
except ImportError:
    try:
        import xml.etree.ElementTree as ET
    except ImportError:
        import elementtree.ElementTree as ET

import anyconfig.backend.base
import anyconfig.compat
import anyconfig.mdicts

_PREFIX = "@"

# It seems that ET.ElementTree.write() cannot process a parameter
# 'xml_declaration' in older python < 2.7:
_IS_OLDER_PYTHON = sys.version_info[0] < 3 and sys.version_info[1] < 7
_ET_NS_RE = re.compile(r"^{(\S+)}(\S+)$")

def _iterparse(xmlfile):
    """
    Avoid bug in python 3.{2,3}. See http://bugs.python.org/issue9257.

    :param xmlfile: XML file or file-like object
    """
    try:
        return ET.iterparse(xmlfile, events=("start-ns", ))
    except TypeError:
        return ET.iterparse(xmlfile, events=(b"start-ns", ))

def flip(tpl):
    """
    >>> flip((1, 2))
    (2, 1)
    """
    return (tpl[1], tpl[0])

def _namespaces_from_file(xmlfile):
    """
    :param xmlfile: XML file or file-like object
    :return: {namespace_uri: namespace_prefix} or {}
    """
    return dict(flip(t) for _, t in _iterparse(xmlfile))

def _gen_tags(pprefix=_PREFIX):
    """
    Generate special prefixed tags.

    :param pprefix: Special parameter name prefix
    :return: A tuple of prefixed (attributes, text, children)
    """
    return tuple(pprefix + x for x in ("attrs", "text", "children"))

def _tweak_ns(tag, nspaces):
    """
    :param tag: XML tag element
    :param nspaces: A namespaces dict, {uri: prefix}

    >>> _tweak_ns("a", {})
    'a'
    >>> _tweak_ns("a", {"http://example.com/ns/val/": "val"})
    'a'
    >>> _tweak_ns("{http://example.com/ns/val/}a",
    ...           {"http://example.com/ns/val/": "val"})
    'val:a'
    """
    if nspaces:
        matched = _ET_NS_RE.match(tag)
        if matched:
            (uri, tag) = matched.groups()
            prefix = nspaces.get(uri, False)
            if prefix:
                return "%s:%s" % (prefix, tag)

    return tag

def root_to_container(root, to_container, nspaces, pprefix=_PREFIX):
    """
    Convert XML ElementTree Root Element to a collection of container objects.

    :param root: etree root object or None
    :param to_container: callble to make a container object
    :param nspaces: A namespaces dict, {uri: prefix}
    :param pprefix: Special parameter name prefix
    """
    tree = to_container()
    if root is None:
        return tree

    if nspaces is None:
        nspaces = dict()

    if nspaces:
        for uri, prefix in nspaces.items():
            root.attrib["xmlns:" + prefix if prefix else "xmlns"] = uri

    return elem_to_container(root, to_container, nspaces, _gen_tags(pprefix))

def elem_to_container(elem, to_container, nspaces, tags=False):
    """
    Convert XML ElementTree Element to a collection of container objects.

    :param elem: etree elem object or None
    :param to_container: callble to make a container object
    :param nspaces: A namespaces dict, {uri: prefix}
    :param tags: (attrs, text, children) parameter names
    """
    tree = to_container()
    if elem is None:
        return tree

    subtree = tree[_tweak_ns(elem.tag, nspaces)] = to_container()
    (attrs, text, children) = tags if tags else _gen_tags()
    _num_of_children = len(elem)

    if elem.attrib:
        subtree[attrs] = to_container(elem.attrib)

    if elem.text:
        elem.text = elem.text.strip()
        if elem.text:
            if not _num_of_children and not elem.attrib:
                # .. note:: Treat as special case for later convenience.
                tree[elem.tag] = elem.text
            else:
                subtree[text] = elem.text

    if _num_of_children:
        # Note: Configuration item cannot have both attributes and values
        # (list) at the same time in current implementation:
        args = (to_container, nspaces, tags)
        if _num_of_children == 1:  # .. note:: Another special case.
            tree[elem.tag] = [elem_to_container(c, *args) for c in elem][0]
        else:
            subtree[children] = [elem_to_container(c, *args) for c in elem]

    return tree

def container_to_etree(obj, parent=None, pprefix=_PREFIX):
    """
    Convert a dict-like object to XML ElementTree.

    :param obj: Container instance to convert to
    :param parent: XML ElementTree parent node object or None
    :param pprefix: Special parameter name prefix
    """
    if not anyconfig.mdicts.is_dict_like(obj):
        if parent is not None and obj:
            parent.text = obj  # Parent is a leaf text node.
        return  # All attributes and text should be set already.

    (attrs, text, children) = _gen_tags(pprefix)
    for key, val in anyconfig.compat.iteritems(obj):
        if key == attrs:
            for attr, aval in anyconfig.compat.iteritems(val):
                parent.set(attr, aval)
        elif key == text:
            parent.text = val
        elif key == children:
            for child in val:  # child should be a dict-like object.
                for ckey, cval in anyconfig.compat.iteritems(child):
                    celem = ET.Element(ckey)
                    container_to_etree(cval, celem, pprefix)
                    parent.append(celem)
        elif parent is not None:
            elem = ET.Element(key)
            container_to_etree(val, elem, pprefix)
            parent.append(elem)
            return ET.ElementTree(parent)
        else:
            elem = ET.Element(key)
            container_to_etree(val, elem, pprefix)
            return ET.ElementTree(elem)

def etree_write(tree, stream):
    """
    Write XML ElementTree `root` content into `stream`.

    :param tree: XML ElementTree object
    :param stream: File or file-like object can write to
    """
    if _IS_OLDER_PYTHON:
        tree.write(stream, encoding='UTF-8')
    else:
        tree.write(stream, encoding='UTF-8', xml_declaration=True)

class Parser(anyconfig.backend.base.ToStreamDumper):
    """
    Parser for XML files.
    """
    _type = "xml"
    _extensions = ["xml"]
    _open_flags = ('rb', 'wb')

    def load_from_string(self, content, to_container, **kwargs):
        """
        Load config from XML snippet (a string `content`).

        :param content: XML snippet (a string)
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters passed to

        :return: Dict-like object holding config parameters
        """
        root = ET.fromstring(content)
        nspaces = _namespaces_from_file(anyconfig.compat.StringIO(content))
        return root_to_container(root, to_container, nspaces)

    def load_from_path(self, filepath, to_container, **kwargs):
        """
        :param filepath: XML file path
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters to be sanitized

        :return: Dict-like object holding config parameters
        """
        root = ET.parse(filepath).getroot()
        nspaces = _namespaces_from_file(filepath)
        return root_to_container(root, to_container, nspaces)

    def load_from_stream(self, stream, to_container, **kwargs):
        """
        :param stream: XML file or file-like object
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters to be sanitized

        :return: Dict-like object holding config parameters
        """
        return self.load_from_path(stream, to_container, **kwargs)

    def dump_to_string(self, cnf, **kwargs):
        """
        :param cnf: Configuration data to dump
        :param kwargs: optional keyword parameters

        :return: string represents the configuration
        """
        tree = container_to_etree(cnf)
        buf = BytesIO()
        etree_write(tree, buf)
        return buf.getvalue()

    def dump_to_stream(self, cnf, stream, **kwargs):
        """
        :param cnf: Configuration data to dump
        :param stream: Config file or file like object write to
        :param kwargs: optional keyword parameters
        """
        tree = container_to_etree(cnf)
        etree_write(tree, stream)

# vim:sw=4:ts=4:et:
<commit_msg>enhancement: try to import cElementTree if available in .backend.xml<commit_after>#
# Copyright (C) 2011 - 2017 Satoru SATOH <ssato @ redhat.com>
# License: MIT
#
# Some XML modules may be missing and Base.{load,dumps}_impl are not overriden:
# pylint: disable=import-error
"""XML files parser backend, should be available always.

.. versionchanged:: 0.1.0
   Added XML dump support.

- Format to support: XML, e.g. http://www.w3.org/TR/xml11/
- Requirements: one of the followings

  - lxml2.etree if available
  - xml.etree.ElementTree in standard lib if python >= 2.5
  - elementtree.ElementTree (otherwise)

- Limitations:

  - '<prefix>attrs', '<prefix>text' and '<prefix>children' are used as special
    parameter to keep XML structure of original data. You have to cusomize
    <prefix> (default: '@') if any config parameters conflict with some of
    them.

  - Some data or structures of original XML file may be lost if make it backed
    to XML file; XML file - (anyconfig.load) -> config - (anyconfig.dump) ->
    XML file

  - XML specific features (namespace, etc.) may not be processed correctly.

- Special Options: None supported
"""
from __future__ import absolute_import
from io import BytesIO

import re
import sys
try:
    # First, try lxml which is compatible with elementtree and looks faster a
    # lot. See also: http://getpython3.com/diveintopython3/xml.html
    from lxml2 import etree as ET
except ImportError:
    try:
        import xml.etree.cElementTree as ET
    except ImportError:
        import xml.etree.ElementTree as ET
    except ImportError:
        import elementtree.ElementTree as ET

import anyconfig.backend.base
import anyconfig.compat
import anyconfig.mdicts

_PREFIX = "@"

# It seems that ET.ElementTree.write() cannot process a parameter
# 'xml_declaration' in older python < 2.7:
_IS_OLDER_PYTHON = sys.version_info[0] < 3 and sys.version_info[1] < 7
_ET_NS_RE = re.compile(r"^{(\S+)}(\S+)$")

def _iterparse(xmlfile):
    """
    Avoid bug in python 3.{2,3}. See http://bugs.python.org/issue9257.

    :param xmlfile: XML file or file-like object
    """
    try:
        return ET.iterparse(xmlfile, events=("start-ns", ))
    except TypeError:
        return ET.iterparse(xmlfile, events=(b"start-ns", ))

def flip(tpl):
    """
    >>> flip((1, 2))
    (2, 1)
    """
    return (tpl[1], tpl[0])

def _namespaces_from_file(xmlfile):
    """
    :param xmlfile: XML file or file-like object
    :return: {namespace_uri: namespace_prefix} or {}
    """
    return dict(flip(t) for _, t in _iterparse(xmlfile))

def _gen_tags(pprefix=_PREFIX):
    """
    Generate special prefixed tags.

    :param pprefix: Special parameter name prefix
    :return: A tuple of prefixed (attributes, text, children)
    """
    return tuple(pprefix + x for x in ("attrs", "text", "children"))

def _tweak_ns(tag, nspaces):
    """
    :param tag: XML tag element
    :param nspaces: A namespaces dict, {uri: prefix}

    >>> _tweak_ns("a", {})
    'a'
    >>> _tweak_ns("a", {"http://example.com/ns/val/": "val"})
    'a'
    >>> _tweak_ns("{http://example.com/ns/val/}a",
    ...           {"http://example.com/ns/val/": "val"})
    'val:a'
    """
    if nspaces:
        matched = _ET_NS_RE.match(tag)
        if matched:
            (uri, tag) = matched.groups()
            prefix = nspaces.get(uri, False)
            if prefix:
                return "%s:%s" % (prefix, tag)

    return tag

def root_to_container(root, to_container, nspaces, pprefix=_PREFIX):
    """
    Convert XML ElementTree Root Element to a collection of container objects.

    :param root: etree root object or None
    :param to_container: callble to make a container object
    :param nspaces: A namespaces dict, {uri: prefix}
    :param pprefix: Special parameter name prefix
    """
    tree = to_container()
    if root is None:
        return tree

    if nspaces is None:
        nspaces = dict()

    if nspaces:
        for uri, prefix in nspaces.items():
            root.attrib["xmlns:" + prefix if prefix else "xmlns"] = uri

    return elem_to_container(root, to_container, nspaces, _gen_tags(pprefix))

def elem_to_container(elem, to_container, nspaces, tags=False):
    """
    Convert XML ElementTree Element to a collection of container objects.

    :param elem: etree elem object or None
    :param to_container: callble to make a container object
    :param nspaces: A namespaces dict, {uri: prefix}
    :param tags: (attrs, text, children) parameter names
    """
    tree = to_container()
    if elem is None:
        return tree

    subtree = tree[_tweak_ns(elem.tag, nspaces)] = to_container()
    (attrs, text, children) = tags if tags else _gen_tags()
    _num_of_children = len(elem)

    if elem.attrib:
        subtree[attrs] = to_container(elem.attrib)

    if elem.text:
        elem.text = elem.text.strip()
        if elem.text:
            if not _num_of_children and not elem.attrib:
                # .. note:: Treat as special case for later convenience.
                tree[elem.tag] = elem.text
            else:
                subtree[text] = elem.text

    if _num_of_children:
        # Note: Configuration item cannot have both attributes and values
        # (list) at the same time in current implementation:
        args = (to_container, nspaces, tags)
        if _num_of_children == 1:  # .. note:: Another special case.
            tree[elem.tag] = [elem_to_container(c, *args) for c in elem][0]
        else:
            subtree[children] = [elem_to_container(c, *args) for c in elem]

    return tree

def container_to_etree(obj, parent=None, pprefix=_PREFIX):
    """
    Convert a dict-like object to XML ElementTree.

    :param obj: Container instance to convert to
    :param parent: XML ElementTree parent node object or None
    :param pprefix: Special parameter name prefix
    """
    if not anyconfig.mdicts.is_dict_like(obj):
        if parent is not None and obj:
            parent.text = obj  # Parent is a leaf text node.
        return  # All attributes and text should be set already.

    (attrs, text, children) = _gen_tags(pprefix)
    for key, val in anyconfig.compat.iteritems(obj):
        if key == attrs:
            for attr, aval in anyconfig.compat.iteritems(val):
                parent.set(attr, aval)
        elif key == text:
            parent.text = val
        elif key == children:
            for child in val:  # child should be a dict-like object.
                for ckey, cval in anyconfig.compat.iteritems(child):
                    celem = ET.Element(ckey)
                    container_to_etree(cval, celem, pprefix)
                    parent.append(celem)
        elif parent is not None:
            elem = ET.Element(key)
            container_to_etree(val, elem, pprefix)
            parent.append(elem)
            return ET.ElementTree(parent)
        else:
            elem = ET.Element(key)
            container_to_etree(val, elem, pprefix)
            return ET.ElementTree(elem)

def etree_write(tree, stream):
    """
    Write XML ElementTree `root` content into `stream`.

    :param tree: XML ElementTree object
    :param stream: File or file-like object can write to
    """
    if _IS_OLDER_PYTHON:
        tree.write(stream, encoding='UTF-8')
    else:
        tree.write(stream, encoding='UTF-8', xml_declaration=True)

class Parser(anyconfig.backend.base.ToStreamDumper):
    """
    Parser for XML files.
    """
    _type = "xml"
    _extensions = ["xml"]
    _open_flags = ('rb', 'wb')

    def load_from_string(self, content, to_container, **kwargs):
        """
        Load config from XML snippet (a string `content`).

        :param content: XML snippet (a string)
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters passed to

        :return: Dict-like object holding config parameters
        """
        root = ET.fromstring(content)
        nspaces = _namespaces_from_file(anyconfig.compat.StringIO(content))
        return root_to_container(root, to_container, nspaces)

    def load_from_path(self, filepath, to_container, **kwargs):
        """
        :param filepath: XML file path
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters to be sanitized

        :return: Dict-like object holding config parameters
        """
        root = ET.parse(filepath).getroot()
        nspaces = _namespaces_from_file(filepath)
        return root_to_container(root, to_container, nspaces)

    def load_from_stream(self, stream, to_container, **kwargs):
        """
        :param stream: XML file or file-like object
        :param to_container: callble to make a container object
        :param kwargs: optional keyword parameters to be sanitized

        :return: Dict-like object holding config parameters
        """
        return self.load_from_path(stream, to_container, **kwargs)

    def dump_to_string(self, cnf, **kwargs):
        """
        :param cnf: Configuration data to dump
        :param kwargs: optional keyword parameters

        :return: string represents the configuration
        """
        tree = container_to_etree(cnf)
        buf = BytesIO()
        etree_write(tree, buf)
        return buf.getvalue()

    def dump_to_stream(self, cnf, stream, **kwargs):
        """
        :param cnf: Configuration data to dump
        :param stream: Config file or file like object write to
        :param kwargs: optional keyword parameters
        """
        tree = container_to_etree(cnf)
        etree_write(tree, stream)

# vim:sw=4:ts=4:et:
<|endoftext|>