<commit_before># This file is part of Miro Community.
# Copyright (C) 2010 Participatory Culture Foundation
#
# Miro Community is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# Miro Community is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Miro Community.  If not, see <http://www.gnu.org/licenses/>.

import os
import simplejson
import multiprocessing
import logging

import feedparser

import eventlet
import eventlet.pools

import httplib

from django.db import transaction
from django.core.management.base import BaseCommand, CommandError
from vidscraper.bulk_import import bulk_import_url_list, bulk_import

from localtv import models
import localtv.util
import localtv.tasks

DEFAULT_HTTPLIB_CACHE_PATH='/tmp/.cache-for-uid-%d' % os.getuid()

def function_for_fork_worker(data_tuple):
    video_id, future_status = data_tuple
    import localtv.management.commands.update_one_thumbnail
    cmd = localtv.management.commands.update_one_thumbnail.Command()
    try:
        cmd.handle(video_id, future_status)
    except models.Video.DoesNotExist:
        logging.warn("For some reason, we failed to find the model with ID %d" % (
                video_id, ))
        return False # Aww, shucks. Maybe after a retry this will work.
    except Exception, e:
        logging.warn("uh, bizarre -- the task fell over. Maybe it will work on a retry.")
        logging.exception(e)
        return False
    return True

class Command(BaseCommand):

    args = '[feed primary key]'

    def handle(self, *args, **options):
        if len(args) != 1:
            raise CommandError('bulk_import takes one argument: '
                               '%i argument(s) given' % len(args))
        try:
            feed = models.Feed.objects.get(pk=args[0])
        except models.Feed.DoesNotExist:
            raise CommandError('Feed with pk %s does not exist' % args[0])

        try:
            self.verbose = (int(options['verbosity']) > 1)
        except (KeyError, ValueError):
            self.verbose = False

        httplib2 = eventlet.import_patched('httplib2')

        # Use httplib2 to GET the feed.
        # This permits us to download it only once, passing a parsed_feed through
        # to the vidscraper functions.
        h = httplib2.Http(DEFAULT_HTTPLIB_CACHE_PATH)
        content = localtv.util.http_get(feed.feed_url, _httplib2=httplib2)

        parsed_feed = feedparser.parse(content)

        # Try to work asynchronously, calling feedparser ourselves. We can do that
        # if the importer supports bulk_import_url_list.
        try:
            feed_urls = bulk_import_url_list(parsed_feed=parsed_feed)
        except ValueError:
            return self.use_old_bulk_import(parsed_feed, feed)
        # Okay, good, we either got the feed_url list, or we passed the work
        # off the old-style function. Proceed.
        self.forked_tasks = {}
        self.forked_task_worker_pool =  multiprocessing.Pool(processes=8)
        # start 8 worker processes. That should be fine, right?
        self.bulk_import_asynchronously(parsed_feed, h, feed_urls, feed)
        self.enqueue_forked_tasks_for_thumbnail_fetches(feed)

    @transaction.commit_manually
    def bulk_import_asynchronously(self, original_parsed_feed, h, feed_urls, feed):
        # This asynchronous bulk_import is a parallelism monster.

        # Also, for some reason, this monster causes MySQL to frequently just
        # disconnect us with an error like:
        # OperationalError: (2013, 'Lost connection to MySQL server during query')
        # so we set this flag called keep_going to True, and repeat the body of this code
        # for up to 8 tries, if that's how long it takes.
        MAX_TRIES_FOR_BULK_IMPORT = 8
        num_tries_so_far = 0
        while num_tries_so_far < MAX_TRIES_FOR_BULK_IMPORT:
            try:
                result = self._bulk_import_asynchronously_for_real(original_parsed_feed, h, feed_urls, feed)
                return result
            except Exception, e:
                logging.warn("oh, sad, the bulk import fell over. Maybe it will work on a retry.")
                logging.exception(e)
                num_tries_so_far += 1

    @transaction.commit_manually
    def _bulk_import_asynchronously_for_real(self, original_parsed_feed, h, feed_urls, feed):
        # We do as much network I/O as we can using eventlet,
        # rather than threads or subprocesses.
        httplib2 = eventlet.import_patched('httplib2')

        httppool = eventlet.pools.Pool(max_size=200)
        httppool.create = lambda: httplib2.Http(DEFAULT_HTTPLIB_CACHE_PATH)

        pool = eventlet.GreenPool(100)

        def get_url(url):
            with httppool.item() as http:
                if self.verbose:
                    print 'getting', url
                try:
                    resp, content = http.request(url, 'GET')
                except httplib.BadStatusLine, e:
                    if not e.args[0]: # timeout
                        resp, content = e, ''
                    else:
                        raise
                return (resp, content)

        def cache_thumbnail_url(url):
            with httppool.item() as http:
                if self.verbose:
                    print 'getting thumb', url
                localtv.util.cache_downloaded_file(url, http)

        stats = {
            'total': 0,
            'imported': 0,
            'skipped': 0
            }

        def handle_one_sub_feed(feed_contents):
            parsed_feed = feedparser.parse(feed_contents)
            # For each feed entry in this small sub-feed, handle the item.
            for index, entry in enumerate(parsed_feed['entries'][::-1]):
                yield handle_one_item(index, parsed_feed, entry)

        def handle_one_item(index, parsed_feed, entry):
            i = feed._handle_one_bulk_import_feed_entry(index, parsed_feed, entry, verbose=self.verbose, clear_rejected=False,
                                                        actually_save_thumbnails=False)
            # Enqueue the work to download the thumbnail
            if i['video']:
                v = i['video']
                thumbnail_url = v.thumbnail_url
                if thumbnail_url:
                    cache_thumbnail_url(thumbnail_url)

                # The _handle_one_bulk_import_feed_entry() method gave the
                # video a status, but we have to take that back for now.
                #
                # We set the status to VIDEO_STATUS_PENDING_THUMBNAIL so that
                # no one can see the video until the thumbnailing process is
                # complete.
                #
                # We pass the thumbnailer the status that the video should get
                # so that it can set that once it is ready.
                i['future_status'] = v.status
                v.status = models.VIDEO_STATUS_PENDING_THUMBNAIL
                v.save()

            stats['total'] += 1
            if i['video'] is not None:
                stats['imported'] += 1
            else:
                stats['skipped'] += 1

            return i

        results = []
        for (response, content) in pool.imap(get_url, feed_urls):
            try:
                # We make it a list so that we can iterate across
                # it more than once.
                result = list(handle_one_sub_feed(content))
            except:
                transaction.rollback()
                raise
            else:
                transaction.commit()

            try:
                results.extend(result)
                # Now that handle_one_sub_feed has finished, it is
                # safe to spawn forked tasks to do thumbnail fetching.
                for i in result:
                    video = i['video']
                    if video:
                        future_status = i['future_status']
                        self._enqueue_one_forked_task_for_thumbnail_fetch(
                            video.id, future_status)
            except:
                transaction.rollback()
                raise
            else:
                transaction.commit()

        # Get all the thumbnail URLs, and once you have them
        pool.waitall() # wait for the thumbnails

        try:
            # Finish marking the Feed as imported.
            feed._mark_bulk_import_as_done(original_parsed_feed)

            feed.status = models.FEED_STATUS_ACTIVE
            feed.save()
            return_me = [i['video'].id for i in results if i['video']]
        except:
            transaction.rollback()
            raise
        else:
            transaction.commit()

        print simplejson.dumps(stats),
        return return_me

    def _enqueue_one_forked_task_for_thumbnail_fetch(self, video_id, future_status):
        if self.verbose:
            print 'Starting thumbnail fetches for', video_id

        task = self.forked_task_worker_pool.apply_async(
            function_for_fork_worker, [(video_id, future_status)])
        self.forked_tasks[video_id] = task

    def enqueue_forked_tasks_for_thumbnail_fetches(self, feed):
        # Make sure that any videos from the feed with
        # status=models.VIDEO_STATUS_PENDING_THUMBNAIL have tasks.
        all_feed_items_pending_thumbnail = feed.video_set.filter(
            status=models.VIDEO_STATUS_PENDING_THUMBNAIL)

        for video in all_feed_items_pending_thumbnail:
            if video.id in self.forked_tasks:
                continue
            # Eek. It seems we have to grab a thumbnail for a video
            # where we misplaced the forked task. In that case, we
            # have to look at the feed to see what status videos
            # should get.
            self._enqueue_one_forked_task_for_thumbnail_fetch(
                video.id, feed.default_video_status())

        if self.verbose:
            print 'Enqueued all thumbnail fetches.'

        # Finally, wait for them all to finish
        for video_id in self.forked_tasks:
            task = self.forked_tasks[video_id]
            success = task.get()
            if not success:
                # Re-enqueue it. :-(
                self._enqueue_one_forked_task_for_thumbnail_fetch(
                    video_id, feed.default_video_status())

        for video_id in self.forked_tasks:
            task = self.forked_tasks[video_id]
            success = task.get()

        if self.verbose:
            print 'Finished thumbnail fetches.'

    def use_old_bulk_import(self, parsed_feed, feed):
        bulk_feed = bulk_import(feed_url=None, parsed_feed=parsed_feed)

        stats = {
            'total': 0,
            'imported': 0,
            'skipped': 0
            }
        try:
            for i in feed._update_items_generator(verbose=self.verbose,
                                                  parsed_feed=bulk_feed,
                                                  clear_rejected=True,
                                                  actually_save_thumbnails=True):
                if not models.Feed.objects.filter(pk=feed.pk).count():
                    # someone deleted the feed, quit
                    break
                stats['total'] += 1
                if i['video'] is not None:
                    stats['imported'] += 1
                else:
                    stats['skipped'] += 1
        finally:
            feed.status = models.FEED_STATUS_ACTIVE
            feed.save()
        print simplejson.dumps(stats),
<commit_msg>Close database connection when forking thumbnail processes<commit_after># This file is part of Miro Community.
# Copyright (C) 2010 Participatory Culture Foundation
#
# Miro Community is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# Miro Community is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Miro Community.  If not, see <http://www.gnu.org/licenses/>.

import os
import simplejson
import multiprocessing
import logging

import feedparser

import eventlet
import eventlet.pools

import httplib

from django.db import transaction, close_connection
from django.core.management.base import BaseCommand, CommandError
from vidscraper.bulk_import import bulk_import_url_list, bulk_import

from localtv import models
import localtv.util
import localtv.tasks

DEFAULT_HTTPLIB_CACHE_PATH='/tmp/.cache-for-uid-%d' % os.getuid()

MAX_TRIES_FOR_BULK_IMPORT = 8

def try_a_few_times(func, count=MAX_TRIES_FOR_BULK_IMPORT):
    def wrapper(*args, **kwargs):
        for i in range(count):
            try:
                return func(*args, **kwargs)
            except Exception, e:
                logging.exception('%s(%s, %s) fell down', func, args, kwargs)
        raise e
    return wrapper

def function_for_fork_worker(data_tuple):
    video_id, future_status = data_tuple
    import localtv.management.commands.update_one_thumbnail
    cmd = localtv.management.commands.update_one_thumbnail.Command()
    try:
        cmd.handle(video_id, future_status)
    except models.Video.DoesNotExist:
        logging.warn("For some reason, we failed to find the model with ID %d" % (
                video_id, ))
        return False # Aww, shucks. Maybe after a retry this will work.
    except Exception:
        logging.exception("uh, bizarre -- the task fell over. Maybe it will work on a retry.")
        return False
    return True

class Command(BaseCommand):

    args = '[feed primary key]'

    def handle(self, *args, **options):
        if len(args) != 1:
            raise CommandError('bulk_import takes one argument: '
                               '%i argument(s) given' % len(args))
        try:
            feed = models.Feed.objects.get(pk=args[0])
        except models.Feed.DoesNotExist:
            raise CommandError('Feed with pk %s does not exist' % args[0])

        try:
            self.verbose = (int(options['verbosity']) > 1)
        except (KeyError, ValueError):
            self.verbose = False

        httplib2 = eventlet.import_patched('httplib2')

        # Use httplib2 to GET the feed.
        # This permits us to download it only once, passing a parsed_feed through
        # to the vidscraper functions.
        h = httplib2.Http(DEFAULT_HTTPLIB_CACHE_PATH)
        content = localtv.util.http_get(feed.feed_url, _httplib2=httplib2)

        parsed_feed = feedparser.parse(content)

        # Try to work asynchronously, calling feedparser ourselves. We can do that
        # if the importer supports bulk_import_url_list.
        try:
            feed_urls = bulk_import_url_list(parsed_feed=parsed_feed)
        except ValueError:
            return self.use_old_bulk_import(parsed_feed, feed)
        # Okay, good, we either got the feed_url list, or we passed the work
        # off the old-style function. Proceed.
        self.forked_tasks = {}
        # close the database connection when we start a new process; otherwise,
        # MySQL falls over because we corrupt the connection
        self.forked_task_worker_pool =  multiprocessing.Pool(processes=8,
                                                             initializer=close_connection)
        # start 8 worker processes. That should be fine, right?
        self.bulk_import_asynchronously(parsed_feed, h, feed_urls, feed)
        self.enqueue_forked_tasks_for_thumbnail_fetches(feed)

    @transaction.commit_manually
    @try_a_few_times
    def bulk_import_asynchronously(self, original_parsed_feed, h, feed_urls, feed):
        # This asynchronous bulk_import is a parallelism monster.

        # We do as much network I/O as we can using eventlet,
        # rather than threads or subprocesses.
        httplib2 = eventlet.import_patched('httplib2')

        httppool = eventlet.pools.Pool(max_size=200)
        httppool.create = lambda: httplib2.Http(DEFAULT_HTTPLIB_CACHE_PATH)

        pool = eventlet.GreenPool(100)

        def get_url(url):
            with httppool.item() as http:
                if self.verbose:
                    print 'getting', url
                try:
                    resp, content = http.request(url, 'GET')
                except httplib.BadStatusLine, e:
                    if not e.args[0]: # timeout
                        resp, content = e, ''
                    else:
                        raise
                return (resp, content)

        def cache_thumbnail_url(url):
            with httppool.item() as http:
                if self.verbose:
                    print 'getting thumb', url
                localtv.util.cache_downloaded_file(url, http)

        stats = {
            'total': 0,
            'imported': 0,
            'skipped': 0
            }

        def handle_one_sub_feed(feed_contents):
            parsed_feed = feedparser.parse(feed_contents)
            # For each feed entry in this small sub-feed, handle the item.
            for index, entry in enumerate(parsed_feed['entries'][::-1]):
                try:
                    yield handle_one_item(index, parsed_feed, entry)
                except Exception:
                    logging.warn('error handling %s', entry.get('link',
                                                                '[NO URL]'))
                    raise

        def handle_one_item(index, parsed_feed, entry):
            i = feed._handle_one_bulk_import_feed_entry(index, parsed_feed, entry, verbose=self.verbose, clear_rejected=False,
                                                        actually_save_thumbnails=False)
            # Enqueue the work to download the thumbnail
            if i['video']:
                v = i['video']
                thumbnail_url = v.thumbnail_url
                if thumbnail_url:
                    cache_thumbnail_url(thumbnail_url)

                # The _handle_one_bulk_import_feed_entry() method gave the
                # video a status, but we have to take that back for now.
                #
                # We set the status to VIDEO_STATUS_PENDING_THUMBNAIL so that
                # no one can see the video until the thumbnailing process is
                # complete.
                #
                # We pass the thumbnailer the status that the video should get
                # so that it can set that once it is ready.
                i['future_status'] = v.status
                v.status = models.VIDEO_STATUS_PENDING_THUMBNAIL
                v.save()

            stats['total'] += 1
            if i['video'] is not None:
                stats['imported'] += 1
            else:
                stats['skipped'] += 1

            return i

        results = []
        for (response, content) in pool.imap(get_url, feed_urls):
            try:
                # We make it a list so that we can iterate across
                # it more than once.
                result = list(handle_one_sub_feed(content))
            except:
                transaction.rollback()
                raise
            else:
                transaction.commit()

            try:
                results.extend(result)
                # Now that handle_one_sub_feed has finished, it is
                # safe to spawn forked tasks to do thumbnail fetching.
                for i in result:
                    video = i['video']
                    if video:
                        future_status = i['future_status']
                        self._enqueue_one_forked_task_for_thumbnail_fetch(
                            video.id, future_status)
            except:
                transaction.rollback()
                raise
            else:
                transaction.commit()

        # Get all the thumbnail URLs, and once you have them
        pool.waitall() # wait for the thumbnails

        try:
            # Finish marking the Feed as imported.
            feed._mark_bulk_import_as_done(original_parsed_feed)

            feed.status = models.FEED_STATUS_ACTIVE
            feed.save()
            return_me = [i['video'].id for i in results if i['video']]
        except:
            transaction.rollback()
            raise
        else:
            transaction.commit()

        print simplejson.dumps(stats),
        return return_me

    def _enqueue_one_forked_task_for_thumbnail_fetch(self, video_id, future_status):
        if self.verbose:
            print 'Starting thumbnail fetches for', video_id

        task = self.forked_task_worker_pool.apply_async(
            function_for_fork_worker, [(video_id, future_status)])
        self.forked_tasks[video_id] = task

    def enqueue_forked_tasks_for_thumbnail_fetches(self, feed):
        # Make sure that any videos from the feed with
        # status=models.VIDEO_STATUS_PENDING_THUMBNAIL have tasks.
        all_feed_items_pending_thumbnail = feed.video_set.filter(
            status=models.VIDEO_STATUS_PENDING_THUMBNAIL)

        for video in all_feed_items_pending_thumbnail:
            if video.id in self.forked_tasks:
                continue
            # Eek. It seems we have to grab a thumbnail for a video
            # where we misplaced the forked task. In that case, we
            # have to look at the feed to see what status videos
            # should get.
            self._enqueue_one_forked_task_for_thumbnail_fetch(
                video.id, feed.default_video_status())

        if self.verbose:
            print 'Enqueued all thumbnail fetches.'

        # Finally, wait for them all to finish
        for video_id in self.forked_tasks:
            task = self.forked_tasks[video_id]
            success = task.get()
            if not success:
                # Re-enqueue it. :-(
                self._enqueue_one_forked_task_for_thumbnail_fetch(
                    video_id, feed.default_video_status())

        for video_id in self.forked_tasks:
            task = self.forked_tasks[video_id]
            success = task.get()

        if self.verbose:
            print 'Finished thumbnail fetches.'

    def use_old_bulk_import(self, parsed_feed, feed):
        bulk_feed = bulk_import(feed_url=None, parsed_feed=parsed_feed)

        stats = {
            'total': 0,
            'imported': 0,
            'skipped': 0
            }
        try:
            for i in feed._update_items_generator(verbose=self.verbose,
                                                  parsed_feed=bulk_feed,
                                                  clear_rejected=True,
                                                  actually_save_thumbnails=True):
                if not models.Feed.objects.filter(pk=feed.pk).count():
                    # someone deleted the feed, quit
                    break
                stats['total'] += 1
                if i['video'] is not None:
                    stats['imported'] += 1
                else:
                    stats['skipped'] += 1
        finally:
            feed.status = models.FEED_STATUS_ACTIVE
            feed.save()
        print simplejson.dumps(stats),
<|endoftext|>


<commit_before>"""High level parallel SNP and indel calling using multiple variant callers.
"""
import os
import copy
import toolz as tz

from bcbio.log import logger
from bcbio import bam, utils
from bcbio.pipeline import config_utils
from bcbio.pipeline import datadict as dd
from bcbio.chipseq import macs2
from bcbio.provenance import do
from bcbio.distributed.transaction import file_transaction

def get_callers():
    from bcbio.chipseq import macs2
    return {"macs2": macs2.run}

def peakcall_prepare(data, run_parallel):
    """Entry point for doing peak calling"""
    caller_fns = get_callers()
    to_process = []
    for sample in data:
        mimic = copy.copy(sample[0])
        callers = dd.get_peakcaller(sample[0])
        if not isinstance(callers, list):
            callers = [callers]
        for caller in callers:
            if caller in caller_fns:
                mimic["peak_fn"] = caller
                name = dd.get_sample_name(mimic)
                mimic = _check(mimic, data)
                if mimic:
                    to_process.append(mimic)
                else:
                    logger.info("Skipping peak calling. No input sample for %s" % name)
    if to_process:
        after_process = run_parallel("peakcalling", to_process)
        data = _sync(data, after_process)
    return data

def calling(data):
    """Main function to parallelize peak calling."""
    chip_bam = dd.get_work_bam(data)
    input_bam = data.get("work_bam_input", None)
    caller_fn = get_callers()[data["peak_fn"]]
    name = dd.get_sample_name(data)
    out_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), data["peak_fn"], name ))
    encode_bed = tz.get_in(["genome_resources", "variation", "encode_blacklist"], data)
    # lcr_bed = utils.get_in(data, ("genome_resources", "variation", "lcr"))
    chip_bam = _prepare_bam(chip_bam, encode_bed, data['config'])
    input_bam = _prepare_bam(input_bam, encode_bed, data['config'])
    out_files = caller_fn(name, chip_bam, input_bam, dd.get_genome_build(data), out_dir,
                         dd.get_chip_method(data), data["config"])
    data.update({"peaks_files": out_files})
    return [[data]]

def _prepare_bam(bam_file, bed_file, config):
    """Remove regions from bed files"""
    if not bam_file or not bed_file:
        return bam_file
    out_file = utils.append_stem(bam_file, '_filter')
    samtools = config_utils.get_program("samtools", config)
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out:
            cmd = "{samtools} view -bh -L {bed_file} {bam_file} > {tx_out}"
            do.run(cmd.format(**locals()), "Clean %s" % bam_file)
    return out_file

def _sync(original, processed):
    """
    Add output to data if run sucessfully.
    For now only macs2 is available, so no need
    to consider multiple callers.
    """
    for original_sample in original:
        original_sample[0]["peaks_files"] = {}
        for processs_sample in processed:
            if dd.get_sample_name(original_sample[0]) == dd.get_sample_name(processs_sample[0]):
                if processs_sample[0]["peaks_files"]:
                    original_sample[0]["peaks_files"].update(processs_sample[0]["peaks_files"])
    return original

def _check(sample, data):
    """Get input sample for each chip bam file."""
    if dd.get_chip_method(sample).lower() == "atac":
        return [sample]
    if dd.get_phenotype(sample) == "input":
        return None
    for origin in data:
        if  dd.get_batch(sample) in dd.get_batch(origin[0]) and dd.get_phenotype(origin[0]) == "input":
            sample["work_bam_input"] = dd.get_work_bam(origin[0])
            return [sample]
    return [sample]

def _get_multiplier(samples):
    """Get multiplier to get jobs
       only for samples that have input
    """
    to_process = 1.0
    to_skip = 0
    for sample in samples:
        if dd.get_phenotype(sample[0]) == "chip":
            to_process += 1.0
        elif dd.get_chip_method(sample[0]).lower() == "atac":
            to_process += 1.0
        else:
            to_skip += 1.0
    mult = (to_process - to_skip) / len(samples)
    if mult <= 0:
        mult = 1 / len(samples)
    return max(mult, 1)
<commit_msg>CHIPSEQ: fix filtering of bam files<commit_after>"""High level parallel SNP and indel calling using multiple variant callers.
"""
import os
import copy
import toolz as tz

from bcbio.log import logger
from bcbio import bam, utils
from bcbio.pipeline import config_utils
from bcbio.pipeline import datadict as dd
from bcbio.chipseq import macs2
from bcbio.provenance import do
from bcbio.distributed.transaction import file_transaction

def get_callers():
    from bcbio.chipseq import macs2
    return {"macs2": macs2.run}

def peakcall_prepare(data, run_parallel):
    """Entry point for doing peak calling"""
    caller_fns = get_callers()
    to_process = []
    for sample in data:
        mimic = copy.copy(sample[0])
        callers = dd.get_peakcaller(sample[0])
        if not isinstance(callers, list):
            callers = [callers]
        for caller in callers:
            if caller in caller_fns:
                mimic["peak_fn"] = caller
                name = dd.get_sample_name(mimic)
                mimic = _check(mimic, data)
                if mimic:
                    to_process.append(mimic)
                else:
                    logger.info("Skipping peak calling. No input sample for %s" % name)
    if to_process:
        after_process = run_parallel("peakcalling", to_process)
        data = _sync(data, after_process)
    return data

def calling(data):
    """Main function to parallelize peak calling."""
    chip_bam = dd.get_work_bam(data)
    input_bam = data.get("work_bam_input", None)
    caller_fn = get_callers()[data["peak_fn"]]
    name = dd.get_sample_name(data)
    out_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), data["peak_fn"], name ))
    encode_bed = tz.get_in(["genome_resources", "variation", "encode_blacklist"], data)
    # lcr_bed = utils.get_in(data, ("genome_resources", "variation", "lcr"))
    chip_bam = _prepare_bam(chip_bam, encode_bed, data['config'])
    input_bam = _prepare_bam(input_bam, encode_bed, data['config'])
    out_files = caller_fn(name, chip_bam, input_bam, dd.get_genome_build(data), out_dir,
                         dd.get_chip_method(data), data["config"])
    data.update({"peaks_files": out_files})
    return [[data]]

def _prepare_bam(bam_file, bed_file, config):
    """Remove regions from bed files"""
    if not bam_file or not bed_file:
        return bam_file
    out_file = utils.append_stem(bam_file, '_filter')
    bedtools = config_utils.get_program("bedtools", config)
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out:
            cmd = "{bedtools} subtract -nonamecheck -A -a {bam_file} -b {bed_file} > {tx_out}"
            do.run(cmd.format(**locals()), "Clean %s" % bam_file)
    return out_file

def _sync(original, processed):
    """
    Add output to data if run sucessfully.
    For now only macs2 is available, so no need
    to consider multiple callers.
    """
    for original_sample in original:
        original_sample[0]["peaks_files"] = {}
        for processs_sample in processed:
            if dd.get_sample_name(original_sample[0]) == dd.get_sample_name(processs_sample[0]):
                if processs_sample[0]["peaks_files"]:
                    original_sample[0]["peaks_files"].update(processs_sample[0]["peaks_files"])
    return original

def _check(sample, data):
    """Get input sample for each chip bam file."""
    if dd.get_chip_method(sample).lower() == "atac":
        return [sample]
    if dd.get_phenotype(sample) == "input":
        return None
    for origin in data:
        if  dd.get_batch(sample) in dd.get_batch(origin[0]) and dd.get_phenotype(origin[0]) == "input":
            sample["work_bam_input"] = dd.get_work_bam(origin[0])
            return [sample]
    return [sample]

def _get_multiplier(samples):
    """Get multiplier to get jobs
       only for samples that have input
    """
    to_process = 1.0
    to_skip = 0
    for sample in samples:
        if dd.get_phenotype(sample[0]) == "chip":
            to_process += 1.0
        elif dd.get_chip_method(sample[0]).lower() == "atac":
            to_process += 1.0
        else:
            to_skip += 1.0
    mult = (to_process - to_skip) / len(samples)
    if mult <= 0:
        mult = 1 / len(samples)
    return max(mult, 1)
<|endoftext|>


<commit_before># This file is part of beets.
# Copyright 2013, Verrus, <github.com/Verrus/beets-plugin-featInTitle>
#
# 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.

"""Moves "featured" artists to the title from the artist field.
"""
from beets.plugins import BeetsPlugin
from beets import ui
import re

def split_on_feat(artist):
    """Given an artist string, split the "main" artist from any artist
    on the right-hand side of a string like "feat". Return the main
    artist, which is always a string, and the featuring artist, which
    may be a string or None if none is present.
    """
    parts = re.split(
        r'[fF]t\.|[fF]eaturing|[fF]eat\.|\b[wW]ith\b|&|vs\.|and',
        artist,
        1,  # Only split on the first "feat".
    )
    parts = [s.strip() for s in parts]
    if len(parts) == 1:
        return parts[0], None
    else:
        return parts

def contains_feat(title):
    """Determine whether the title contains a "featured" marker.
    """
    return bool(re.search(
        r'[fF]t\.|[fF]eaturing|[fF]eat\.|\b[wW]ith\b|&',
        title,
    ))

def update_metadata(item, feat_part):
    """Choose how to add new artists to the title and write the new
    metadata.
    """
    print item.path

    # In all cases, update the artist fields.
    item.artist = item.albumartist
    item.artist_sort, _ = split_on_feat(item.artist_sort)  # Strip featured.

    # If the title already contains a featured artist, leave it alone.
    if contains_feat(item.title):
        print u"new artist field", item.artist

    # Otherwise, add "feat. (artist)" to the title.
    else:
        # do replace title.
        print u"artist:", item.artist
        print u"title:", item.title
        print u"featured artist:", feat_part
        item.title = u"{0} feat. {1}".format(item.title, feat_part)

    item.write()

def ft_in_title(item):
    """Look for featured artists in the item's artist fields and move
    them to the title.
    """
    artist = item.artist.strip()
    albumartist = item.albumartist.strip()

    # Check whether there is a featured artist on this track and the
    # artist field does not exactly match the album artist field. In
    # that case, we attempt to move the featured artist to the title.
    _, featured = split_on_feat(artist)
    if featured and albumartist != artist:
        feat_part = None

        # Look for the album artist in the artist field. If it's not
        # present, give up.
        albumartist_split = artist.split(albumartist)
        if len(albumartist_split) <= 1:
            print 'album artist not present in artist; skipping:', item.path

        # If the last element of the split (the right-hand side of the
        # album artist) is nonempty, then it probably contains the
        # featured artist.
        elif albumartist_split[-1] != '':
            # Extract the featured artist from the right-hand side.
            _, feat_part = split_on_feat(albumartist_split[-1])

        # Otherwise, if there's nothing on the right-hand side, look for a
        # featuring artist on the left-hand side.
        else:
            lhs, rhs = split_on_feat(albumartist_split[0])
            if rhs:
                feat_part = lhs

        # If we have a featuring artist, move it to the title.
        if feat_part:
            update_metadata(item, feat_part)
        else:
            print 'found no featuring artists:', item.path

class FtInTitlePlugin(BeetsPlugin):
    def commands(self):
        cmd = ui.Subcommand('ftintitle',
                            help='move featured artists to the title field')
        def func(lib, opts, args):
            for item in lib.items():
                ft_in_title(item)
            print "Manual 'beet update' run is recommended. "
        cmd.func = func
        return [cmd]
<commit_msg>ftintitle: use beets' ui.print_ function<commit_after># This file is part of beets.
# Copyright 2013, Verrus, <github.com/Verrus/beets-plugin-featInTitle>
#
# 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.

"""Moves "featured" artists to the title from the artist field.
"""
from beets.plugins import BeetsPlugin
from beets import ui
from beets.util import displayable_path
import re

def split_on_feat(artist):
    """Given an artist string, split the "main" artist from any artist
    on the right-hand side of a string like "feat". Return the main
    artist, which is always a string, and the featuring artist, which
    may be a string or None if none is present.
    """
    parts = re.split(
        r'[fF]t\.|[fF]eaturing|[fF]eat\.|\b[wW]ith\b|&|vs\.|and',
        artist,
        1,  # Only split on the first "feat".
    )
    parts = [s.strip() for s in parts]
    if len(parts) == 1:
        return parts[0], None
    else:
        return parts

def contains_feat(title):
    """Determine whether the title contains a "featured" marker.
    """
    return bool(re.search(
        r'[fF]t\.|[fF]eaturing|[fF]eat\.|\b[wW]ith\b|&',
        title,
    ))

def update_metadata(item, feat_part):
    """Choose how to add new artists to the title and write the new
    metadata. Also, print out messages about any changes that are made.
    """
    # In all cases, update the artist fields.
    ui.print_(u'artist: {0} -> {1}'.format(item.artist, item.albumartist))
    item.artist = item.albumartist
    item.artist_sort, _ = split_on_feat(item.artist_sort)  # Strip featured.

    # Only update the title if it does not already contain a featured
    # artist.
    if not contains_feat(item.title):
        new_title = u"{0} feat. {1}".format(item.title, feat_part)
        ui.print_(u'title: {0} -> {1}'.format(item.title, new_title))
        item.title = new_title

    item.write()

def ft_in_title(item):
    """Look for featured artists in the item's artist fields and move
    them to the title.
    """
    artist = item.artist.strip()
    albumartist = item.albumartist.strip()

    # Check whether there is a featured artist on this track and the
    # artist field does not exactly match the album artist field. In
    # that case, we attempt to move the featured artist to the title.
    _, featured = split_on_feat(artist)
    if featured and albumartist != artist:
        ui.print_(displayable_path(item.path))
        feat_part = None

        # Look for the album artist in the artist field. If it's not
        # present, give up.
        albumartist_split = artist.split(albumartist)
        if len(albumartist_split) <= 1:
            ui.print_('album artist not present in artist')

        # If the last element of the split (the right-hand side of the
        # album artist) is nonempty, then it probably contains the
        # featured artist.
        elif albumartist_split[-1] != '':
            # Extract the featured artist from the right-hand side.
            _, feat_part = split_on_feat(albumartist_split[-1])

        # Otherwise, if there's nothing on the right-hand side, look for a
        # featuring artist on the left-hand side.
        else:
            lhs, rhs = split_on_feat(albumartist_split[0])
            if rhs:
                feat_part = lhs

        # If we have a featuring artist, move it to the title.
        if feat_part:
            update_metadata(item, feat_part)
        else:
            ui.print_(u'no featuring artists found')

        ui.print_()

class FtInTitlePlugin(BeetsPlugin):
    def commands(self):
        cmd = ui.Subcommand('ftintitle',
                            help='move featured artists to the title field')
        def func(lib, opts, args):
            for item in lib.items():
                ft_in_title(item)
            print "Manual 'beet update' run is recommended. "
        cmd.func = func
        return [cmd]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# pylint: disable=no-member,line-too-long

from __future__ import print_function

from builtins import str # pylint: disable=redefined-builtin

import datetime
import importlib
import json
import os
import tempfile
import traceback
import zipfile

import zipstream

import pytz

from django.conf import settings
from django.core.files import File
from django.core.mail import send_mail
from django.core.management.base import BaseCommand
from django.template.loader import render_to_string
from django.utils import timezone
from django.utils.encoding import force_text

from ...decorators import handle_lock, log_scheduled_event
from ...models import DataPoint, ReportJob, ReportJobBatchRequest, DataGeneratorDefinition, DataSourceReference, DataSource, install_supports_jsonfield

REMOVE_SLEEP_MAX = 60 # Added to avoid "WindowsError: [Error 32] The process cannot access the file because it is being used by another process"

class Command(BaseCommand):
    help = 'Compiles data reports requested by end users.'

    def add_arguments(self, parser):
        pass

    @handle_lock
    @log_scheduled_event
    def handle(self, *args, **options): # pylint: disable=too-many-locals,too-many-branches,too-many-statements
        os.umask(000)

        pending = ReportJob.objects.filter(started=None, completed=None)

        while pending.count() > 0:
            report = ReportJob.objects.filter(started=None, completed=None)\
                                      .order_by('requested', 'pk')\
                                      .first()

            if report is not None:
                report.started = timezone.now()
                report.save()

                parameters = {}

                if install_supports_jsonfield():
                    parameters = report.parameters
                else:
                    parameters = json.loads(report.parameters)

                sources = parameters['sources']
                generators = parameters['generators']

                data_start = None
                data_end = None

                tz_info = pytz.timezone(settings.TIME_ZONE)

                if 'data_start' in parameters and parameters['data_start']:
                    tokens = parameters['data_start'].split('/')

                    data_start = datetime.datetime(int(tokens[2]), \
                                                   int(tokens[0]), \
                                                   int(tokens[1]), \
                                                   0, \
                                                   0, \
                                                   0, \
                                                   0, \
                                                   tz_info)

                if 'data_end' in parameters and parameters['data_end']:
                    tokens = parameters['data_end'].split('/')

                    data_end = datetime.datetime(int(tokens[2]), \
                                                 int(tokens[0]), \
                                                 int(tokens[1]), \
                                                 23, \
                                                 59, \
                                                 59, \
                                                 999999, \
                                                 tz_info)

                date_type = 'created'

                if 'date_type' in parameters and parameters['date_type']:
                    date_type = parameters['date_type']

                raw_json = False

                if ('raw_data' in parameters) and parameters['raw_data'] is True:
                    raw_json = True

                prefix = 'pdk_export_final'

                if 'prefix' in parameters:
                    prefix = parameters['prefix']

                suffix = report.started.date().isoformat()

                if 'suffix' in parameters:
                    suffix = parameters['suffix']

                filename = tempfile.gettempdir() + os.path.sep + prefix + '_' + str(report.pk) + '_' + suffix + '.zip'

                zips_to_merge = []

                with open(filename, 'wb') as final_output_file:
                    with zipstream.ZipFile(mode='w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as export_stream: # pylint: disable=line-too-long
                        to_delete = []

                        for generator in generators: # pylint: disable=too-many-nested-blocks
                            if raw_json:
                                for source in sources:
                                    data_source = DataSource.objects.filter(identifier=source).first()

                                    if data_source is not None and data_source.server is None:
                                        generator_definition = DataGeneratorDefinition.definition_for_identifier(generator)
                                        source_reference = DataSourceReference.reference_for_source(source)

                                        points = DataPoint.objects.filter(source_reference=source_reference, generator_definition=generator_definition)

                                        if data_start is not None:
                                            if date_type == 'recorded':
                                                points = points.filter(recorded__gte=data_start)
                                            else:
                                                points = points.filter(created__gte=data_start)

                                        if data_end is not None:
                                            if date_type == 'recorded':
                                                points = points.filter(recorded__lte=data_end)
                                            else:
                                                points = points.filter(created__lte=data_end)

                                        points = points.order_by('created')

                                        first = points.first() # pylint: disable=line-too-long
                                        last = points.last() # pylint: disable=line-too-long

                                        if first is not None:
                                            first_create = first.created
                                            last_create = last.created

                                            start = datetime.datetime(first_create.year, \
                                                                      first_create.month, \
                                                                      first_create.day, \
                                                                      0, \
                                                                      0, \
                                                                      0, \
                                                                      0, \
                                                                      first_create.tzinfo)

                                            end = datetime.datetime(last_create.year, \
                                                                    last_create.month, \
                                                                    last_create.day, \
                                                                    0, \
                                                                    0, \
                                                                    0, \
                                                                    0, \
                                                                    first_create.tzinfo) + \
                                                                    datetime.timedelta(days=1)

                                            if data_start is not None and data_start > start:
                                                start = data_start

                                            if data_end is not None and data_end < end:
                                                end = data_end

                                            while start <= end:
                                                day_end = start + datetime.timedelta(days=1)

                                                day_filename = source + '__' + generator + '__' + \
                                                               start.date().isoformat() + '.json'

                                                points = DataPoint.objects.filter(source_reference=source_reference, generator_definition=generator_definition, created__gte=start, created__lt=day_end).order_by('created') # pylint: disable=line-too-long

                                                out_points = []

                                                for point in points:
                                                    out_points.append(point.fetch_properties())

                                                if out_points:
                                                    export_stream.writestr(day_filename, str(json.dumps(out_points, indent=2)).encode("utf-8")) # pylint: disable=line-too-long

                                                start = day_end
                            else:
                                output_file = None

                                for app in settings.INSTALLED_APPS:
                                    if output_file is None:
                                        try:
                                            pdk_api = importlib.import_module(app + '.pdk_api')

                                            try:
                                                output_file = pdk_api.compile_report(generator, sources, data_start=data_start, data_end=data_end, date_type=date_type)

                                                if output_file is not None:
                                                    if generator != 'pdk-personal-data':
                                                        if output_file.lower().endswith('.zip'):
                                                            zips_to_merge.append(output_file)

                                                            # with zipfile.ZipFile(output_file, 'r') as source_file:
                                                            #    for name in source_file.namelist():
                                                            #        data_file = source_file.open(name)

                                                            #        export_stream.write_iter(name, data_file, compress_type=zipfile.ZIP_DEFLATED)
                                                        else:
                                                            name = os.path.basename(os.path.normpath(output_file))

                                                            export_stream.write(output_file, name, compress_type=zipfile.ZIP_DEFLATED)

                                                            to_delete.append(output_file)
                                                    else:
                                                        name = os.path.basename(os.path.normpath(output_file))

                                                        export_stream.write(output_file, name, compress_type=zipfile.ZIP_DEFLATED)

                                                    # to_delete.append(output_file)
                                            except TypeError as exception:
                                                traceback.print_exc()
                                                print('Verify that ' + app + '.' + generator + ' implements all compile_report arguments!')
                                                raise exception
                                        except ImportError:
                                            output_file = None
                                        except AttributeError:
                                            output_file = None

                        for data in export_stream:
                            final_output_file.write(data)

                        for output_file in to_delete:
                            remove_sleep = 1.0

                            while remove_sleep < REMOVE_SLEEP_MAX:
                                try:
                                    os.remove(output_file)

                                    remove_sleep = REMOVE_SLEEP_MAX
                                except OSError:
                                    remove_sleep = remove_sleep * 2

                                    if remove_sleep >= REMOVE_SLEEP_MAX:
                                        traceback.print_exc()

                if zips_to_merge:
                    with zipfile.ZipFile(filename, 'a') as zip_output:
                        for zip_filename in zips_to_merge:
                            zip_file = zipfile.ZipFile(zip_filename, 'r')

                            for child_file in zip_file.namelist():
                                with zip_file.open(child_file) as child_stream:
                                    zip_output.writestr(child_file, child_stream.read())

                report.report.save(filename.split(os.path.sep)[-1], File(open(filename, 'rb')))
                report.completed = timezone.now()
                report.save()

                if report.requester.email is not None:
                    subject = render_to_string('pdk_report_subject.txt', {
                        'report': report,
                        'url': settings.SITE_URL
                    })

                    if 'email_subject' in parameters:
                        subject = parameters['email_subject']

                    message = render_to_string('pdk_report_message.txt', {
                        'report': report,
                        'url': settings.SITE_URL
                    })

                    tokens = settings.SITE_URL.split('/')
                    host = ''

                    while tokens and tokens[-1] == '':
                        tokens.pop()

                    if tokens:
                        host = tokens[-1]

                    send_mail(subject, \
                              message, \
                              'Petey Kay <noreply@' + host + '>', \
                              [report.requester.email], \
                              fail_silently=False)

                for extra_destination in report.requester.pdk_report_destinations.all():
                    extra_destination.transmit(filename)

                remove_sleep = 1.0

                while remove_sleep < REMOVE_SLEEP_MAX:
                    try:
                        os.remove(filename)

                        remove_sleep = REMOVE_SLEEP_MAX
                    except OSError:
                        remove_sleep = remove_sleep * 2

                        if remove_sleep >= REMOVE_SLEEP_MAX:
                            traceback.print_exc()

        request = ReportJobBatchRequest.objects.filter(started=None, completed=None)\
                      .order_by('requested', 'pk')\
                      .first()

        if request is not None:
            request.process()
<commit_msg>Removing unused import<commit_after># -*- coding: utf-8 -*-
# pylint: disable=no-member,line-too-long

from __future__ import print_function

from builtins import str # pylint: disable=redefined-builtin

import datetime
import importlib
import json
import os
import tempfile
import traceback
import zipfile

import zipstream

import pytz

from django.conf import settings
from django.core.files import File
from django.core.mail import send_mail
from django.core.management.base import BaseCommand
from django.template.loader import render_to_string
from django.utils import timezone

from ...decorators import handle_lock, log_scheduled_event
from ...models import DataPoint, ReportJob, ReportJobBatchRequest, DataGeneratorDefinition, DataSourceReference, DataSource, install_supports_jsonfield

REMOVE_SLEEP_MAX = 60 # Added to avoid "WindowsError: [Error 32] The process cannot access the file because it is being used by another process"

class Command(BaseCommand):
    help = 'Compiles data reports requested by end users.'

    def add_arguments(self, parser):
        pass

    @handle_lock
    @log_scheduled_event
    def handle(self, *args, **options): # pylint: disable=too-many-locals,too-many-branches,too-many-statements
        os.umask(000)

        pending = ReportJob.objects.filter(started=None, completed=None)

        while pending.count() > 0:
            report = ReportJob.objects.filter(started=None, completed=None)\
                                      .order_by('requested', 'pk')\
                                      .first()

            if report is not None:
                report.started = timezone.now()
                report.save()

                parameters = {}

                if install_supports_jsonfield():
                    parameters = report.parameters
                else:
                    parameters = json.loads(report.parameters)

                sources = parameters['sources']
                generators = parameters['generators']

                data_start = None
                data_end = None

                tz_info = pytz.timezone(settings.TIME_ZONE)

                if 'data_start' in parameters and parameters['data_start']:
                    tokens = parameters['data_start'].split('/')

                    data_start = datetime.datetime(int(tokens[2]), \
                                                   int(tokens[0]), \
                                                   int(tokens[1]), \
                                                   0, \
                                                   0, \
                                                   0, \
                                                   0, \
                                                   tz_info)

                if 'data_end' in parameters and parameters['data_end']:
                    tokens = parameters['data_end'].split('/')

                    data_end = datetime.datetime(int(tokens[2]), \
                                                 int(tokens[0]), \
                                                 int(tokens[1]), \
                                                 23, \
                                                 59, \
                                                 59, \
                                                 999999, \
                                                 tz_info)

                date_type = 'created'

                if 'date_type' in parameters and parameters['date_type']:
                    date_type = parameters['date_type']

                raw_json = False

                if ('raw_data' in parameters) and parameters['raw_data'] is True:
                    raw_json = True

                prefix = 'pdk_export_final'

                if 'prefix' in parameters:
                    prefix = parameters['prefix']

                suffix = report.started.date().isoformat()

                if 'suffix' in parameters:
                    suffix = parameters['suffix']

                filename = tempfile.gettempdir() + os.path.sep + prefix + '_' + str(report.pk) + '_' + suffix + '.zip'

                zips_to_merge = []

                with open(filename, 'wb') as final_output_file:
                    with zipstream.ZipFile(mode='w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as export_stream: # pylint: disable=line-too-long
                        to_delete = []

                        for generator in generators: # pylint: disable=too-many-nested-blocks
                            if raw_json:
                                for source in sources:
                                    data_source = DataSource.objects.filter(identifier=source).first()

                                    if data_source is not None and data_source.server is None:
                                        generator_definition = DataGeneratorDefinition.definition_for_identifier(generator)
                                        source_reference = DataSourceReference.reference_for_source(source)

                                        points = DataPoint.objects.filter(source_reference=source_reference, generator_definition=generator_definition)

                                        if data_start is not None:
                                            if date_type == 'recorded':
                                                points = points.filter(recorded__gte=data_start)
                                            else:
                                                points = points.filter(created__gte=data_start)

                                        if data_end is not None:
                                            if date_type == 'recorded':
                                                points = points.filter(recorded__lte=data_end)
                                            else:
                                                points = points.filter(created__lte=data_end)

                                        points = points.order_by('created')

                                        first = points.first() # pylint: disable=line-too-long
                                        last = points.last() # pylint: disable=line-too-long

                                        if first is not None:
                                            first_create = first.created
                                            last_create = last.created

                                            start = datetime.datetime(first_create.year, \
                                                                      first_create.month, \
                                                                      first_create.day, \
                                                                      0, \
                                                                      0, \
                                                                      0, \
                                                                      0, \
                                                                      first_create.tzinfo)

                                            end = datetime.datetime(last_create.year, \
                                                                    last_create.month, \
                                                                    last_create.day, \
                                                                    0, \
                                                                    0, \
                                                                    0, \
                                                                    0, \
                                                                    first_create.tzinfo) + \
                                                                    datetime.timedelta(days=1)

                                            if data_start is not None and data_start > start:
                                                start = data_start

                                            if data_end is not None and data_end < end:
                                                end = data_end

                                            while start <= end:
                                                day_end = start + datetime.timedelta(days=1)

                                                day_filename = source + '__' + generator + '__' + \
                                                               start.date().isoformat() + '.json'

                                                points = DataPoint.objects.filter(source_reference=source_reference, generator_definition=generator_definition, created__gte=start, created__lt=day_end).order_by('created') # pylint: disable=line-too-long

                                                out_points = []

                                                for point in points:
                                                    out_points.append(point.fetch_properties())

                                                if out_points:
                                                    export_stream.writestr(day_filename, str(json.dumps(out_points, indent=2)).encode("utf-8")) # pylint: disable=line-too-long

                                                start = day_end
                            else:
                                output_file = None

                                for app in settings.INSTALLED_APPS:
                                    if output_file is None:
                                        try:
                                            pdk_api = importlib.import_module(app + '.pdk_api')

                                            try:
                                                output_file = pdk_api.compile_report(generator, sources, data_start=data_start, data_end=data_end, date_type=date_type)

                                                if output_file is not None:
                                                    if generator != 'pdk-personal-data':
                                                        if output_file.lower().endswith('.zip'):
                                                            zips_to_merge.append(output_file)

                                                            # with zipfile.ZipFile(output_file, 'r') as source_file:
                                                            #    for name in source_file.namelist():
                                                            #        data_file = source_file.open(name)

                                                            #        export_stream.write_iter(name, data_file, compress_type=zipfile.ZIP_DEFLATED)
                                                        else:
                                                            name = os.path.basename(os.path.normpath(output_file))

                                                            export_stream.write(output_file, name, compress_type=zipfile.ZIP_DEFLATED)

                                                            to_delete.append(output_file)
                                                    else:
                                                        name = os.path.basename(os.path.normpath(output_file))

                                                        export_stream.write(output_file, name, compress_type=zipfile.ZIP_DEFLATED)

                                                    # to_delete.append(output_file)
                                            except TypeError as exception:
                                                traceback.print_exc()
                                                print('Verify that ' + app + '.' + generator + ' implements all compile_report arguments!')
                                                raise exception
                                        except ImportError:
                                            output_file = None
                                        except AttributeError:
                                            output_file = None

                        for data in export_stream:
                            final_output_file.write(data)

                        for output_file in to_delete:
                            remove_sleep = 1.0

                            while remove_sleep < REMOVE_SLEEP_MAX:
                                try:
                                    os.remove(output_file)

                                    remove_sleep = REMOVE_SLEEP_MAX
                                except OSError:
                                    remove_sleep = remove_sleep * 2

                                    if remove_sleep >= REMOVE_SLEEP_MAX:
                                        traceback.print_exc()

                if zips_to_merge:
                    with zipfile.ZipFile(filename, 'a') as zip_output:
                        for zip_filename in zips_to_merge:
                            zip_file = zipfile.ZipFile(zip_filename, 'r')

                            for child_file in zip_file.namelist():
                                with zip_file.open(child_file) as child_stream:
                                    zip_output.writestr(child_file, child_stream.read())

                report.report.save(filename.split(os.path.sep)[-1], File(open(filename, 'rb')))
                report.completed = timezone.now()
                report.save()

                if report.requester.email is not None:
                    subject = render_to_string('pdk_report_subject.txt', {
                        'report': report,
                        'url': settings.SITE_URL
                    })

                    if 'email_subject' in parameters:
                        subject = parameters['email_subject']

                    message = render_to_string('pdk_report_message.txt', {
                        'report': report,
                        'url': settings.SITE_URL
                    })

                    tokens = settings.SITE_URL.split('/')
                    host = ''

                    while tokens and tokens[-1] == '':
                        tokens.pop()

                    if tokens:
                        host = tokens[-1]

                    send_mail(subject, \
                              message, \
                              'Petey Kay <noreply@' + host + '>', \
                              [report.requester.email], \
                              fail_silently=False)

                for extra_destination in report.requester.pdk_report_destinations.all():
                    extra_destination.transmit(filename)

                remove_sleep = 1.0

                while remove_sleep < REMOVE_SLEEP_MAX:
                    try:
                        os.remove(filename)

                        remove_sleep = REMOVE_SLEEP_MAX
                    except OSError:
                        remove_sleep = remove_sleep * 2

                        if remove_sleep >= REMOVE_SLEEP_MAX:
                            traceback.print_exc()

        request = ReportJobBatchRequest.objects.filter(started=None, completed=None)\
                      .order_by('requested', 'pk')\
                      .first()

        if request is not None:
            request.process()
<|endoftext|>


<commit_before>'''
Created on 10 dec. 2013

@author: TuRz4m
'''
import ConfigParser
import logging
import sys
from api.APIUtils import BetaSerieAPI, BadLoginException, SickBeardAPI
import os.path

logger = logging.getLogger(__name__)
logging.getLogger(__name__).setLevel(logging.INFO)
logging.getLogger(__name__).addHandler(logging.StreamHandler())
logging.getLogger(__name__).addHandler(logging.FileHandler("logs/BetaBeard.log"))

configFile = "BetaBeard.ini"
configDbFile = "BetaBeard.db"
param = {}
paramDb = {}

"""
Load the config file & fill all the var.
"""
def checkConfig(config):
    try:
        global param

        param['login'] = config.get("BetaSeries", "login")
        param['password'] = config.get("BetaSeries", "password")

        param['sburl'] = config.get("SickBeard", "url")
        if (config.getboolean("SickBeard", "https")):
            param['scheme'] = "https"

        param['apikey'] = config.get("SickBeard", "apikey")

        param['location'] = config.get("SickBeard", "location")
        if (param['location'] == ""):
            param['location'] = None
        param['lang'] = config.get("SickBeard", "lang")
        if (param['lang'] == ""):
            param['lang'] = None
        param['flatten_folder'] = config.get("SickBeard", "flatten_folder")
        if (param['flatten_folder'] == ""):
            param['flatten_folder'] = None
        param['status'] = config.get("SickBeard", "status")
        if (param['status'] == ""):
            param['status'] = None
        param['initial'] = config.get("SickBeard", "initial")
        if (param['initial'] == ""):
            param['initial'] = None
        param['archive'] = config.get("SickBeard", "archive")
        if (param['archive'] == ""):
            param['archive'] = None

        param['fullUpdate'] = config.getboolean("BetaBeard", "fullUpdate")
        param['checkTimeLine'] = config.getboolean("BetaBeard", "checkTimeLine")
        param['demoMode'] = config.getboolean("BetaBeard", "demoMode")

    except ConfigParser.NoOptionError as ex:
        logger.error("[BetaBeard] Error in config file : %s", ex)
        return False
    except ConfigParser.NoSectionError as ex:
        logger.error("[BetaBeard] Error in config file : %s", ex)
        return False

    return True

def loadDb(configToLoad):
    global paramDb

    if (os.path.exists(configDbFile)):
        configToLoad.read(configDbFile)

    try:
        paramDb['last_event_id'] = configToLoad.get("BetaBeard", "last_event_id")
        if (paramDb['last_event_id'] == ""):
            paramDb['last_event_id'] = None
    except ConfigParser.NoOptionError:
        logger.debug("[BetaBeard] Config file Tech not found. Use default.")
        paramDb['last_event_id'] = None
    except ConfigParser.NoSectionError:
        logger.debug("[BetaBeard] Config file Tech not found. Use default.")
        configToLoad.add_section("BetaBeard")
        paramDb['last_event_id'] = None

"""
Update the BetaBeard-tech.ini
"""
def updateDb(configToSave):
    logger.debug("[BetaBeard] Update file %s", configDbFile)
    cfgfile = open(configDbFile,'w')
    configToSave.write(cfgfile)
    cfgfile.close()
    logger.debug("[BetaBeard] File %s updated.", configDbFile)

if __name__ == '__main__':
    # First of all, we need to reed the BetaBeard.ini config file.
    config = ConfigParser.SafeConfigParser()
    configDb = ConfigParser.SafeConfigParser()

    if (os.path.exists(configFile) == False):
        logger.error("[BetaBeard] Config file %s not found.", configFile)
        sys.exit(0)

    config.read(configFile)
    loadDb(configDb)

    if checkConfig(config) == False:
        sys.exit(0)

    # ----------- Init BetaSeries ----------- #
    try:
        beta = BetaSerieAPI(param['login'], param['password'])
    except BadLoginException as ex:
        logger.error("[BetaBeard] can't log into BetaSeries.com : %s", ex.value['text'])
        sys.exit(0)

    logger.info("[BetaBeard] Login successfull.")
    # ----------- Init SickBeard ----------- #

    sickBeard = SickBeardAPI(param['sburl'], param['scheme'], param['apikey'])

    # ----------- Test SickBeard ----------- #
    if (sickBeard.ping() == False):
        logger.error("[BetaBeard] Can't ping SickBeard on url : %s://%s with apikey = %s",param['scheme'], param['sburl'], param['apikey'])
        sys.exit(0)
    logger.info("[BetaBeard] Ping SickBeard successfull.")

    # ----------- If fullUpdate, we retrieve all the current show and add them to sickbear.----------- #
    if paramDb['last_event_id'] == None:
        logger.debug("[BetaBeard] last_index_id is None")
        if param['fullUpdate'] == True:
            shows = beta.show_list();
            logger.debug("[BetaBeard] shows : %s", shows)
            logger.info("[BetaBeard] Start processing shows.")
            for show in shows:
                logger.info("[BetaBeard] Add show in SickBeard : %s (%s)", show[1], show[0])
                if (param['demoMode'] == False):
                    success = sickBeard.add_show(show[0], param['location'],  param['lang'],  param['flatten_folder'],  param['status'],  param['initial'],  param['archive'])
                    if (success == False):
                        logger.error("[BetaBeard] Can't add show %s (%s) to sickbeard.", show[1], show[0])

        # ----------- retrieve last  event processed in betaseries----------- #
        param['last_event_id'], emptyList = beta.timeline_since(None)
    elif param['checkTimeLine']:
        logger.info("[BetaBeard] Start processing timeline.")
        param['last_event_id'], events = beta.timeline_since(paramDb['last_event_id'])
        logger.debug("[BetaBeard] Processing timeline : %s", events)
        if (events != None):
            for event in events:
                logger.debug("[BetaBeard] Event : %s", event)

                if (event['type'] == 'add_serie'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Add Show to sickbeard : %s (%s)", title, tvdbid)
                elif (event['type'] == 'del_serie'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Delete Show from sickbeard :  %s (%s)", title, tvdbid)
                elif (event['type'] == 'archive'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Archive Show on sickbeard : %s (%s)", title, tvdbid)
                elif (event['type'] == 'unarchive'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] UnArchive Show on sickbeard :  %s (%s)", title, tvdbid)

    # ----------- Update Last_event_id in config file.----------- #
    if (param['last_event_id'] != None):
        logger.debug("[BetaBeard] update config with last_event_id=%s", param['last_event_id'])
        configDb.set("BetaBeard", "last_event_id", str(param['last_event_id']));
        updateDb(configDb);
    else:
        logger.debug("[BetaBeard] Can't update config file because last_event_id is null")

<commit_msg>Add DemoMode on timeline process.<commit_after>'''
Created on 10 dec. 2013

@author: TuRz4m
'''
import ConfigParser
import logging
import sys
from api.APIUtils import BetaSerieAPI, BadLoginException, SickBeardAPI
import os.path

logger = logging.getLogger(__name__)
logging.getLogger(__name__).setLevel(logging.INFO)
logging.getLogger(__name__).addHandler(logging.StreamHandler())
logging.getLogger(__name__).addHandler(logging.FileHandler("logs/BetaBeard.log"))

configFile = "BetaBeard.ini"
configDbFile = "BetaBeard.db"
param = {}
paramDb = {}

"""
Load the config file & fill all the var.
"""
def checkConfig(config):
    try:
        global param

        param['login'] = config.get("BetaSeries", "login")
        param['password'] = config.get("BetaSeries", "password")

        param['sburl'] = config.get("SickBeard", "url")
        if (config.getboolean("SickBeard", "https")):
            param['scheme'] = "https"

        param['apikey'] = config.get("SickBeard", "apikey")

        param['location'] = config.get("SickBeard", "location")
        if (param['location'] == ""):
            param['location'] = None
        param['lang'] = config.get("SickBeard", "lang")
        if (param['lang'] == ""):
            param['lang'] = None
        param['flatten_folder'] = config.get("SickBeard", "flatten_folder")
        if (param['flatten_folder'] == ""):
            param['flatten_folder'] = None
        param['status'] = config.get("SickBeard", "status")
        if (param['status'] == ""):
            param['status'] = None
        param['initial'] = config.get("SickBeard", "initial")
        if (param['initial'] == ""):
            param['initial'] = None
        param['archive'] = config.get("SickBeard", "archive")
        if (param['archive'] == ""):
            param['archive'] = None

        param['fullUpdate'] = config.getboolean("BetaBeard", "fullUpdate")
        param['checkTimeLine'] = config.getboolean("BetaBeard", "checkTimeLine")
        param['demoMode'] = config.getboolean("BetaBeard", "demoMode")

    except ConfigParser.NoOptionError as ex:
        logger.error("[BetaBeard] Error in config file : %s", ex)
        return False
    except ConfigParser.NoSectionError as ex:
        logger.error("[BetaBeard] Error in config file : %s", ex)
        return False

    return True

def loadDb(configToLoad):
    global paramDb

    if (os.path.exists(configDbFile)):
        configToLoad.read(configDbFile)

    try:
        paramDb['last_event_id'] = configToLoad.get("BetaBeard", "last_event_id")
        if (paramDb['last_event_id'] == ""):
            paramDb['last_event_id'] = None
    except ConfigParser.NoOptionError:
        logger.debug("[BetaBeard] Config file Tech not found. Use default.")
        paramDb['last_event_id'] = None
    except ConfigParser.NoSectionError:
        logger.debug("[BetaBeard] Config file Tech not found. Use default.")
        configToLoad.add_section("BetaBeard")
        paramDb['last_event_id'] = None

"""
Update the BetaBeard-tech.ini
"""
def updateDb(configToSave):
    logger.debug("[BetaBeard] Update file %s", configDbFile)
    cfgfile = open(configDbFile,'w')
    configToSave.write(cfgfile)
    cfgfile.close()
    logger.debug("[BetaBeard] File %s updated.", configDbFile)

if __name__ == '__main__':
    # First of all, we need to reed the BetaBeard.ini config file.
    config = ConfigParser.SafeConfigParser()
    configDb = ConfigParser.SafeConfigParser()

    if (os.path.exists(configFile) == False):
        logger.error("[BetaBeard] Config file %s not found.", configFile)
        sys.exit(0)

    config.read(configFile)
    loadDb(configDb)

    if checkConfig(config) == False:
        sys.exit(0)

    # ----------- Init BetaSeries ----------- #
    try:
        beta = BetaSerieAPI(param['login'], param['password'])
    except BadLoginException as ex:
        logger.error("[BetaBeard] can't log into BetaSeries.com : %s", ex.value['text'])
        sys.exit(0)

    logger.info("[BetaBeard] Login successfull.")
    # ----------- Init SickBeard ----------- #

    sickBeard = SickBeardAPI(param['sburl'], param['scheme'], param['apikey'])

    # ----------- Test SickBeard ----------- #
    if (sickBeard.ping() == False):
        logger.error("[BetaBeard] Can't ping SickBeard on url : %s://%s with apikey = %s",param['scheme'], param['sburl'], param['apikey'])
        sys.exit(0)
    logger.info("[BetaBeard] Ping SickBeard successfull.")

    # ----------- If fullUpdate, we retrieve all the current show and add them to sickbear.----------- #
    if paramDb['last_event_id'] == None:
        logger.debug("[BetaBeard] last_index_id is None")
        if param['fullUpdate'] == True:
            shows = beta.show_list();
            logger.debug("[BetaBeard] shows : %s", shows)
            logger.info("[BetaBeard] Start processing shows.")
            for show in shows:
                logger.info("[BetaBeard] Add show in SickBeard : %s (%s)", show[1], show[0])
                if (param['demoMode'] == False):
                    success = sickBeard.add_show(show[0], param['location'],  param['lang'],  param['flatten_folder'],  param['status'],  param['initial'],  param['archive'])
                    if (success == False):
                        logger.error("[BetaBeard] Can't add show %s (%s) to sickbeard.", show[1], show[0])

        # ----------- retrieve last  event processed in betaseries----------- #
        param['last_event_id'], emptyList = beta.timeline_since(None)
    elif param['checkTimeLine']:
        logger.info("[BetaBeard] Start processing timeline.")
        param['last_event_id'], events = beta.timeline_since(paramDb['last_event_id'])
        logger.debug("[BetaBeard] Processing timeline : %s", events)
        if (events != None):
            for event in events:
                logger.debug("[BetaBeard] Event : %s", event)

                # - ADD SERIE - #
                if (event['type'] == 'add_serie'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Add Show to sickbeard : %s (%s)", title, tvdbid)

                    if (param['demoMode'] == False):
                        success = sickBeard.add_show(tvdbid, param['location'],  param['lang'],  param['flatten_folder'],  param['status'],  param['initial'],  param['archive'])
                        if (success == False):
                            logger.error("[BetaBeard] Can't add show %s (%s) to sickbeard.", title, tvdbid)

                # - DELETE SERIE - #
                elif (event['type'] == 'del_serie'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Delete Show from sickbeard :  %s (%s)", title, tvdbid)

                    if (param['demoMode'] == False):
                        success =  sickBeard.del_show(tvdbid)
                        if (success == False):
                            logger.error("[BetaBeard] Can't delete show %s (%s) from sickbeard.", title, tvdbid)

                # - PAUSE SERIE - #
                elif (event['type'] == 'archive'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] Archive Show on sickbeard : %s (%s)", title, tvdbid)

                    if (param['demoMode'] == False):
                        success =  sickBeard.pause_show(tvdbid, 1)
                        if (success == False):
                            logger.error("[BetaBeard] Can't pause show %s (%s) on sickbeard.", title, tvdbid)

                # - UNPAUSE SERIE - #
                elif (event['type'] == 'unarchive'):
                    betaid = str(event['ref_id']);
                    tvdbid, title = beta.shows_tvdbid(betaid)
                    logger.info("[BetaBeard] UnArchive Show on sickbeard :  %s (%s)", title, tvdbid)

                    if (param['demoMode'] == False):
                        success =  sickBeard.pause_show(tvdbid, 0)
                        if (success == False):
                            logger.error("[BetaBeard] Can't unpause show %s (%s) on sickbeard.", title, tvdbid)

        logger.info("[BetaBeard] Timeline processing done.")

    # ----------- Update Last_event_id in config file.----------- #
    if (param['last_event_id'] != None):
        logger.debug("[BetaBeard] update config with last_event_id=%s", param['last_event_id'])
        configDb.set("BetaBeard", "last_event_id", str(param['last_event_id']));
        updateDb(configDb);
    else:
        logger.debug("[BetaBeard] Can't update config file because last_event_id is null")

<|endoftext|>


<commit_before>import os
import json

'''
data_dir = os.path.join(
    os.path.split(
        os.path.split(os.path.dirname(__file__))[0])[0], 'data')
'''
data_dir = os.path.join(os.environ['PORTER'], 'data')

def read_and_parse_json(data_type):
    sub_dir = os.path.join(data_dir, data_type)
    elements = []

    def full_path(file_name):
        return os.path.join(sub_dir, file_name)

    def only_json(file_name):
        return file_name.endswith('.json')

    for json_file_name in filter(only_json, map(full_path, os.listdir(sub_dir))):
        with open(json_file_name) as json_file:
            elements.append(json.load(json_file))
    return elements
<commit_msg>Remove commented file loading code<commit_after>import os
import json

data_dir = os.path.join(os.environ['PORTER'], 'data')

def read_and_parse_json(data_type):
    sub_dir = os.path.join(data_dir, data_type)
    elements = []

    def full_path(file_name):
        return os.path.join(sub_dir, file_name)

    def only_json(file_name):
        return file_name.endswith('.json')

    for json_file_name in filter(only_json, map(full_path, os.listdir(sub_dir))):
        with open(json_file_name) as json_file:
            elements.append(json.load(json_file))
    return elements
<|endoftext|>


<commit_before>from __future__ import absolute_import

import itertools
import os
import tempfile

from celery import shared_task
from PIL import Image as PILImage
from PIL import JpegImagePlugin

from betty.conf.app import settings
from betty.cropper.models import Image

COLOR_DENSITY_RATIO = 0.11

def get_color_density(im):
    area = im.size[0] * im.size[1]
    unique_colors = len(filter(None, im.histogram()))
    return unique_colors / float(area)

def get_error(a, b):
    assert a.size == b.size
    difference = 0
    for color_sets in itertools.izip(a.getdata(), b.getdata()):
        distance = 0
        for color_pair in zip(color_sets[0], color_sets[1]):
            distance += ((color_pair[0] - color_pair[1]) ** 2)
        difference += (distance ** 0.5)

    pixel_error = difference / float(b.size[0] * b.size[1])
    return pixel_error

@shared_task
def search_image_quality(image_id):

    image = Image.objects.get(id=image_id)

    im = PILImage.open(image.optimized.path)
    search_im = im.copy()

    area = search_im.size[0] * search_im.size[1]
    max_area = (1000.0 * 1000.0)
    if area > max_area:
        scale = max_area / area
        new_size = (search_im.size[0] * scale, search_im.size[1] * scale)
        search_im = search_im.resize(map(int, new_size), PILImage.ANTIALIAS)
    if im.mode != "RGB":
        im = im.convert("RGB", palette=PILImage.ADAPTIVE)

    original_density = get_color_density(search_im)
    icc_profile = im.info.get("icc_profile")

    search_range = settings.BETTY_JPEG_QUALITY_RANGE

    # First, let's check to make sure that this image isn't already an optimized JPEG
    if im.format == "JPEG":
        original_path = tempfile.mkstemp()[1]
        search_im.save(
            original_path,
            "JPEG",
            quantization=im.quantization,
            subsampling=JpegImagePlugin.get_sampling(im),
            icc_profile=icc_profile)
        optimized_path = tempfile.mkstemp()[1]
        search_im.save(
            optimized_path,
            "JPEG",
            quality=settings.BETTY_DEFAULT_JPEG_QUALITY,
            icc_profile=icc_profile,
            optimize=True)
        if os.stat(original_path).st_size < os.stat(optimized_path).st_size:
            # Looks like the original was already compressed, let's bail.
            return

    while (search_range[1] - search_range[0]) > 1:
        quality = int(round(search_range[0] + (search_range[1] - search_range[0]) / 2.0))
        print("Searching: {}".format(quality))

        output_filepath = tempfile.mkstemp()[1]
        search_im.save(output_filepath, "jpeg", quality=quality, icc_profile=icc_profile, optimize=True)
        saved = PILImage.open(output_filepath)

        pixel_error = get_error(saved, search_im)
        density_ratio = (get_color_density(saved) - original_density) / original_density

        if pixel_error > settings.BETTY_JPEG_MAX_ERROR or density_ratio > COLOR_DENSITY_RATIO:
            search_range = (quality, search_range[1])
        else:
            search_range = (search_range[0], quality)
        os.remove(output_filepath)

    image.jpeg_quality = search_range[1]
    image.save()
<commit_msg>Fixing up the image conversion stuff<commit_after>from __future__ import absolute_import

import itertools
import os
import tempfile

from celery import shared_task
from PIL import Image as PILImage
from PIL import JpegImagePlugin

from betty.conf.app import settings
from betty.cropper.models import Image

COLOR_DENSITY_RATIO = 0.11

def get_color_density(im):
    area = im.size[0] * im.size[1]
    unique_colors = len(filter(None, im.histogram()))
    return unique_colors / float(area)

def get_error(a, b):
    assert a.size == b.size
    difference = 0
    for color_sets in itertools.izip(a.getdata(), b.getdata()):
        distance = 0
        for color_pair in zip(color_sets[0], color_sets[1]):
            distance += ((color_pair[0] - color_pair[1]) ** 2)
        difference += (distance ** 0.5)

    pixel_error = difference / float(b.size[0] * b.size[1])
    return pixel_error

@shared_task
def search_image_quality(image_id):

    image = Image.objects.get(id=image_id)

    im = PILImage.open(image.optimized.path)
    search_im = im.copy()

    area = search_im.size[0] * search_im.size[1]
    max_area = (1000.0 * 1000.0)
    if area > max_area:
        scale = max_area / area
        new_size = (search_im.size[0] * scale, search_im.size[1] * scale)
        search_im = search_im.resize(map(int, new_size), PILImage.ANTIALIAS)
    if search_im.mode != "RGB":
        search_im = search_im.convert("RGB", palette=PILImage.ADAPTIVE)

    original_density = get_color_density(search_im)
    icc_profile = im.info.get("icc_profile")

    search_range = settings.BETTY_JPEG_QUALITY_RANGE

    # First, let's check to make sure that this image isn't already an optimized JPEG
    if im.format == "JPEG":
        original_path = tempfile.mkstemp()[1]
        search_im.save(
            original_path,
            "JPEG",
            quantization=im.quantization,
            subsampling=JpegImagePlugin.get_sampling(im),
            icc_profile=icc_profile)
        optimized_path = tempfile.mkstemp()[1]
        search_im.save(
            optimized_path,
            "JPEG",
            quality=settings.BETTY_DEFAULT_JPEG_QUALITY,
            icc_profile=icc_profile,
            optimize=True)
        if os.stat(original_path).st_size < os.stat(optimized_path).st_size:
            # Looks like the original was already compressed, let's bail.
            return

    while (search_range[1] - search_range[0]) > 1:
        quality = int(round(search_range[0] + (search_range[1] - search_range[0]) / 2.0))

        output_filepath = tempfile.mkstemp()[1]
        search_im.save(output_filepath, "jpeg", quality=quality, icc_profile=icc_profile, optimize=True)
        saved = PILImage.open(output_filepath)

        pixel_error = get_error(saved, search_im)
        density_ratio = (get_color_density(saved) - original_density) / original_density

        if pixel_error > settings.BETTY_JPEG_MAX_ERROR or density_ratio > COLOR_DENSITY_RATIO:
            search_range = (quality, search_range[1])
        else:
            search_range = (search_range[0], quality)
        os.remove(output_filepath)

    image.jpeg_quality = search_range[1]
    image.save()
<|endoftext|>