<commit_before>import os
from urllib.parse import urljoin
from zipfile import ZipFile

import click
import requests
from lxml import html

class Comic:
    def __init__(self, start_url, next_page_selector, comic_image_selector):
        self.url = start_url
        self.next_page_selector = next_page_selector
        self.comic_image_selector = comic_image_selector
        self.current_page = 1

    def download(self, directory_name="finalComic"):
        """
        Downloads an entire Webcomic page by page starting from the first one
        and saves them in the directory_name created in the current working directory
        """
        os.makedirs(directory_name)
        while True:
            click.echo("Downloading page {}".format(self.url))
            response = requests.get(self.url)
            parsed_html = html.fromstring(response.content)

            image_element = parsed_html.xpath(self.comic_image_selector)
            next_link = parsed_html.xpath(self.next_page_selector)

            if image_element == []:
                click.echo("Could not find comic image.")
            else:
                try:
                    image_url = urljoin(self.url, image_element[0])
                    self.save_image(image_url, directory_name)
                except:
                    click.echo("The image couldn't be downloaded.")

            self.current_page += 1
            if next_link == [] or next_link[0].endswith("#"):
                break
            self.url = urljoin(self.url, next_link[0])
        click.echo("Finished downloading the images.")

    def save_image(self, image_url, directory_name):
        """
        Gets the image from the image_url and saves it in the directory_name
        """
        click.echo("Saving image {}".format(image_url))
        res = requests.get(image_url)
        res.raise_for_status()
        image_path = self.save_image_location(image_url, directory_name)
        if os.path.isfile(image_path):
            click.echo("The image was already downloaded. Skipping...")
        else:
            # Save the image
            with open(image_path, 'wb') as image_file:
                image_file.write(res.content)

    def save_image_location(self, url, directory):
        """
        Returns the location in the filesystem under which the webcomic will be saved
        """
        if url.count(".") <= 1:
            # No file extension (only dot in url is domain name)
            file_name = str(self.current_page)
        else:
            file_name = "{}{}".format(self.current_page, url[url.rindex("."):])
        return "/".join([directory, file_name])

    @staticmethod
    def make_cbz(comic_name, source_directory="finalComic"):
        """
        Takes all of the previously downloaded pages and compresses them in
        a .cbz file, erasing them afterwards.
        """
        cbz_file = ZipFile("{}.cbz".format(comic_name), mode="w")
        images = os.listdir("{}".format(source_directory))
        for image in images:
            image_location = "{}/{}".format(source_directory, image)
            cbz_file.write(image_location)
            os.remove(image_location)
        os.rmdir(source_directory)
        if cbz_file.testzip() != None:
            click.echo("Error while testing the archive; it might be corrupted.")
            cbz_file.close()
        else:
            cbz_file.close()
    @staticmethod
    def verify_xpath(url, next_page, image):
        """
        Takes a url and the XPath expressions for the next_page and image to go three pages
        into the comic. It returns a tuple containing the url of each page and their respective
        image urls.
        """
        verification = []
        for _ in range(3):
            response = requests.get(url)
            parsed_html = html.fromstring(response.content)

            image_element = parsed_html.xpath(image)[0]
            image_url = urljoin(url, image_element)
            next_link = parsed_html.xpath(next_page)[0]
            verification.append((url, image_url))
            url = urljoin(url, next_link)
        return verification
<commit_msg>Make code a bit more concise<commit_after>import os
from urllib.parse import urljoin
from zipfile import ZipFile

import click
import requests
from lxml import html

class Comic:
    def __init__(self, start_url, next_page_selector, comic_image_selector):
        self.url = start_url
        self.next_page_selector = next_page_selector
        self.comic_image_selector = comic_image_selector
        self.current_page = 1

    def download(self, directory_name="finalComic"):
        """
        Downloads an entire Webcomic page by page starting from the first one
        and saves them in the directory_name created in the current working directory
        """
        os.makedirs(directory_name)
        while True:
            click.echo("Downloading page {}".format(self.url))
            response = requests.get(self.url)
            parsed_html = html.fromstring(response.content)

            image_element = parsed_html.xpath(self.comic_image_selector)
            next_link = parsed_html.xpath(self.next_page_selector)

            if image_element == []:
                click.echo("Could not find comic image.")
            else:
                try:
                    image_url = urljoin(self.url, image_element[0])
                    self.save_image(image_url, directory_name)
                except:
                    click.echo("The image couldn't be downloaded.")

            self.current_page += 1
            if next_link == [] or next_link[0].endswith("#"):
                break
            self.url = urljoin(self.url, next_link[0])
        click.echo("Finished downloading the images.")

    def save_image(self, image_url, directory_name):
        """
        Gets the image from the image_url and saves it in the directory_name
        """
        click.echo("Saving image {}".format(image_url))
        res = requests.get(image_url)
        res.raise_for_status()
        image_path = self.save_image_location(image_url, directory_name)
        if os.path.isfile(image_path):
            click.echo("The image was already downloaded. Skipping...")
        else:
            with open(image_path, 'wb') as image_file:
                image_file.write(res.content)

    def save_image_location(self, url, directory):
        """
        Returns the location in the filesystem under which the webcomic will be saved
        """
        if url.count(".") <= 1:
            # No file extension (only dot in url is domain name)
            file_name = str(self.current_page)
        else:
            file_name = "{}{}".format(self.current_page, url[url.rindex("."):])
        return "/".join([directory, file_name])

    @staticmethod
    def make_cbz(comic_name, source_directory="finalComic"):
        """
        Takes all of the previously downloaded pages and compresses them in
        a .cbz file, erasing them afterwards.
        """
        cbz_file = ZipFile("{}.cbz".format(comic_name), mode="w")
        images = os.listdir("{}".format(source_directory))
        for image in images:
            image_location = "{}/{}".format(source_directory, image)
            cbz_file.write(image_location)
            os.remove(image_location)
        os.rmdir(source_directory)
        if cbz_file.testzip() != None:
            click.echo("Error while testing the archive; it might be corrupted.")
        cbz_file.close()

    @staticmethod
    def verify_xpath(url, next_page, image):
        """
        Takes a url and the XPath expressions for the next_page and image to go three pages
        into the comic. It returns a tuple containing the url of each page and their respective
        image urls.
        """
        verification = []
        for _ in range(3):
            response = requests.get(url)
            parsed_html = html.fromstring(response.content)

            image_element = parsed_html.xpath(image)[0]
            image_url = urljoin(url, image_element)
            next_link = parsed_html.xpath(next_page)[0]
            verification.append((url, image_url))
            url = urljoin(url, next_link)
        return verification
<|endoftext|>


<commit_before>import datetime

from icalendar import Calendar, Event

from django.db.models import Q
from django.views.generic import DetailView, TemplateView, View
from django.contrib.sites.shortcuts import get_current_site
from django.http import HttpResponse

from rest_framework import viewsets
from rest_framework.permissions import IsAdminUser
from wafer.pages.models import Page
from wafer.schedule.models import Venue, Slot, Day
from wafer.schedule.admin import check_schedule, validate_schedule
from wafer.schedule.models import ScheduleItem
from wafer.schedule.serializers import ScheduleItemSerializer
from wafer.talks.models import ACCEPTED, CANCELLED
from wafer.talks.models import Talk

class ScheduleRow(object):
    """This is a helpful containter for the schedule view to keep sanity"""
    def __init__(self, schedule_day, slot):
        self.schedule_day = schedule_day
        self.slot = slot
        self.items = {}

    def get_sorted_items(self):
        sorted_items = []
        for venue in self.schedule_day.venues:
            if venue in self.items:
                sorted_items.append(self.items[venue])
        return sorted_items

    def __repr__(self):
        """Debugging aid"""
        return '%s - %s' % (self.slot, self.get_sorted_items())

class ScheduleDay(object):
    """A helpful container for information a days in a schedule view."""
    def __init__(self, day):
        self.day = day
        self.venues = list(day.venue_set.all())
        self.rows = []

class VenueView(DetailView):
    template_name = 'wafer.schedule/venue.html'
    model = Venue

def make_schedule_row(schedule_day, slot, seen_items):
    """Create a row for the schedule table."""
    row = ScheduleRow(schedule_day, slot)
    skip = {}
    expanding = {}
    all_items = list(slot.scheduleitem_set
                     .select_related('talk', 'page', 'venue')
                     .all())

    for item in all_items:
        if item in seen_items:
            # Inc rowspan
            seen_items[item]['rowspan'] += 1
            # Note that we need to skip this during colspan checks
            skip[item.venue] = seen_items[item]
            continue
        scheditem = {'item': item, 'rowspan': 1, 'colspan': 1}
        row.items[item.venue] = scheditem
        seen_items[item] = scheditem
        if item.expand:
            expanding[item.venue] = []

    empty = []
    expanding_right = None
    skipping = 0
    skip_item = None
    for venue in schedule_day.venues:
        if venue in skip:
            # We need to skip all the venues this item spans over
            skipping = 1
            skip_item = skip[venue]
            continue
        if venue in expanding:
            item = row.items[venue]
            for empty_venue in empty:
                row.items.pop(empty_venue)
                item['colspan'] += 1
            empty = []
            expanding_right = item
        elif venue in row.items:
            empty = []
            expanding_right = None
        elif expanding_right:
            expanding_right['colspan'] += 1
        elif skipping > 0 and skipping < skip_item['colspan']:
            skipping += 1
        else:
            skipping = 0
            empty.append(venue)
            row.items[venue] = {'item': None, 'rowspan': 1, 'colspan': 1}

    return row

def generate_schedule(today=None):
    """Helper function which creates an ordered list of schedule days"""
    # We create a list of slots and schedule items
    schedule_days = {}
    seen_items = {}
    for slot in Slot.objects.all().order_by('end_time', 'start_time', 'day'):
        day = slot.get_day()
        if today and day != today:
            # Restrict ourselves to only today
            continue
        schedule_day = schedule_days.get(day)
        if schedule_day is None:
            schedule_day = schedule_days[day] = ScheduleDay(day)
        row = make_schedule_row(schedule_day, slot, seen_items)
        schedule_day.rows.append(row)
    return sorted(schedule_days.values(), key=lambda x: x.day.date)

class ScheduleView(TemplateView):
    template_name = 'wafer.schedule/full_schedule.html'

    def get_context_data(self, **kwargs):
        context = super(ScheduleView, self).get_context_data(**kwargs)
        # Check if the schedule is valid
        context['active'] = False
        if not check_schedule():
            return context
        context['active'] = True
        day = self.request.GET.get('day', None)
        dates = dict([(x.date.strftime('%Y-%m-%d'), x) for x in
                      Day.objects.all()])
        # We choose to return the full schedule if given an invalid date
        day = dates.get(day, None)
        context['schedule_days'] = generate_schedule(day)
        return context

class ScheduleXmlView(ScheduleView):
    template_name = 'wafer.schedule/penta_schedule.xml'
    content_type = 'application/xml'

    def get_context_data(self, **kwargs):
        """Allow adding a 'render_description' parameter"""
        context = super(ScheduleXmlView, self).get_context_data(**kwargs)
        if self.request.GET.get('render_description', None) == '1':
            context['render_description'] = True
        else:
            context['render_description'] = False
        return context

class CurrentView(TemplateView):
    template_name = 'wafer.schedule/current.html'

    def _parse_today(self, day):
        if day is None:
            day = str(datetime.date.today())
        dates = dict([(x.date.strftime('%Y-%m-%d'), x) for x in
                      Day.objects.all()])
        if day not in dates:
            return None
        return ScheduleDay(dates[day])

    def _parse_time(self, time):
        now = datetime.datetime.now().time()
        if time is None:
            return now
        try:
            return datetime.datetime.strptime(time, '%H:%M').time()
        except ValueError:
            pass
        return now

    def _add_note(self, row, note, overlap_note):
        for item in row.items.values():
            if item['rowspan'] == 1:
                item['note'] = note
            else:
                # Must overlap with current slot
                item['note'] = overlap_note

    def _current_slots(self, schedule_day, time):
        today = schedule_day.day
        cur_slot, prev_slot, next_slot = None, None, None
        for slot in Slot.objects.all():
            if slot.get_day() != today:
                continue
            if slot.get_start_time() <= time and slot.end_time > time:
                cur_slot = slot
            elif slot.end_time <= time:
                if not prev_slot or prev_slot.end_time < slot.end_time:
                    prev_slot = slot
            elif slot.get_start_time() >= time:
                if not next_slot or next_slot.end_time > slot.end_time:
                    next_slot = slot
        cur_rows = self._current_rows(
            schedule_day, cur_slot, prev_slot, next_slot)
        return cur_slot, cur_rows

    def _current_rows(self, schedule_day, cur_slot, prev_slot, next_slot):
        seen_items = {}
        rows = []
        for slot in (prev_slot, cur_slot, next_slot):
            if slot:
                row = make_schedule_row(schedule_day, slot, seen_items)
            else:
                row = None
            rows.append(row)
        # Add styling hints. Needs to be after all the schedule rows are
        # created so the spans are set correctly
        if prev_slot:
            self._add_note(rows[0], 'complete', 'current')
        if cur_slot:
            self._add_note(rows[1], 'current', 'current')
        if next_slot:
            self._add_note(rows[2], 'forthcoming', 'current')
        return [r for r in rows if r]

    def get_context_data(self, **kwargs):
        context = super(CurrentView, self).get_context_data(**kwargs)
        # If the schedule is invalid, return a context with active=False
        context['active'] = False
        if not check_schedule():
            return context
        # The schedule is valid, so add active=True and empty slots
        context['active'] = True
        context['slots'] = []
        # Allow refresh time to be overridden
        context['refresh'] = self.request.GET.get('refresh', None)
        # If there are no items scheduled for today, return an empty slots list
        schedule_day = self._parse_today(self.request.GET.get('day', None))
        if schedule_day is None:
            return context
        context['schedule_day'] = schedule_day
        # Allow current time to be overridden
        time = self._parse_time(self.request.GET.get('time', None))

        cur_slot, current_rows = self._current_slots(schedule_day, time)
        context['cur_slot'] = cur_slot
        context['slots'].extend(current_rows)

        return context

class ScheduleItemViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows groups to be viewed or edited.
    """
    queryset = ScheduleItem.objects.all().order_by('id')
    serializer_class = ScheduleItemSerializer
    permission_classes = (IsAdminUser, )

class ScheduleEditView(TemplateView):
    template_name = 'wafer.schedule/edit_schedule.html'

    def _slot_context(self, slot, venues):
        slot_context = {
            'name': slot.name,
            'start_time': slot.get_start_time(),
            'end_time': slot.end_time,
            'id': slot.id,
            'venues': []
        }
        for venue in venues:
            venue_context = {
                'name': venue.name,
                'id': venue.id,
            }
            for schedule_item in slot.scheduleitem_set.all():
                if schedule_item.venue.name == venue.name:
                    venue_context['scheduleitem_id'] = schedule_item.id
                    if schedule_item.talk:
                        talk = schedule_item.talk
                        venue_context['title'] = talk.title
                        venue_context['talk'] = talk
                    if (schedule_item.page and
                            not schedule_item.page.exclude_from_static):
                        page = schedule_item.page
                        venue_context['title'] = page.name
                        venue_context['page'] = page
            slot_context['venues'].append(venue_context)
        return slot_context

    def get_context_data(self, day_id=None, **kwargs):
        context = super(ScheduleEditView, self).get_context_data(**kwargs)

        days = Day.objects.all()
        if day_id:
            day = days.get(id=day_id)
        else:
            day = days.first()

        public_talks = Talk.objects.filter(Q(status=ACCEPTED) |
                                           Q(status=CANCELLED))
        venues = Venue.objects.filter(days__in=[day])
        slots = Slot.objects.all().select_related(
            'day', 'previous_slot').prefetch_related(
            'scheduleitem_set', 'slot_set').order_by(
                'end_time', 'start_time', 'day')
        aggregated_slots = []

        for slot in slots:
            if day != slot.get_day():
                continue
            aggregated_slots.append(self._slot_context(slot, venues))

        context['day'] = day
        context['venues'] = venues
        context['slots'] = aggregated_slots
        context['talks_all'] = public_talks
        context['talks_unassigned'] = public_talks.filter(scheduleitem=None)
        context['pages'] = Page.objects.all()
        context['days'] = days
        context['validation_errors'] = validate_schedule()
        return context

class ICalView(View):

    def get(self, request):
        """Create a iCal file from the schedule"""
        # Heavily inspired by https://djangosnippets.org/snippets/2223/ and
        # the icalendar documentation
        calendar = Calendar()
        site = get_current_site(request)
        calendar.add('prodid', '-//%s Schedule//%s//' % (site.name, site.domain))
        calendar.add('version', '2.0')

        # Since we don't need to format anything here, we can just use a list
        # of schedule items
        for item in ScheduleItem.objects.all():
            sched_event = Event()
            sched_event.add('dtstamp', item.last_updated)
            sched_event.add('summary', item.get_title())
            sched_event.add('location', item.venue.name)
            sched_event.add('dtstart', item.get_start_datetime())
            sched_event.add('duration', datetime.timedelta(minutes=item.get_duration_minutes()))
            calendar.add_component(sched_event)
        response = HttpResponse(calendar.to_ical(), content_type="text/calendar")
        response['Content-Disposition'] = 'attachment; filename=schedule.ics'
        return response
<commit_msg>Add some missing requirements<commit_after>import datetime

from icalendar import Calendar, Event

from django.db.models import Q
from django.views.generic import DetailView, TemplateView, View
from django.contrib.sites.shortcuts import get_current_site
from django.http import HttpResponse

from rest_framework import viewsets
from rest_framework.permissions import IsAdminUser
from wafer.pages.models import Page
from wafer.schedule.models import Venue, Slot, Day
from wafer.schedule.admin import check_schedule, validate_schedule
from wafer.schedule.models import ScheduleItem
from wafer.schedule.serializers import ScheduleItemSerializer
from wafer.talks.models import ACCEPTED, CANCELLED
from wafer.talks.models import Talk

class ScheduleRow(object):
    """This is a helpful containter for the schedule view to keep sanity"""
    def __init__(self, schedule_day, slot):
        self.schedule_day = schedule_day
        self.slot = slot
        self.items = {}

    def get_sorted_items(self):
        sorted_items = []
        for venue in self.schedule_day.venues:
            if venue in self.items:
                sorted_items.append(self.items[venue])
        return sorted_items

    def __repr__(self):
        """Debugging aid"""
        return '%s - %s' % (self.slot, self.get_sorted_items())

class ScheduleDay(object):
    """A helpful container for information a days in a schedule view."""
    def __init__(self, day):
        self.day = day
        self.venues = list(day.venue_set.all())
        self.rows = []

class VenueView(DetailView):
    template_name = 'wafer.schedule/venue.html'
    model = Venue

def make_schedule_row(schedule_day, slot, seen_items):
    """Create a row for the schedule table."""
    row = ScheduleRow(schedule_day, slot)
    skip = {}
    expanding = {}
    all_items = list(slot.scheduleitem_set
                     .select_related('talk', 'page', 'venue')
                     .all())

    for item in all_items:
        if item in seen_items:
            # Inc rowspan
            seen_items[item]['rowspan'] += 1
            # Note that we need to skip this during colspan checks
            skip[item.venue] = seen_items[item]
            continue
        scheditem = {'item': item, 'rowspan': 1, 'colspan': 1}
        row.items[item.venue] = scheditem
        seen_items[item] = scheditem
        if item.expand:
            expanding[item.venue] = []

    empty = []
    expanding_right = None
    skipping = 0
    skip_item = None
    for venue in schedule_day.venues:
        if venue in skip:
            # We need to skip all the venues this item spans over
            skipping = 1
            skip_item = skip[venue]
            continue
        if venue in expanding:
            item = row.items[venue]
            for empty_venue in empty:
                row.items.pop(empty_venue)
                item['colspan'] += 1
            empty = []
            expanding_right = item
        elif venue in row.items:
            empty = []
            expanding_right = None
        elif expanding_right:
            expanding_right['colspan'] += 1
        elif skipping > 0 and skipping < skip_item['colspan']:
            skipping += 1
        else:
            skipping = 0
            empty.append(venue)
            row.items[venue] = {'item': None, 'rowspan': 1, 'colspan': 1}

    return row

def generate_schedule(today=None):
    """Helper function which creates an ordered list of schedule days"""
    # We create a list of slots and schedule items
    schedule_days = {}
    seen_items = {}
    for slot in Slot.objects.all().order_by('end_time', 'start_time', 'day'):
        day = slot.get_day()
        if today and day != today:
            # Restrict ourselves to only today
            continue
        schedule_day = schedule_days.get(day)
        if schedule_day is None:
            schedule_day = schedule_days[day] = ScheduleDay(day)
        row = make_schedule_row(schedule_day, slot, seen_items)
        schedule_day.rows.append(row)
    return sorted(schedule_days.values(), key=lambda x: x.day.date)

class ScheduleView(TemplateView):
    template_name = 'wafer.schedule/full_schedule.html'

    def get_context_data(self, **kwargs):
        context = super(ScheduleView, self).get_context_data(**kwargs)
        # Check if the schedule is valid
        context['active'] = False
        if not check_schedule():
            return context
        context['active'] = True
        day = self.request.GET.get('day', None)
        dates = dict([(x.date.strftime('%Y-%m-%d'), x) for x in
                      Day.objects.all()])
        # We choose to return the full schedule if given an invalid date
        day = dates.get(day, None)
        context['schedule_days'] = generate_schedule(day)
        return context

class ScheduleXmlView(ScheduleView):
    template_name = 'wafer.schedule/penta_schedule.xml'
    content_type = 'application/xml'

    def get_context_data(self, **kwargs):
        """Allow adding a 'render_description' parameter"""
        context = super(ScheduleXmlView, self).get_context_data(**kwargs)
        if self.request.GET.get('render_description', None) == '1':
            context['render_description'] = True
        else:
            context['render_description'] = False
        return context

class CurrentView(TemplateView):
    template_name = 'wafer.schedule/current.html'

    def _parse_today(self, day):
        if day is None:
            day = str(datetime.date.today())
        dates = dict([(x.date.strftime('%Y-%m-%d'), x) for x in
                      Day.objects.all()])
        if day not in dates:
            return None
        return ScheduleDay(dates[day])

    def _parse_time(self, time):
        now = datetime.datetime.now().time()
        if time is None:
            return now
        try:
            return datetime.datetime.strptime(time, '%H:%M').time()
        except ValueError:
            pass
        return now

    def _add_note(self, row, note, overlap_note):
        for item in row.items.values():
            if item['rowspan'] == 1:
                item['note'] = note
            else:
                # Must overlap with current slot
                item['note'] = overlap_note

    def _current_slots(self, schedule_day, time):
        today = schedule_day.day
        cur_slot, prev_slot, next_slot = None, None, None
        for slot in Slot.objects.all():
            if slot.get_day() != today:
                continue
            if slot.get_start_time() <= time and slot.end_time > time:
                cur_slot = slot
            elif slot.end_time <= time:
                if not prev_slot or prev_slot.end_time < slot.end_time:
                    prev_slot = slot
            elif slot.get_start_time() >= time:
                if not next_slot or next_slot.end_time > slot.end_time:
                    next_slot = slot
        cur_rows = self._current_rows(
            schedule_day, cur_slot, prev_slot, next_slot)
        return cur_slot, cur_rows

    def _current_rows(self, schedule_day, cur_slot, prev_slot, next_slot):
        seen_items = {}
        rows = []
        for slot in (prev_slot, cur_slot, next_slot):
            if slot:
                row = make_schedule_row(schedule_day, slot, seen_items)
            else:
                row = None
            rows.append(row)
        # Add styling hints. Needs to be after all the schedule rows are
        # created so the spans are set correctly
        if prev_slot:
            self._add_note(rows[0], 'complete', 'current')
        if cur_slot:
            self._add_note(rows[1], 'current', 'current')
        if next_slot:
            self._add_note(rows[2], 'forthcoming', 'current')
        return [r for r in rows if r]

    def get_context_data(self, **kwargs):
        context = super(CurrentView, self).get_context_data(**kwargs)
        # If the schedule is invalid, return a context with active=False
        context['active'] = False
        if not check_schedule():
            return context
        # The schedule is valid, so add active=True and empty slots
        context['active'] = True
        context['slots'] = []
        # Allow refresh time to be overridden
        context['refresh'] = self.request.GET.get('refresh', None)
        # If there are no items scheduled for today, return an empty slots list
        schedule_day = self._parse_today(self.request.GET.get('day', None))
        if schedule_day is None:
            return context
        context['schedule_day'] = schedule_day
        # Allow current time to be overridden
        time = self._parse_time(self.request.GET.get('time', None))

        cur_slot, current_rows = self._current_slots(schedule_day, time)
        context['cur_slot'] = cur_slot
        context['slots'].extend(current_rows)

        return context

class ScheduleItemViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows groups to be viewed or edited.
    """
    queryset = ScheduleItem.objects.all().order_by('id')
    serializer_class = ScheduleItemSerializer
    permission_classes = (IsAdminUser, )

class ScheduleEditView(TemplateView):
    template_name = 'wafer.schedule/edit_schedule.html'

    def _slot_context(self, slot, venues):
        slot_context = {
            'name': slot.name,
            'start_time': slot.get_start_time(),
            'end_time': slot.end_time,
            'id': slot.id,
            'venues': []
        }
        for venue in venues:
            venue_context = {
                'name': venue.name,
                'id': venue.id,
            }
            for schedule_item in slot.scheduleitem_set.all():
                if schedule_item.venue.name == venue.name:
                    venue_context['scheduleitem_id'] = schedule_item.id
                    if schedule_item.talk:
                        talk = schedule_item.talk
                        venue_context['title'] = talk.title
                        venue_context['talk'] = talk
                    if (schedule_item.page and
                            not schedule_item.page.exclude_from_static):
                        page = schedule_item.page
                        venue_context['title'] = page.name
                        venue_context['page'] = page
            slot_context['venues'].append(venue_context)
        return slot_context

    def get_context_data(self, day_id=None, **kwargs):
        context = super(ScheduleEditView, self).get_context_data(**kwargs)

        days = Day.objects.all()
        if day_id:
            day = days.get(id=day_id)
        else:
            day = days.first()

        public_talks = Talk.objects.filter(Q(status=ACCEPTED) |
                                           Q(status=CANCELLED))
        venues = Venue.objects.filter(days__in=[day])
        slots = Slot.objects.all().select_related(
            'day', 'previous_slot').prefetch_related(
            'scheduleitem_set', 'slot_set').order_by(
                'end_time', 'start_time', 'day')
        aggregated_slots = []

        for slot in slots:
            if day != slot.get_day():
                continue
            aggregated_slots.append(self._slot_context(slot, venues))

        context['day'] = day
        context['venues'] = venues
        context['slots'] = aggregated_slots
        context['talks_all'] = public_talks
        context['talks_unassigned'] = public_talks.filter(scheduleitem=None)
        context['pages'] = Page.objects.all()
        context['days'] = days
        context['validation_errors'] = validate_schedule()
        return context

class ICalView(View):

    def get(self, request):
        """Create a iCal file from the schedule"""
        # Heavily inspired by https://djangosnippets.org/snippets/2223/ and
        # the icalendar documentation
        calendar = Calendar()
        site = get_current_site(request)
        calendar.add('prodid', '-//%s Schedule//%s//' % (site.name, site.domain))
        calendar.add('version', '2.0')

        # Since we don't need to format anything here, we can just use a list
        # of schedule items
        for item in ScheduleItem.objects.all():
            sched_event = Event()
            sched_event.add('dtstamp', item.last_updated)
            sched_event.add('summary', item.get_title())
            sched_event.add('location', item.venue.name)
            sched_event.add('dtstart', item.get_start_datetime())
            sched_event.add('duration', datetime.timedelta(minutes=item.get_duration_minutes()))
            sched_event.add('class', 'PUBLIC')
            sched_event.add('uid', '%s@%s' % (item.pk, site.domain))
            calendar.add_component(sched_event)
        response = HttpResponse(calendar.to_ical(), content_type="text/calendar")
        response['Content-Disposition'] = 'attachment; filename=schedule.ics'
        return response
<|endoftext|>


<commit_before>from __future__ import absolute_import, print_function, unicode_literals

import datetime
import logging
import mimetypes
import os
import os.path
import re
import shutil
import subprocess
import tempfile
import threading

from django.conf import settings
from django.core.exceptions import SuspiciousFileOperation
from django.core.files.base import ContentFile
from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save, pre_delete
from django.dispatch.dispatcher import receiver
from django.forms.utils import flatatt
from django.utils.encoding import python_2_unicode_compatible
from django.utils.text import mark_safe
from django.utils.translation import ugettext_lazy as _
from enumchoicefield import ChoiceEnum, EnumChoiceField
from taggit.managers import TaggableManager
from wagtail.wagtailadmin.utils import get_object_usage
from wagtail.wagtailcore.models import CollectionMember
from wagtail.wagtailsearch import index
from wagtail.wagtailsearch.queryset import SearchableQuerySetMixin

from wagtailvideos import ffmpeg

logger = logging.getLogger(__name__)

class VideoQuality(ChoiceEnum):
    default = 'Default'
    lowest = 'Low'
    highest = 'High'

class MediaFormats(ChoiceEnum):
    webm = 'VP8 and Vorbis in WebM'
    mp4 = 'H.264 and MP3 in Mp4'
    ogg = 'Theora and Voris in Ogg'

    def get_quality_param(self, quality):
        if self is MediaFormats.webm:
            return {
                VideoQuality.lowest: '50',
                VideoQuality.default: '22',
                VideoQuality.highest: '4'
            }[quality]
        elif self is MediaFormats.mp4:
            return {
                VideoQuality.lowest: '28',
                VideoQuality.default: '24',
                VideoQuality.highest: '18'
            }[quality]
        elif self is MediaFormats.ogg:
            return {
                VideoQuality.lowest: '5',
                VideoQuality.default: '7',
                VideoQuality.highest: '9'
            }[quality]

class VideoQuerySet(SearchableQuerySetMixin, models.QuerySet):
    pass

def get_upload_to(instance, filename):
    # Dumb proxy to instance method.
    return instance.get_upload_to(filename)

@python_2_unicode_compatible
class AbstractVideo(CollectionMember, index.Indexed, models.Model):
    title = models.CharField(max_length=255, verbose_name=_('title'))
    file = models.FileField(
        verbose_name=_('file'), upload_to=get_upload_to)
    thumbnail = models.ImageField(upload_to=get_upload_to, null=True, blank=True)
    created_at = models.DateTimeField(verbose_name=_('created at'), auto_now_add=True, db_index=True)
    duration = models.DurationField(blank=True, null=True)
    uploaded_by_user = models.ForeignKey(
        settings.AUTH_USER_MODEL, verbose_name=_('uploaded by user'),
        null=True, blank=True, editable=False, on_delete=models.SET_NULL
    )

    tags = TaggableManager(help_text=None, blank=True, verbose_name=_('tags'))

    file_size = models.PositiveIntegerField(null=True, editable=False)

    objects = VideoQuerySet.as_manager()

    search_fields = list(CollectionMember.search_fields) + [
        index.SearchField('title', partial_match=True, boost=10),
        index.RelatedFields('tags', [
            index.SearchField('name', partial_match=True, boost=10),
        ]),
        index.FilterField('uploaded_by_user'),
    ]

    def is_stored_locally(self):
        """
        Returns True if the image is hosted on the local filesystem
        """
        try:
            self.file.path
            return True
        except NotImplementedError:
            return False

    def get_file_size(self):
        if self.file_size is None:
            try:
                self.file_size = self.file.size
            except OSError:
                # File doesn't exist
                return

            self.save(update_fields=['file_size'])

        return self.file_size

    def get_upload_to(self, filename):
        folder_name = 'original_videos'
        filename = self.file.field.storage.get_valid_name(filename)
        max_length = self._meta.get_field('file').max_length

        # Truncate filename so it fits in the 100 character limit
        # https://code.djangoproject.com/ticket/9893
        file_path = os.path.join(folder_name, filename)
        too_long = len(file_path) - max_length
        if too_long > 0:
            head, ext = os.path.splitext(filename)
            if too_long > len(head) + 1:
                raise SuspiciousFileOperation('File name can not be shortened to a safe length')
            filename = head[:-too_long] + ext
            file_path = os.path.join(folder_name, filename)
        return os.path.join(folder_name, filename)

    def get_usage(self):
        return get_object_usage(self)

    @property
    def usage_url(self):
        return reverse('wagtailvideos:video_usage', args=(self.id,))

    @property
    def formatted_duration(self):
        if(self.duration):
            hours, remainder = divmod(self.duration.seconds, 3600)
            minutes, seconds = divmod(remainder, 60)
            return "%d:%02d:%02d" % (hours, minutes, seconds)
        return ''

    def __str__(self):
        return self.title

    def get_duration(self):
        if self.duration:
            return self.duration

        if not ffmpeg_installed():
            return None

        file_path = self.file.path
        try:
            # FIXME prints out extra stuff on travis, pip stderr to dev/null
            show_format = subprocess.check_output(['ffprobe', file_path, '-show_format', '-v', 'quiet'])
            show_format = show_format.decode("utf-8")
            # show_format comes out in key=value pairs seperated by newlines
            duration = re.findall(r'([duration^=]+)=([^=]+)(?:\n|$)', show_format)[0][1]
            return datetime.timedelta(seconds=float(duration))
        except subprocess.CalledProcessError:
            logger.exception("Getting video duration failed")
            return None

    def get_thumbnail(self):
        if self.thumbnail:
            return self.thumbnail

        if not ffmpeg_installed():
            return None

        file_path = self.file.path
        file_name = self.filename(include_ext=False) + '_thumb.jpg'

        try:
            output_dir = tempfile.mkdtemp()
            output_file = os.path.join(output_dir, file_name)
            try:
                FNULL = open(os.devnull, 'r')
                subprocess.check_call([
                    'ffmpeg',
                    '-v', 'quiet',
                    '-itsoffset', '-4',
                    '-i', file_path,
                    '-vcodec', 'mjpeg',
                    '-vframes', '1',
                    '-an', '-f', 'rawvideo',
                    '-s', '320x240',
                    output_file,
                ], stdin=FNULL)
            except subprocess.CalledProcessError:
                return None
            return ContentFile(open(output_file, 'rb').read(), file_name)
        finally:
            shutil.rmtree(output_dir, ignore_errors=True)

    def save(self, **kwargs):
        super(AbstractVideo, self).save(**kwargs)

    @property
    def url(self):
        return self.file.url

    def filename(self, include_ext=True):
        if include_ext:
            return os.path.basename(self.file.name)
        else:
            return os.path.splitext(os.path.basename(self.file.name))[0]

    @property
    def file_ext(self):
        return os.path.splitext(self.filename())[1][1:]

    def is_editable_by_user(self, user):
        from wagtailvideos.permissions import permission_policy
        return permission_policy.user_has_permission_for_instance(user, 'change', self)

    @classmethod
    def get_transcode_model(cls):
        return cls.transcodes.rel.related_model

    def get_transcode(self, media_format):
        Transcode = self.get_transcode_model()
        try:
            return self.transcodes.get(media_format=media_format)
        except Transcode.DoesNotExist:
            return self.do_transcode(media_format)

    def video_tag(self, attrs=None):
        if attrs is None:
            attrs = {}
        else:
            attrs = attrs.copy()
        if self.thumbnail:
            attrs['poster'] = self.thumbnail.url

        transcodes = self.transcodes.exclude(processing=True).filter(error_message__exact='')
        sources = []
        for transcode in transcodes:
            sources.append("<source src='{0}' type='video/{1}' >".format(transcode.url, transcode.media_format.name))

        mime = mimetypes.MimeTypes()
        sources.append("<source src='{0}' type='{1}'>"
                       .format(self.url, mime.guess_type(self.url)[0]))

        sources.append("
Sorry, your browser doesn't support playback for this video
")
        return mark_safe(
            "<video {0}>\n{1}\n</video>".format(flatatt(attrs), "\n".join(sources)))

    def do_transcode(self, media_format, quality):
        transcode, created = self.transcodes.get_or_create(
            media_format=media_format,
        )
        if transcode.processing is False:
            transcode.processing = True
            transcode.error_messages = ''
            transcode.quality = quality
            # Lock the transcode model
            transcode.save(update_fields=['processing', 'error_message',
                                          'quality'])
            TranscodingThread(transcode).start()
        else:
            pass  # TODO Queue?

    class Meta:
        abstract = True

class Video(AbstractVideo):
    admin_form_fields = (
        'title',
        'file',
        'collection',
        'thumbnail',
        'tags',
    )

class TranscodingThread(threading.Thread):
    def __init__(self, transcode, **kwargs):
        super(TranscodingThread, self).__init__(**kwargs)
        self.transcode = transcode

    def run(self):
        video = self.transcode.video
        media_format = self.transcode.media_format
        input_file = video.file.path
        output_dir = tempfile.mkdtemp()
        transcode_name = "{0}.{1}".format(
            video.filename(include_ext=False),
            media_format.name)

        output_file = os.path.join(output_dir, transcode_name)
        FNULL = open(os.devnull, 'r')
        quality_param = media_format.get_quality_param(self.transcode.quality)
        args = ['ffmpeg', '-hide_banner', '-i', input_file]
        try:
            if media_format is MediaFormats.ogg:
                subprocess.check_output(args + [
                    '-codec:v', 'libtheora',
                    '-qscale:v', quality_param,
                    '-codec:a', 'libvorbis',
                    '-qscale:a', '5',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            elif media_format is MediaFormats.mp4:
                subprocess.check_output(args + [
                    '-codec:v', 'libx264',
                    '-preset', 'slow',  # TODO Checkout other presets
                    '-crf', quality_param,
                    '-codec:a', 'copy',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            elif media_format is MediaFormats.webm:
                subprocess.check_output(args + [
                    '-codec:v', 'libvpx',
                    '-crf', quality_param,
                    '-codec:a', 'libvorbis',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            self.transcode.file = ContentFile(
                open(output_file, 'rb').read(), transcode_name)
            self.transcode.error_message = ''
        except subprocess.CalledProcessError as error:
            self.transcode.error_message = error.output

        finally:
            self.transcode.processing = False
            self.transcode.save()
            shutil.rmtree(output_dir, ignore_errors=True)

# Delete files when model is deleted
@receiver(pre_delete, sender=Video)
def video_delete(sender, instance, **kwargs):
    instance.thumbnail.delete(False)
    instance.file.delete(False)

# Fields that need the actual video file to create
@receiver(post_save, sender=Video)
def video_saved(sender, instance, **kwargs):
    if not ffmpeg.installed():
        return

    if hasattr(instance, '_from_signal'):
        return
    instance.thumbnail = instance.get_thumbnail()
    instance.duration = instance.get_duration()
    instance.file_size = instance.file.size
    instance._from_signal = True
    instance.save()
    del instance._from_signal

class AbstractVideoTranscode(models.Model):
    media_format = EnumChoiceField(MediaFormats)
    quality = EnumChoiceField(VideoQuality, default=VideoQuality.default)
    processing = models.BooleanField(default=False)
    file = models.FileField(null=True, blank=True, verbose_name=_('file'),
                            upload_to=get_upload_to)
    error_message = models.TextField(blank=True)

    @property
    def url(self):
        return self.file.url

    def get_upload_to(self, filename):
        folder_name = 'video_transcodes'
        filename = self.file.field.storage.get_valid_name(filename)
        return os.path.join(folder_name, filename)

    class Meta:
        abstract = True

class VideoTranscode(AbstractVideoTranscode):
    video = models.ForeignKey(Video, related_name='transcodes')

    class Meta:
        unique_together = (
            ('video', 'media_format')
        )

# Delete files when model is deleted
@receiver(pre_delete, sender=VideoTranscode)
def transcode_delete(sender, instance, **kwargs):
    instance.file.delete(False)
<commit_msg>Set an explicit ordering on Videos<commit_after>from __future__ import absolute_import, print_function, unicode_literals

import datetime
import logging
import mimetypes
import os
import os.path
import re
import shutil
import subprocess
import tempfile
import threading

from django.conf import settings
from django.core.exceptions import SuspiciousFileOperation
from django.core.files.base import ContentFile
from django.core.urlresolvers import reverse
from django.db import models
from django.db.models.signals import post_save, pre_delete
from django.dispatch.dispatcher import receiver
from django.forms.utils import flatatt
from django.utils.encoding import python_2_unicode_compatible
from django.utils.text import mark_safe
from django.utils.translation import ugettext_lazy as _
from enumchoicefield import ChoiceEnum, EnumChoiceField
from taggit.managers import TaggableManager
from wagtail.wagtailadmin.utils import get_object_usage
from wagtail.wagtailcore.models import CollectionMember
from wagtail.wagtailsearch import index
from wagtail.wagtailsearch.queryset import SearchableQuerySetMixin

from wagtailvideos import ffmpeg

logger = logging.getLogger(__name__)

class VideoQuality(ChoiceEnum):
    default = 'Default'
    lowest = 'Low'
    highest = 'High'

class MediaFormats(ChoiceEnum):
    webm = 'VP8 and Vorbis in WebM'
    mp4 = 'H.264 and MP3 in Mp4'
    ogg = 'Theora and Voris in Ogg'

    def get_quality_param(self, quality):
        if self is MediaFormats.webm:
            return {
                VideoQuality.lowest: '50',
                VideoQuality.default: '22',
                VideoQuality.highest: '4'
            }[quality]
        elif self is MediaFormats.mp4:
            return {
                VideoQuality.lowest: '28',
                VideoQuality.default: '24',
                VideoQuality.highest: '18'
            }[quality]
        elif self is MediaFormats.ogg:
            return {
                VideoQuality.lowest: '5',
                VideoQuality.default: '7',
                VideoQuality.highest: '9'
            }[quality]

class VideoQuerySet(SearchableQuerySetMixin, models.QuerySet):
    pass

def get_upload_to(instance, filename):
    # Dumb proxy to instance method.
    return instance.get_upload_to(filename)

@python_2_unicode_compatible
class AbstractVideo(CollectionMember, index.Indexed, models.Model):
    title = models.CharField(max_length=255, verbose_name=_('title'))
    file = models.FileField(
        verbose_name=_('file'), upload_to=get_upload_to)
    thumbnail = models.ImageField(upload_to=get_upload_to, null=True, blank=True)
    created_at = models.DateTimeField(verbose_name=_('created at'), auto_now_add=True, db_index=True)
    duration = models.DurationField(blank=True, null=True)
    uploaded_by_user = models.ForeignKey(
        settings.AUTH_USER_MODEL, verbose_name=_('uploaded by user'),
        null=True, blank=True, editable=False, on_delete=models.SET_NULL
    )

    tags = TaggableManager(help_text=None, blank=True, verbose_name=_('tags'))

    file_size = models.PositiveIntegerField(null=True, editable=False)

    objects = VideoQuerySet.as_manager()

    search_fields = list(CollectionMember.search_fields) + [
        index.SearchField('title', partial_match=True, boost=10),
        index.RelatedFields('tags', [
            index.SearchField('name', partial_match=True, boost=10),
        ]),
        index.FilterField('uploaded_by_user'),
    ]

    def is_stored_locally(self):
        """
        Returns True if the image is hosted on the local filesystem
        """
        try:
            self.file.path
            return True
        except NotImplementedError:
            return False

    def get_file_size(self):
        if self.file_size is None:
            try:
                self.file_size = self.file.size
            except OSError:
                # File doesn't exist
                return

            self.save(update_fields=['file_size'])

        return self.file_size

    def get_upload_to(self, filename):
        folder_name = 'original_videos'
        filename = self.file.field.storage.get_valid_name(filename)
        max_length = self._meta.get_field('file').max_length

        # Truncate filename so it fits in the 100 character limit
        # https://code.djangoproject.com/ticket/9893
        file_path = os.path.join(folder_name, filename)
        too_long = len(file_path) - max_length
        if too_long > 0:
            head, ext = os.path.splitext(filename)
            if too_long > len(head) + 1:
                raise SuspiciousFileOperation('File name can not be shortened to a safe length')
            filename = head[:-too_long] + ext
            file_path = os.path.join(folder_name, filename)
        return os.path.join(folder_name, filename)

    def get_usage(self):
        return get_object_usage(self)

    @property
    def usage_url(self):
        return reverse('wagtailvideos:video_usage', args=(self.id,))

    @property
    def formatted_duration(self):
        if(self.duration):
            hours, remainder = divmod(self.duration.seconds, 3600)
            minutes, seconds = divmod(remainder, 60)
            return "%d:%02d:%02d" % (hours, minutes, seconds)
        return ''

    def __str__(self):
        return self.title

    def get_duration(self):
        if self.duration:
            return self.duration

        if not ffmpeg_installed():
            return None

        file_path = self.file.path
        try:
            # FIXME prints out extra stuff on travis, pip stderr to dev/null
            show_format = subprocess.check_output(['ffprobe', file_path, '-show_format', '-v', 'quiet'])
            show_format = show_format.decode("utf-8")
            # show_format comes out in key=value pairs seperated by newlines
            duration = re.findall(r'([duration^=]+)=([^=]+)(?:\n|$)', show_format)[0][1]
            return datetime.timedelta(seconds=float(duration))
        except subprocess.CalledProcessError:
            logger.exception("Getting video duration failed")
            return None

    def get_thumbnail(self):
        if self.thumbnail:
            return self.thumbnail

        if not ffmpeg_installed():
            return None

        file_path = self.file.path
        file_name = self.filename(include_ext=False) + '_thumb.jpg'

        try:
            output_dir = tempfile.mkdtemp()
            output_file = os.path.join(output_dir, file_name)
            try:
                FNULL = open(os.devnull, 'r')
                subprocess.check_call([
                    'ffmpeg',
                    '-v', 'quiet',
                    '-itsoffset', '-4',
                    '-i', file_path,
                    '-vcodec', 'mjpeg',
                    '-vframes', '1',
                    '-an', '-f', 'rawvideo',
                    '-s', '320x240',
                    output_file,
                ], stdin=FNULL)
            except subprocess.CalledProcessError:
                return None
            return ContentFile(open(output_file, 'rb').read(), file_name)
        finally:
            shutil.rmtree(output_dir, ignore_errors=True)

    def save(self, **kwargs):
        super(AbstractVideo, self).save(**kwargs)

    @property
    def url(self):
        return self.file.url

    def filename(self, include_ext=True):
        if include_ext:
            return os.path.basename(self.file.name)
        else:
            return os.path.splitext(os.path.basename(self.file.name))[0]

    @property
    def file_ext(self):
        return os.path.splitext(self.filename())[1][1:]

    def is_editable_by_user(self, user):
        from wagtailvideos.permissions import permission_policy
        return permission_policy.user_has_permission_for_instance(user, 'change', self)

    @classmethod
    def get_transcode_model(cls):
        return cls.transcodes.rel.related_model

    def get_transcode(self, media_format):
        Transcode = self.get_transcode_model()
        try:
            return self.transcodes.get(media_format=media_format)
        except Transcode.DoesNotExist:
            return self.do_transcode(media_format)

    def video_tag(self, attrs=None):
        if attrs is None:
            attrs = {}
        else:
            attrs = attrs.copy()
        if self.thumbnail:
            attrs['poster'] = self.thumbnail.url

        transcodes = self.transcodes.exclude(processing=True).filter(error_message__exact='')
        sources = []
        for transcode in transcodes:
            sources.append("<source src='{0}' type='video/{1}' >".format(transcode.url, transcode.media_format.name))

        mime = mimetypes.MimeTypes()
        sources.append("<source src='{0}' type='{1}'>"
                       .format(self.url, mime.guess_type(self.url)[0]))

        sources.append("
Sorry, your browser doesn't support playback for this video
")
        return mark_safe(
            "<video {0}>\n{1}\n</video>".format(flatatt(attrs), "\n".join(sources)))

    def do_transcode(self, media_format, quality):
        transcode, created = self.transcodes.get_or_create(
            media_format=media_format,
        )
        if transcode.processing is False:
            transcode.processing = True
            transcode.error_messages = ''
            transcode.quality = quality
            # Lock the transcode model
            transcode.save(update_fields=['processing', 'error_message',
                                          'quality'])
            TranscodingThread(transcode).start()
        else:
            pass  # TODO Queue?

    class Meta:
        abstract = True
        ordering = ['-created_at']

class Video(AbstractVideo):
    admin_form_fields = (
        'title',
        'file',
        'collection',
        'thumbnail',
        'tags',
    )

class TranscodingThread(threading.Thread):
    def __init__(self, transcode, **kwargs):
        super(TranscodingThread, self).__init__(**kwargs)
        self.transcode = transcode

    def run(self):
        video = self.transcode.video
        media_format = self.transcode.media_format
        input_file = video.file.path
        output_dir = tempfile.mkdtemp()
        transcode_name = "{0}.{1}".format(
            video.filename(include_ext=False),
            media_format.name)

        output_file = os.path.join(output_dir, transcode_name)
        FNULL = open(os.devnull, 'r')
        quality_param = media_format.get_quality_param(self.transcode.quality)
        args = ['ffmpeg', '-hide_banner', '-i', input_file]
        try:
            if media_format is MediaFormats.ogg:
                subprocess.check_output(args + [
                    '-codec:v', 'libtheora',
                    '-qscale:v', quality_param,
                    '-codec:a', 'libvorbis',
                    '-qscale:a', '5',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            elif media_format is MediaFormats.mp4:
                subprocess.check_output(args + [
                    '-codec:v', 'libx264',
                    '-preset', 'slow',  # TODO Checkout other presets
                    '-crf', quality_param,
                    '-codec:a', 'copy',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            elif media_format is MediaFormats.webm:
                subprocess.check_output(args + [
                    '-codec:v', 'libvpx',
                    '-crf', quality_param,
                    '-codec:a', 'libvorbis',
                    output_file,
                ], stdin=FNULL, stderr=subprocess.STDOUT)
            self.transcode.file = ContentFile(
                open(output_file, 'rb').read(), transcode_name)
            self.transcode.error_message = ''
        except subprocess.CalledProcessError as error:
            self.transcode.error_message = error.output

        finally:
            self.transcode.processing = False
            self.transcode.save()
            shutil.rmtree(output_dir, ignore_errors=True)

# Delete files when model is deleted
@receiver(pre_delete, sender=Video)
def video_delete(sender, instance, **kwargs):
    instance.thumbnail.delete(False)
    instance.file.delete(False)

# Fields that need the actual video file to create
@receiver(post_save, sender=Video)
def video_saved(sender, instance, **kwargs):
    if not ffmpeg.installed():
        return

    if hasattr(instance, '_from_signal'):
        return
    instance.thumbnail = instance.get_thumbnail()
    instance.duration = instance.get_duration()
    instance.file_size = instance.file.size
    instance._from_signal = True
    instance.save()
    del instance._from_signal

class AbstractVideoTranscode(models.Model):
    media_format = EnumChoiceField(MediaFormats)
    quality = EnumChoiceField(VideoQuality, default=VideoQuality.default)
    processing = models.BooleanField(default=False)
    file = models.FileField(null=True, blank=True, verbose_name=_('file'),
                            upload_to=get_upload_to)
    error_message = models.TextField(blank=True)

    @property
    def url(self):
        return self.file.url

    def get_upload_to(self, filename):
        folder_name = 'video_transcodes'
        filename = self.file.field.storage.get_valid_name(filename)
        return os.path.join(folder_name, filename)

    class Meta:
        abstract = True

class VideoTranscode(AbstractVideoTranscode):
    video = models.ForeignKey(Video, related_name='transcodes')

    class Meta:
        unique_together = (
            ('video', 'media_format')
        )

# Delete files when model is deleted
@receiver(pre_delete, sender=VideoTranscode)
def transcode_delete(sender, instance, **kwargs):
    instance.file.delete(False)
<|endoftext|>


<commit_before>with open(r"{{cookiecutter.project_slug}}/wscript") as f:
    lines = f.read()

out = lines.replace('ctx.load("sphinx_build")', "")
out = out.replace(r'ctx.load("biber")', "")
out = out.replace(r'ctx.load("tex")', "")

with open(r"{{cookiecutter.project_slug}}/wscript", "w") as f:
    f.write(out)

for folder in ["matlab", "stata", "r", "python"]:
    with open(r"{{cookiecutter.project_slug}}" + f"/src_{folder}/wscript") as f:
        lines = f.read()
    out = lines.replace('ctx.recurse("paper")', "")
    out = out.replace('ctx.recurse("documentation")', "")

    with open(r"{{cookiecutter.project_slug}}" + f"/src_{folder}/wscript", "w") as f:
        f.write(out)
<commit_msg>Remove old file.<commit_after><|endoftext|>


<commit_before>import os
import json
import logging
import logging.config

PROJECT_NAME = 'waterbutler'
PROJECT_CONFIG_PATH = '~/.cos'

try:
    import colorlog  # noqa
    DEFAULT_FORMATTER = {
        '()': 'colorlog.ColoredFormatter',
        'format': '%(cyan)s[%(asctime)s]%(log_color)s[%(levelname)s][%(name)s]: %(reset)s%(message)s'
    }
except ImportError:
    DEFAULT_FORMATTER = {
        '()': 'waterbutler.core.logging.MaskFormatter',
        'format': '[%(asctime)s][%(levelname)s][%(name)s]: %(message)s',
        'pattern': '(?<=cookie=)(.*?)(?=&|$)',
        'mask': '***'
    }
DEFAULT_LOGGING_CONFIG = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'console': DEFAULT_FORMATTER,
    },
    'handlers': {
        'console': {
            'class': 'logging.StreamHandler',
            'level': 'INFO',
            'formatter': 'console'
        },
        'syslog': {
            'class': 'logging.handlers.SysLogHandler',
            'level': 'INFO'
        }
    },
    'loggers': {
        '': {
            'handlers': ['console'],
            'level': 'INFO',
            'propagate': False
        }
    },
    'root': {
        'level': 'INFO',
        'handlers': ['console']
    }
}

try:
    config_path = os.environ['{}_CONFIG'.format(PROJECT_NAME.upper())]
except KeyError:
    env = os.environ.get('ENV', 'test')
    config_path = '{}/{}-{}.json'.format(PROJECT_CONFIG_PATH, PROJECT_NAME, env)

config = {}
config_path = os.path.expanduser(config_path)
if not os.path.exists(config_path):
    logging.warning('No \'{}\' configuration file found'.format(config_path))
else:
    with open(os.path.expanduser(config_path)) as fp:
        config = json.load(fp)

def get(key, default):
    return config.get(key, default)

DEBUG = get('DEBUG', True)
REQUEST_LIMIT = get('REQUEST_LIMIT', 25)

logging_config = get('LOGGING', DEFAULT_LOGGING_CONFIG)
logging.config.dictConfig(logging_config)

SENTRY_DSN = get('SENTRY_DSN', None)
<commit_msg>Drop max requests per provider to 10<commit_after>import os
import json
import logging
import logging.config

PROJECT_NAME = 'waterbutler'
PROJECT_CONFIG_PATH = '~/.cos'

try:
    import colorlog  # noqa
    DEFAULT_FORMATTER = {
        '()': 'colorlog.ColoredFormatter',
        'format': '%(cyan)s[%(asctime)s]%(log_color)s[%(levelname)s][%(name)s]: %(reset)s%(message)s'
    }
except ImportError:
    DEFAULT_FORMATTER = {
        '()': 'waterbutler.core.logging.MaskFormatter',
        'format': '[%(asctime)s][%(levelname)s][%(name)s]: %(message)s',
        'pattern': '(?<=cookie=)(.*?)(?=&|$)',
        'mask': '***'
    }
DEFAULT_LOGGING_CONFIG = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'console': DEFAULT_FORMATTER,
    },
    'handlers': {
        'console': {
            'class': 'logging.StreamHandler',
            'level': 'INFO',
            'formatter': 'console'
        },
        'syslog': {
            'class': 'logging.handlers.SysLogHandler',
            'level': 'INFO'
        }
    },
    'loggers': {
        '': {
            'handlers': ['console'],
            'level': 'INFO',
            'propagate': False
        }
    },
    'root': {
        'level': 'INFO',
        'handlers': ['console']
    }
}

try:
    config_path = os.environ['{}_CONFIG'.format(PROJECT_NAME.upper())]
except KeyError:
    env = os.environ.get('ENV', 'test')
    config_path = '{}/{}-{}.json'.format(PROJECT_CONFIG_PATH, PROJECT_NAME, env)

config = {}
config_path = os.path.expanduser(config_path)
if not os.path.exists(config_path):
    logging.warning('No \'{}\' configuration file found'.format(config_path))
else:
    with open(os.path.expanduser(config_path)) as fp:
        config = json.load(fp)

def get(key, default):
    return config.get(key, default)

DEBUG = get('DEBUG', True)
REQUEST_LIMIT = get('REQUEST_LIMIT', 10)

logging_config = get('LOGGING', DEFAULT_LOGGING_CONFIG)
logging.config.dictConfig(logging_config)

SENTRY_DSN = get('SENTRY_DSN', None)
<|endoftext|>


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

"""
Copyright(C) 2010  Romain Bignon

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

"""

try:
    from xml.etree import cElementTree as ElementTree
except ImportError:
    from xml.etree import ElementTree

try:
    from elementtidy import TidyHTMLTreeBuilder
    TidyHTMLTreeBuilder.ElementTree = ElementTree # force cElementTree if using it.
    HTMLTreeBuilder = TidyHTMLTreeBuilder.TidyHTMLTreeBuilder
except ImportError:
    from HTMLParser import HTMLParser

    class HTMLTreeBuilder(HTMLParser):
        def __init__(self, encoding=None):
            HTMLParser.__init__(self)
            self._target = ElementTree.TreeBuilder()

        def doctype(self, name, pubid, system):
            pass

        def close(self):
            tree = self._target.close()
            return tree

        def handle_starttag(self, tag, attrs):
            self._target.start(tag, dict(attrs))

        def handle_startendtag(self, tag, attrs):
            self._target.start(tag, dict(attrs))
            self._target.end(tag)

        def handle_data(self, data):
            self._target.data(data)

        def handle_endtag(self, tag):
            self._target.end(tag)

class StandardParser(object):
    def parse(self, data, encoding=None):
        parser = HTMLTreeBuilder(encoding)
        tree = ElementTree.parse(data, parser)

        for elem in tree.getiterator():
            if elem.tag.startswith('{'):
                elem.tag = elem.tag[elem.tag.find('}')+1:]
        return tree

def tostring(element):
    e = ElementTree.Element('body')
    e.text = element.text
    e.tail = element.tail
    for sub in element.getchildren():
        e.append(sub)
    s = ElementTree.tostring(e, 'utf-8')
    return unicode(s)
<commit_msg>while elementtidy segfaults, force using internal HTMLTreeBuilder<commit_after># -*- coding: utf-8 -*-

"""
Copyright(C) 2010  Romain Bignon

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

"""

try:
    from xml.etree import cElementTree as ElementTree
except ImportError:
    from xml.etree import ElementTree

try:
    raise ImportError
    from elementtidy import TidyHTMLTreeBuilder
    TidyHTMLTreeBuilder.ElementTree = ElementTree # force cElementTree if using it.
    HTMLTreeBuilder = TidyHTMLTreeBuilder.TidyHTMLTreeBuilder
except ImportError:
    from HTMLParser import HTMLParser

    class HTMLTreeBuilder(HTMLParser):
        def __init__(self, encoding=None):
            HTMLParser.__init__(self)
            self._target = ElementTree.TreeBuilder()

        def doctype(self, name, pubid, system):
            pass

        def close(self):
            tree = self._target.close()
            return tree

        def handle_starttag(self, tag, attrs):
            self._target.start(tag, dict(attrs))

        def handle_startendtag(self, tag, attrs):
            self._target.start(tag, dict(attrs))
            self._target.end(tag)

        def handle_data(self, data):
            self._target.data(data)

        def handle_endtag(self, tag):
            self._target.end(tag)

class StandardParser(object):
    def parse(self, data, encoding=None):
        parser = HTMLTreeBuilder(encoding)
        tree = ElementTree.parse(data, parser)

        for elem in tree.getiterator():
            if elem.tag.startswith('{'):
                elem.tag = elem.tag[elem.tag.find('}')+1:]
        return tree

def tostring(element):
    e = ElementTree.Element('body')
    e.text = element.text
    e.tail = element.tail
    for sub in element.getchildren():
        e.append(sub)
    s = ElementTree.tostring(e, 'utf-8')
    return unicode(s)
<|endoftext|>


<commit_before>from datetime import datetime

import numpy
from PIL import Image, ImageFilter

import imagehash

def find_region(pixels, segmented_pixels, segment_threshold_ratio, segmentation_img_size, from_highest=True):
    in_region = set()
    not_in_region = set()
    highest = numpy.unravel_index(numpy.nanargmax(pixels, axis=None), pixels.shape)
    lowest = numpy.unravel_index(numpy.nanargmin(pixels, axis=None), pixels.shape)
    peak = pixels[highest]
    trough = pixels[lowest]
    if from_highest:
        threshold = (peak - ((peak - trough) * segment_threshold_ratio))
        in_region.add(highest)
    else:
        threshold = (trough + ((peak - trough) * segment_threshold_ratio))
        in_region.add(lowest)
    new_pixels = in_region.copy()
    while True:
        try_next = set()
        # Find surrounding pixels
        for pixel in new_pixels:
            x, y = pixel
            new_pixels = []
            if x > 0:
                new_pixels.append((x - 1, y))
            if x < segmentation_img_size - 1:
                new_pixels.append((x + 1, y))
            if y > 0:
                new_pixels.append((x, y - 1))
            if y < segmentation_img_size - 1:
                new_pixels.append((x, y + 1))
            try_next.update(new_pixels)
        try_next -= in_region
        try_next -= not_in_region
        try_next -= segmented_pixels
        if not try_next:
            break
        # Check those
        if from_highest:
            upper, lower = peak, threshold
        else:
            upper, lower = threshold, trough
        new_pixels = set()
        for pixel in try_next:
            if upper > pixels[pixel] > lower:
                in_region.add(pixel)
                new_pixels.add(pixel)
            else:
                not_in_region.add(pixel)
    return in_region

def watershed_hash(image):
    # Define some tunable params
    segmentation_img_size = 300
    gaussian_blur = 2
    median_filter = 3
    segment_threshold_ratio = 0.3
    segmentation_required = 0.5
    min_segment_size = 100
    hash_func = imagehash.dhash

    orig_image = image.copy()
    orig_w, orig_h = orig_image.size
    # Convert to gray scale and resize
    image = image.convert("L").resize((segmentation_img_size, segmentation_img_size), Image.ANTIALIAS)
    # Add filters
    image = image.filter(ImageFilter.GaussianBlur(gaussian_blur)).filter(ImageFilter.MedianFilter(median_filter))
    pixels = numpy.array(image).astype(numpy.float32)

    segments = []
    already_segmented = set()
    while len(already_segmented) < segmentation_img_size * segmentation_img_size * segmentation_required:
        segment = find_region(pixels, already_segmented, segment_threshold_ratio, segmentation_img_size, True)
        # Apply segment
        if len(segment) > min_segment_size:
            segments.append(segment)
        for pix in segment:
            pixels[pix] = numpy.NaN
        already_segmented.update(segment)
        segment = find_region(pixels, already_segmented, segment_threshold_ratio, segmentation_img_size, False)
        # Apply segment
        if len(segment) > min_segment_size:
            segments.append(segment)
        for pix in segment:
            pixels[pix] = numpy.NaN
        already_segmented.update(segment)

    # Create bounding box for each segment
    hashes = []
    for segment in segments:
        scale_w = 1  # orig_w / segmentation_img_size
        scale_h = 1  # orig_h / segmentation_img_size
        min_y = min(coord[0] for coord in segment) * scale_h
        min_x = min(coord[1] for coord in segment) * scale_w
        max_y = (max(coord[0] for coord in segment)+1) * scale_h
        max_x = (max(coord[1] for coord in segment)+1) * scale_w
        # Compute robust hash for each bounding box
        bounding_box = image.crop((min_x, min_y, max_x, max_y))
        hashes.append(hash_func(bounding_box))
        # Show bounding box
        #show_segments(image, segment)
        #bounding_box.show()

    return imagehash.CropResistantHash(hashes)

def show_segments(image, segment):
    im_segment = image.copy()
    for pix in segment:
        im_segment.putpixel(pix[::-1], 255)
    im_segment.show()

if __name__ == "__main__":
    start_time = datetime.now()
    img = Image.open("alpacabread_Lorg.png")
    result = watershed_hash(img)
    print(datetime.now() - start_time)

    start_time = datetime.now()
    img2 = Image.open("alpacabread_Lorg-crop.png")
    result2 = watershed_hash(img2)
    print(datetime.now() - start_time)
    print(result)
    print(result2)
    print(result == result2)
    print(result.matches(result2, 10))
    print("Done")
<commit_msg>Swapping to have a singular segmentation threshold, instead of an upper and lower, and allowing some pixels to be unsegmented<commit_after>from datetime import datetime

import numpy
from PIL import Image, ImageFilter

import imagehash

def find_region(pixels, segmented_pixels, segment_threshold, segmentation_img_size, from_highest=True):
    in_region = set()
    not_in_region = set()
    highest = numpy.unravel_index(numpy.nanargmax(pixels, axis=None), pixels.shape)
    lowest = numpy.unravel_index(numpy.nanargmin(pixels, axis=None), pixels.shape)
    peak = pixels[highest]
    trough = pixels[lowest]
    if from_highest:
        in_region.add(highest)
    else:
        in_region.add(lowest)
    new_pixels = in_region.copy()
    while True:
        try_next = set()
        # Find surrounding pixels
        for pixel in new_pixels:
            x, y = pixel
            new_pixels = []
            if x > 0:
                new_pixels.append((x - 1, y))
            if x < segmentation_img_size - 1:
                new_pixels.append((x + 1, y))
            if y > 0:
                new_pixels.append((x, y - 1))
            if y < segmentation_img_size - 1:
                new_pixels.append((x, y + 1))
            try_next.update(new_pixels)
        try_next -= in_region
        try_next -= not_in_region
        try_next -= segmented_pixels
        if not try_next:
            break
        # Check those
        if from_highest:
            upper, lower = peak, segment_threshold
        else:
            upper, lower = segment_threshold, trough
        new_pixels = set()
        for pixel in try_next:
            if upper > pixels[pixel] > lower:
                in_region.add(pixel)
                new_pixels.add(pixel)
            else:
                not_in_region.add(pixel)
    return in_region

def watershed_hash(image):
    # Define some tunable params
    segmentation_img_size = 300
    gaussian_blur = 2
    median_filter = 3
    segment_threshold = 128
    min_segment_size = 100
    hash_func = imagehash.dhash

    orig_image = image.copy()
    # Convert to gray scale and resize
    image = image.convert("L").resize((segmentation_img_size, segmentation_img_size), Image.ANTIALIAS)
    # Add filters
    image = image.filter(ImageFilter.GaussianBlur(gaussian_blur)).filter(ImageFilter.MedianFilter(median_filter))
    pixels = numpy.array(image).astype(numpy.float32)

    segments = []
    already_segmented = set()
    while numpy.nanmax(pixels) > segment_threshold:
        segment = find_region(pixels, already_segmented, segment_threshold, segmentation_img_size, True)
        # Apply segment
        if len(segment) > min_segment_size:
            segments.append(segment)
        for pix in segment:
            pixels[pix] = numpy.NaN
        already_segmented.update(segment)

    while len(already_segmented) < segmentation_img_size * segmentation_img_size:
        segment = find_region(pixels, already_segmented, segment_threshold, segmentation_img_size, False)
        # Apply segment
        if len(segment) > min_segment_size:
            segments.append(segment)
        for pix in segment:
            pixels[pix] = numpy.NaN
        already_segmented.update(segment)

    # Create bounding box for each segment
    hashes = []
    for segment in segments:
        # orig_w, orig_h = orig_image.size
        scale_w = 1  # orig_w / segmentation_img_size
        scale_h = 1  # orig_h / segmentation_img_size
        min_y = min(coord[0] for coord in segment) * scale_h
        min_x = min(coord[1] for coord in segment) * scale_w
        max_y = (max(coord[0] for coord in segment)+1) * scale_h
        max_x = (max(coord[1] for coord in segment)+1) * scale_w
        # Compute robust hash for each bounding box
        bounding_box = image.crop((min_x, min_y, max_x, max_y))
        hashes.append(hash_func(bounding_box))
        # Show bounding box
        #show_segments(image, segment)
        #bounding_box.show()

    return imagehash.CropResistantHash(hashes)

def show_segments(image, segment):
    im_segment = image.copy()
    for pix in segment:
        im_segment.putpixel(pix[::-1], 255)
    im_segment.show()

if __name__ == "__main__":
    start_time = datetime.now()
    img = Image.open("alpacabread_Lorg.png")
    result = watershed_hash(img)
    print(datetime.now() - start_time)

    start_time = datetime.now()
    img2 = Image.open("alpacabread_Lorg-crop.png")
    result2 = watershed_hash(img2)
    print(datetime.now() - start_time)
    print(result)
    print(result2)
    print(result == result2)
    print(result.matches(result2, 10))
    print("Done")
<|endoftext|>


<commit_before>"""Connect DataTables server side to a MongoDB database.

Supports column sorting and filtering by search terms.  Also supports custom filtering in case you want to manipulate
the search server side, for example send a value like days_til_expiration=60 and convert to Mongo search like
{'ExpiryDate': {'$gt': ts, '$lt': ds}} where ts is today's date like 2017-09-12 and ds is today's date plus 60 days.
"""
import json
from bson.objectid import ObjectId

class DataTables(object):
    def __init__(self, pymongo_object, collection, request_args, **custom_filter):
        """

        :param pymongo_object: The PyMongo object representing the connection to a Mongo instance.
        :param collection: The Mongo collection
        :param request_args: The Flask request args, from request.args.to_dict()
        :param custom_filter: kwargs to be used as a custom Mongo filter, like key=value
        """

        self.mongo = pymongo_object
        self.collection = collection
        self.request_args = request_args
        self.custom_filter = custom_filter

    @property
    def db(self):
        return self.mongo.db

    @property
    def search_terms(self):
        return str(self.request_args.get("search")["value"]).split()

    @property
    def search_terms_without_a_colon(self):
        return [term for term in self.search_terms if ":" not in term]

    @property
    def search_terms_with_a_colon(self):
        return [term for term in self.search_terms if term.count(":") == 1]

    @property
    def requested_columns(self):
        return [column["data"] for column in self.request_args.get("columns")]

    @property
    def draw(self):
        return self.request_args.get("draw")

    @property
    def start(self):
        return self.request_args.get("start")

    @property
    def limit(self):
        _length = self.request_args.get("length")
        if _length == -1:
            return None
        return _length

    @property
    def cardinality(self):
        return self.db[self.collection].count()

    @property
    def cardinality_filtered(self):
        return self.db[self.collection].find(self.filter).count()

    @property
    def order_dir(self):
        """
        Return '1' for 'asc' or '-1' for 'desc'
        :return:
        """
        _dir = self.request_args.get("order")[0]["dir"]
        _MONGO_ORDER = {'asc': 1, 'desc': -1}
        return _MONGO_ORDER[_dir]

    @property
    def order_column(self):
        """DataTables provides the index of the order column, but Mongo .sort wants its name.

        :return:
        """
        _order_col = self.request_args.get("order")[0]["column"]
        return self.requested_columns[_order_col]

    @property
    def projection(self):
        p = {}
        for key in self.requested_columns:
            p.update({key: {'$ifNull': ['$' + key, '']}})
        return p

    def search_specific_key(self):
        """Search specific keys (columns) like 'key:value'.

        :return:
        """
        _col_specific_search = {}
        for term in self.search_terms_with_a_colon:
            col, term = term.split(':')
            _col_specific_search.update({col: {'$regex': term, '$options': 'i'}})

        return _col_specific_search

    def search_query(self):
        """Build the MongoDB query, searching every column for every term (case insensitive regex).

        :return:
        """
        # D3
        _search_query = {}
        if self.search_terms_without_a_colon:
            # L2
            and_filter_on_all_terms = []
            for term in self.search_terms_without_a_colon:
                # D2
                or_filter = {}
                # L1
                or_filter_on_all_columns = []
                for column in self.requested_columns:
                    # D1
                    column_filter = {
                        column: {'$regex': term, '$options': 'i'}
                    }
                    or_filter_on_all_columns.append(column_filter)
                or_filter['$or'] = or_filter_on_all_columns
                and_filter_on_all_terms.append(or_filter)

            _search_query['$and'] = and_filter_on_all_terms

        return _search_query

    @property
    def filter(self):
        _filter = {}
        _filter.update(self.custom_filter)
        _filter.update(self.search_query())
        _filter.update(self.search_specific_key())
        return _filter

    def results(self):

        _agg = [
                {'$match': self.filter},
                {'$sort': {self.order_column: self.order_dir}},
                {'$skip': self.start},
                {'$project': self.projection}
            ]

        if self.limit:
            _agg.append({'$limit': self.limit})

        _results = list(self.db[self.collection].aggregate(_agg))

        processed_results = []
        for result in _results:
            result = dict(result)
            result["DT_RowId"] = str(result.pop('_id'))  # rename the _id and convert ObjectId to str

            # go through every val in result and try to json.dumps objects and arrays - skip this if strings are okay
            for key, val in result.items():
                if type(val) not in [str, int, ObjectId]:
                    result[key] = json.dumps(val)

            processed_results.append(result)

        return processed_results

    def get_rows(self):
        return {
            'recordsTotal': str(self.cardinality),
            'recordsFiltered': str(self.cardinality_filtered),
            'draw': int(str(self.draw)),  # cast draw as integer to prevent XSS
            'data': self.results()
        }
<commit_msg>revert docstring mistake<commit_after>"""Connect DataTables server side to a MongoDB database.

Supports column sorting and filtering by search terms.  Also supports custom filtering in case you want to manipulate
the search server side, for example send a value like days_til_expiration=60 and convert to Mongo search like
{'ExpiryDate': {'$gt': ts, '$lt': ds}} where ts is today's date like 2017-09-12 and ds is today's date plus 60 days.
"""
import json
from bson.objectid import ObjectId

class DataTables(object):
    def __init__(self, pymongo_object, collection, request_args, **custom_filter):
        """

        :param pymongo_object: The PyMongo object representing the connection to a Mongo instance.
        :param collection: The Mongo collection
        :param request_args: The args from DataTables, passed as Flask request.values.get('args')
        :param custom_filter: kwargs to be used as a custom Mongo filter, like key=value
        """

        self.mongo = pymongo_object
        self.collection = collection
        self.request_args = request_args
        self.custom_filter = custom_filter

    @property
    def db(self):
        return self.mongo.db

    @property
    def search_terms(self):
        return str(self.request_args.get("search")["value"]).split()

    @property
    def search_terms_without_a_colon(self):
        return [term for term in self.search_terms if ":" not in term]

    @property
    def search_terms_with_a_colon(self):
        return [term for term in self.search_terms if term.count(":") == 1]

    @property
    def requested_columns(self):
        return [column["data"] for column in self.request_args.get("columns")]

    @property
    def draw(self):
        return self.request_args.get("draw")

    @property
    def start(self):
        return self.request_args.get("start")

    @property
    def limit(self):
        _length = self.request_args.get("length")
        if _length == -1:
            return None
        return _length

    @property
    def cardinality(self):
        return self.db[self.collection].count()

    @property
    def cardinality_filtered(self):
        return self.db[self.collection].find(self.filter).count()

    @property
    def order_dir(self):
        """
        Return '1' for 'asc' or '-1' for 'desc'
        :return:
        """
        _dir = self.request_args.get("order")[0]["dir"]
        _MONGO_ORDER = {'asc': 1, 'desc': -1}
        return _MONGO_ORDER[_dir]

    @property
    def order_column(self):
        """DataTables provides the index of the order column, but Mongo .sort wants its name.

        :return:
        """
        _order_col = self.request_args.get("order")[0]["column"]
        return self.requested_columns[_order_col]

    @property
    def projection(self):
        p = {}
        for key in self.requested_columns:
            p.update({key: {'$ifNull': ['$' + key, '']}})
        return p

    def search_specific_key(self):
        """Search specific keys (columns) like 'key:value'.

        :return:
        """
        _col_specific_search = {}
        for term in self.search_terms_with_a_colon:
            col, term = term.split(':')
            _col_specific_search.update({col: {'$regex': term, '$options': 'i'}})

        return _col_specific_search

    def search_query(self):
        """Build the MongoDB query, searching every column for every term (case insensitive regex).

        :return:
        """
        # D3
        _search_query = {}
        if self.search_terms_without_a_colon:
            # L2
            and_filter_on_all_terms = []
            for term in self.search_terms_without_a_colon:
                # D2
                or_filter = {}
                # L1
                or_filter_on_all_columns = []
                for column in self.requested_columns:
                    # D1
                    column_filter = {
                        column: {'$regex': term, '$options': 'i'}
                    }
                    or_filter_on_all_columns.append(column_filter)
                or_filter['$or'] = or_filter_on_all_columns
                and_filter_on_all_terms.append(or_filter)

            _search_query['$and'] = and_filter_on_all_terms

        return _search_query

    @property
    def filter(self):
        _filter = {}
        _filter.update(self.custom_filter)
        _filter.update(self.search_query())
        _filter.update(self.search_specific_key())
        return _filter

    def results(self):

        _agg = [
                {'$match': self.filter},
                {'$sort': {self.order_column: self.order_dir}},
                {'$skip': self.start},
                {'$project': self.projection}
            ]

        if self.limit:
            _agg.append({'$limit': self.limit})

        _results = list(self.db[self.collection].aggregate(_agg))

        processed_results = []
        for result in _results:
            result = dict(result)
            result["DT_RowId"] = str(result.pop('_id'))  # rename the _id and convert ObjectId to str

            # go through every val in result and try to json.dumps objects and arrays - skip this if strings are okay
            for key, val in result.items():
                if type(val) not in [str, int, ObjectId]:
                    result[key] = json.dumps(val)

            processed_results.append(result)

        return processed_results

    def get_rows(self):
        return {
            'recordsTotal': str(self.cardinality),
            'recordsFiltered': str(self.cardinality_filtered),
            'draw': int(str(self.draw)),  # cast draw as integer to prevent XSS
            'data': self.results()
        }
<|endoftext|>