<commit_before>import os
import re
import subprocess
import sys

from sgactions.utils import notify, alert
from sgfs import SGFS

PLAYABLE_EXTS = set(('.dpx', '.mov', '.jpg', '.jpeg', '.exr', '.mp4', '.aif'))

def bake_url(args):
    return 'rvlink://baked/%s' % ((' ' + ' '.join(args)).encode('hex'))

def rvlink(cmd):
    url = bake_url(cmd)
    if sys.platform == 'darwin':
        subprocess.check_call(['open', url])
    else:
        subprocess.check_call(['xdg-open', url])

def run_play(entity_type, selected_ids, **kwargs):

    sgfs = SGFS()
    entities = sgfs.session.get(entity_type, selected_ids)

    if entity_type == 'PublishEvent':
        sgfs.session.fetch(entities, ('code', 'sg_type', 'path_to_frames', 'path_to_movie', 'path', 'link.Task.entity'))
    elif entity_type == 'Version':
        sgfs.session.fetch(entities, ('code', 'path_to_frames', 'path_to_movie', 'entity'))
    else:
        alert('Cannot play %s entities in RV.' % entity_type)
        return

    chunks = []

    for entity in entities:

        skipped_exts = []

        for path_key in ('path_to_frames', 'path_to_movie', 'path'):

            path = (entity.get(path_key) or '').strip()
            if not path:
                continue

            ext = os.path.splitext(path)[1]
            if ext not in PLAYABLE_EXTS:
                skipped_exts.append(ext)
                continue

            notify('Opening %s in RV...' % path)

            # Convert any %04d into ####
            rv_path = re.sub(r'(?:%0?(\d)[sd])', lambda m: '#' * int(m.group(1)), path)

            # Go looking for audio.
            if entity_type == 'PublishEvent':
                shot = entity['link.Task.entity']
            else:
                shot = entity['entity']
            if shot:
                audio = sgfs.session.find_one('PublishEvent', [
                    ('sg_type', 'is', 'audio'),
                    ('link.Task.entity', 'is', shot),
                ], ['path'])
            else:
                audio = None

            # Open it (optionally with audio).
            if audio:
                chunks.extend(('[', rv_path, audio['path'], ']'))
            else:
                chunks.append(rv_path)
            break

        else:
            if skipped_exts:
                alert("""We don't know how to play %s %d ("%s") with %s extensions.""" % (
                    entity['sg_type'] + ' Publish' if entity_type == 'PublishEvent' else entity_type,
                    entity['id'],
                    entity['code'],
                    '/'.join(sorted(skipped_exts))
                ))
            else:
                alert("""%s %d ("%s") has nothing to play.""" % (
                    entity['sg_type'].title() + ' Publish' if entity_type == 'PublishEvent' else entity_type,
                    entity['id'],
                    entity['code'],
                ))
            return

    if chunks:
        # -l -> use lookahead cache
        # -play -> play immediately
        rvlink(['-l', '-play'] + chunks)
    else:
        alert("We don't know how to play %s entities." % entity_type)
<commit_msg>Remove autoplay for RV<commit_after>import os
import re
import subprocess
import sys

from sgactions.utils import notify, alert
from sgfs import SGFS

PLAYABLE_EXTS = set(('.dpx', '.mov', '.jpg', '.jpeg', '.exr', '.mp4', '.aif'))

def bake_url(args):
    return 'rvlink://baked/%s' % ((' ' + ' '.join(args)).encode('hex'))

def rvlink(cmd):
    url = bake_url(cmd)
    if sys.platform == 'darwin':
        subprocess.check_call(['open', url])
    else:
        subprocess.check_call(['xdg-open', url])

def run_play(entity_type, selected_ids, **kwargs):

    sgfs = SGFS()
    entities = sgfs.session.get(entity_type, selected_ids)

    if entity_type == 'PublishEvent':
        sgfs.session.fetch(entities, ('code', 'sg_type', 'path_to_frames', 'path_to_movie', 'path', 'link.Task.entity'))
    elif entity_type == 'Version':
        sgfs.session.fetch(entities, ('code', 'path_to_frames', 'path_to_movie', 'entity'))
    else:
        alert('Cannot play %s entities in RV.' % entity_type)
        return

    chunks = []

    for entity in entities:

        skipped_exts = []

        for path_key in ('path_to_frames', 'path_to_movie', 'path'):

            path = (entity.get(path_key) or '').strip()
            if not path:
                continue

            ext = os.path.splitext(path)[1]
            if ext not in PLAYABLE_EXTS:
                skipped_exts.append(ext)
                continue

            notify('Opening %s in RV...' % path)

            # Convert any %04d into ####
            rv_path = re.sub(r'(?:%0?(\d)[sd])', lambda m: '#' * int(m.group(1)), path)

            # Go looking for audio.
            if entity_type == 'PublishEvent':
                shot = entity['link.Task.entity']
            else:
                shot = entity['entity']
            if shot:
                audio = sgfs.session.find_one('PublishEvent', [
                    ('sg_type', 'is', 'audio'),
                    ('link.Task.entity', 'is', shot),
                ], ['path'])
            else:
                audio = None

            # Open it (optionally with audio).
            if audio:
                chunks.extend(('[', rv_path, audio['path'], ']'))
            else:
                chunks.append(rv_path)
            break

        else:
            if skipped_exts:
                alert("""We don't know how to play %s %d ("%s") with %s extensions.""" % (
                    entity['sg_type'] + ' Publish' if entity_type == 'PublishEvent' else entity_type,
                    entity['id'],
                    entity['code'],
                    '/'.join(sorted(skipped_exts))
                ))
            else:
                alert("""%s %d ("%s") has nothing to play.""" % (
                    entity['sg_type'].title() + ' Publish' if entity_type == 'PublishEvent' else entity_type,
                    entity['id'],
                    entity['code'],
                ))
            return

    if chunks:
        # -l -> use lookahead cache
        # -play -> play immediately
        rvlink(['-l'] + chunks)
    else:
        alert("We don't know how to play %s entities." % entity_type)
<|endoftext|>


<commit_before>import datetime
from furl import furl

from django import forms
from django.conf.urls import url
from django.contrib import admin
from django.contrib.admin import SimpleListFilter
from django.contrib.admin.widgets import AdminDateWidget
from django.http import HttpResponseRedirect
from django.template.response import TemplateResponse
from django.urls import reverse
from django.utils import timezone
from django.utils.html import format_html

from oauth2_provider.models import AccessToken

from share import tasks
from share.admin.celery import CeleryTaskResultAdmin
from share.admin.readonly import ReadOnlyAdmin
from share.admin.share_objects import CreativeWorkAdmin
from share.harvest.scheduler import HarvestScheduler
from share.models.banner import SiteBanner
from share.models.celery import CeleryTaskResult
from share.models.change import ChangeSet
from share.models.core import NormalizedData, ShareUser
from share.models.creative import AbstractCreativeWork
from share.models.ingest import RawDatum, Source, SourceConfig, Harvester, Transformer
from share.models.logs import HarvestLog
from share.models.registration import ProviderRegistration

admin.site.register(AbstractCreativeWork, CreativeWorkAdmin)
admin.site.register(CeleryTaskResult, CeleryTaskResultAdmin)

class NormalizedDataAdmin(admin.ModelAdmin):
    date_hierarchy = 'created_at'
    list_filter = ['source', ]
    raw_id_fields = ('raw', 'tasks',)

class ChangeSetSubmittedByFilter(SimpleListFilter):
    title = 'Source'
    parameter_name = 'source_id'

    def lookups(self, request, model_admin):
        return ShareUser.objects.filter(is_active=True).values_list('id', 'username')

    def queryset(self, request, queryset):
        if self.value():
            return queryset.filter(normalized_data__source_id=self.value())
        return queryset

class ChangeSetAdmin(admin.ModelAdmin):
    list_display = ('status_', 'count_changes', 'submitted_by', 'submitted_at')
    actions = ['accept_changes']
    list_filter = ['status', ChangeSetSubmittedByFilter]
    raw_id_fields = ('normalized_data',)

    def submitted_by(self, obj):
        return obj.normalized_data.source
    submitted_by.short_description = 'submitted by'

    def count_changes(self, obj):
        return obj.changes.count()
    count_changes.short_description = 'number of changes'

    def status_(self, obj):
        return ChangeSet.STATUS[obj.status].title()

class RawDatumAdmin(admin.ModelAdmin):
    raw_id_fields = ()

class AccessTokenAdmin(admin.ModelAdmin):
    list_display = ('token', 'user', 'scope')

class ProviderRegistrationAdmin(ReadOnlyAdmin):
    list_display = ('source_name', 'status_', 'submitted_at', 'submitted_by', 'direct_source')
    list_filter = ('direct_source', 'status',)
    readonly_fields = ('submitted_at', 'submitted_by',)

    def status_(self, obj):
        return ProviderRegistration.STATUS[obj.status].title()

class SiteBannerAdmin(admin.ModelAdmin):
    list_display = ('title', 'color', 'icon', 'active')
    list_editable = ('active',)
    ordering = ('-active', '-last_modified_at')
    readonly_fields = ('created_at', 'created_by', 'last_modified_at', 'last_modified_by')

    def save_model(self, request, obj, form, change):
        if not change:
            obj.created_by = request.user
        obj.last_modified_by = request.user
        super().save_model(request, obj, form, change)

class SourceConfigFilter(admin.SimpleListFilter):
    title = 'Source Config'
    parameter_name = 'source_config'

    def lookups(self, request, model_admin):
        # TODO make this into a cool hierarchy deal
        # return SourceConfig.objects.select_related('source').values_list('
        return SourceConfig.objects.order_by('label').values_list('id', 'label')

    def queryset(self, request, queryset):
        if self.value():
            return queryset.filter(source_config=self.value())

class HarvestLogAdmin(admin.ModelAdmin):
    list_display = ('id', 'source', 'label', 'share_version', 'status_', 'start_date_', 'end_date_', 'harvest_log_actions', )
    list_filter = ('status', SourceConfigFilter, )
    list_select_related = ('source_config__source', )
    readonly_fields = ('harvest_log_actions',)
    actions = ('restart_tasks', )

    STATUS_COLORS = {
        HarvestLog.STATUS.created: 'blue',
        HarvestLog.STATUS.started: 'cyan',
        HarvestLog.STATUS.failed: 'red',
        HarvestLog.STATUS.succeeded: 'green',
        HarvestLog.STATUS.rescheduled: 'goldenrod',
        HarvestLog.STATUS.forced: 'maroon',
        HarvestLog.STATUS.skipped: 'orange',
        HarvestLog.STATUS.retried: 'darkseagreen',
        HarvestLog.STATUS.cancelled: 'grey',
    }

    def source(self, obj):
        return obj.source_config.source.long_title

    def label(self, obj):
        return obj.source_config.label

    def start_date_(self, obj):
        return obj.start_date.isoformat()

    def end_date_(self, obj):
        return obj.end_date.isoformat()

    def status_(self, obj):
        return format_html(
            '{}',
            self.STATUS_COLORS[obj.status],
            HarvestLog.STATUS[obj.status].title(),
        )

    def restart_tasks(self, request, queryset):
        queryset.update(status=HarvestLog.STATUS.created)
    restart_tasks.short_description = 'Re-enqueue'

    def harvest_log_actions(self, obj):
        url = furl(reverse('admin:source-config-harvest', args=[obj.source_config_id]))
        url.args['start'] = self.start_date_(obj)
        url.args['end'] = self.end_date_(obj)
        url.args['superfluous'] = True
        return format_html('Restart', url.url)
    harvest_log_actions.short_description = 'Actions'

class HarvestForm(forms.Form):
    start = forms.DateField(widget=AdminDateWidget())
    end = forms.DateField(widget=AdminDateWidget())
    superfluous = forms.BooleanField(required=False)

    def clean(self):
        super().clean()
        if self.cleaned_data['start'] > self.cleaned_data['end']:
            raise forms.ValidationError('Start date cannot be after end date.')

class SourceConfigAdmin(admin.ModelAdmin):
    list_display = ('label', 'source_', 'version', 'enabled', 'source_config_actions')
    list_select_related = ('source',)
    readonly_fields = ('source_config_actions',)
    search_fields = ['label', 'source__name', 'source__long_title']

    def source_(self, obj):
        return obj.source.long_title

    def enabled(self, obj):
        return not obj.disabled
    enabled.boolean = True

    def get_urls(self):
        return [
            url(
                r'^(?P<config_id>.+)/harvest/$',
                self.admin_site.admin_view(self.harvest),
                name='source-config-harvest'
            )
        ] + super().get_urls()

    def source_config_actions(self, obj):
        if obj.harvester_id is None:
            return ''
        return format_html(
            'Harvest',
            reverse('admin:source-config-harvest', args=[obj.pk]),
        )
    source_config_actions.short_description = 'Actions'

    def harvest(self, request, config_id):
        config = self.get_object(request, config_id)
        if config.harvester_id is None:
            raise ValueError('You need a harvester to harvest.')

        if request.method == 'POST':
            form = HarvestForm(request.POST)
            if form.is_valid():
                for log in HarvestScheduler(config).range(form.cleaned_data['start'], form.cleaned_data['end']):
                    tasks.harvest.apply_async((), {'log_id': log.id, 'superfluous': form.cleaned_data['superfluous']})

                self.message_user(request, 'Started harvesting {}!'.format(config.label))
                url = reverse('admin:share_harvestlog_changelist', current_app=self.admin_site.name)
                return HttpResponseRedirect(url)
        else:
            initial = {'start': config.earliest_date, 'end': timezone.now().date()}
            for field in HarvestForm.base_fields.keys():
                if field in request.GET:
                    initial[field] = request.GET[field]
            form = HarvestForm(initial=initial)

        context = self.admin_site.each_context(request)
        context['opts'] = self.model._meta
        context['form'] = form
        context['source_config'] = config
        context['title'] = 'Harvest {}'.format(config.label)
        return TemplateResponse(request, 'admin/harvest.html', context)

class SourceAdminInline(admin.StackedInline):
    model = Source

class ShareUserAdmin(admin.ModelAdmin):
    inlines = (SourceAdminInline,)

class SourceAddForm(forms.ModelForm):

    title = forms.CharField(min_length=3, max_length=255, help_text='What this source will be displayed as to the end user. Must be unique.')
    url = forms.URLField(min_length=11, max_length=255, help_text='The home page or canonical URL for this source, make sure it is unique!.\nThe reverse DNS notation prepended with "sources." will be used to create a user for this source. IE share.osf.io -> sources.io.osf.share')

    class Meta:
        model = Source
        exclude = ('access_token', )
        fields = ('title', 'url', 'icon', )

    def validate_unique(self):
        # Forces validation checks on every field
        try:
            self.instance.validate_unique()
        except forms.ValidationError as e:
            # Translate field names because I'm a bad person
            if 'long_title' in e.error_dict:
                e.error_dict['title'] = e.error_dict.pop('long_title')
            if 'name' in e.error_dict:
                e.error_dict['url'] = e.error_dict.pop('name')
            if 'home_page' in e.error_dict:
                e.error_dict['url'] = e.error_dict.pop('home_page')
            self._update_errors(e)

    def _post_clean(self):
        if not self._errors:
            self.instance.home_page = self.cleaned_data['url'].lower().strip('/')
            self.instance.long_title = self.cleaned_data.pop('title')
            self.instance.name = '.'.join(reversed(self.instance.home_page.split('//')[1].split('.')))
        return super()._post_clean()

    def save(self, commit=True):
        instance = super().save(commit=False)

        user = ShareUser.objects.create_user(username='sources.' + instance.name, save=commit)
        user.set_unusable_password()
        instance.user = user

        if commit:
            instance.save()

        return instance

class SourceAdmin(admin.ModelAdmin):
    search_fields = ('name', 'long_title')
    readonly_fields = ('access_token', )

    def add_view(self, *args, **kwargs):
        self.form = SourceAddForm
        return super().add_view(*args, **kwargs)

    def save_model(self, request, obj, form, change):
        if obj.user and not (obj.user_id or obj.user.id):
            obj.user.save()
            obj.user_id = obj.user.id  # Django is weird
        obj.save()

    def access_token(self, obj):
        tokens = obj.user.accesstoken_set.all()
        if tokens:
            return tokens[0].token
        return None

admin.site.unregister(AccessToken)
admin.site.register(AccessToken, AccessTokenAdmin)

admin.site.register(ChangeSet, ChangeSetAdmin)
admin.site.register(HarvestLog, HarvestLogAdmin)
admin.site.register(NormalizedData, NormalizedDataAdmin)
admin.site.register(ProviderRegistration, ProviderRegistrationAdmin)
admin.site.register(RawDatum, RawDatumAdmin)
admin.site.register(SiteBanner, SiteBannerAdmin)

admin.site.register(Harvester)
admin.site.register(ShareUser, ShareUserAdmin)
admin.site.register(Source, SourceAdmin)
admin.site.register(SourceConfig, SourceConfigAdmin)
admin.site.register(Transformer)
<commit_msg>[Improvement] Update RawDatumAdmin<commit_after>from furl import furl

from django import forms
from django.conf.urls import url
from django.contrib import admin
from django.contrib.admin import SimpleListFilter
from django.contrib.admin.widgets import AdminDateWidget
from django.http import HttpResponseRedirect
from django.template.response import TemplateResponse
from django.urls import reverse
from django.utils import timezone
from django.utils.html import format_html

from oauth2_provider.models import AccessToken

from share import tasks
from share.admin.celery import CeleryTaskResultAdmin
from share.admin.readonly import ReadOnlyAdmin
from share.admin.share_objects import CreativeWorkAdmin
from share.harvest.scheduler import HarvestScheduler
from share.models.banner import SiteBanner
from share.models.celery import CeleryTaskResult
from share.models.change import ChangeSet
from share.models.core import NormalizedData, ShareUser
from share.models.creative import AbstractCreativeWork
from share.models.ingest import RawDatum, Source, SourceConfig, Harvester, Transformer
from share.models.logs import HarvestLog
from share.models.registration import ProviderRegistration

admin.site.register(AbstractCreativeWork, CreativeWorkAdmin)
admin.site.register(CeleryTaskResult, CeleryTaskResultAdmin)

class NormalizedDataAdmin(admin.ModelAdmin):
    date_hierarchy = 'created_at'
    list_filter = ['source', ]
    raw_id_fields = ('raw', 'tasks',)

class ChangeSetSubmittedByFilter(SimpleListFilter):
    title = 'Source'
    parameter_name = 'source_id'

    def lookups(self, request, model_admin):
        return ShareUser.objects.filter(is_active=True).values_list('id', 'username')

    def queryset(self, request, queryset):
        if self.value():
            return queryset.filter(normalized_data__source_id=self.value())
        return queryset

class ChangeSetAdmin(admin.ModelAdmin):
    list_display = ('status_', 'count_changes', 'submitted_by', 'submitted_at')
    actions = ['accept_changes']
    list_filter = ['status', ChangeSetSubmittedByFilter]
    raw_id_fields = ('normalized_data',)

    def submitted_by(self, obj):
        return obj.normalized_data.source
    submitted_by.short_description = 'submitted by'

    def count_changes(self, obj):
        return obj.changes.count()
    count_changes.short_description = 'number of changes'

    def status_(self, obj):
        return ChangeSet.STATUS[obj.status].title()

class RawDatumAdmin(admin.ModelAdmin):
    show_full_result_count = False
    list_select_related = ('suid__source_config', )
    list_display = ('id', 'identifier', 'source_config_label', 'datestamp', 'date_created', 'date_modified', )

    def identifier(self, obj):
        return obj.suid.identifier

    def source_config_label(self, obj):
        return obj.suid.source_config.label

class AccessTokenAdmin(admin.ModelAdmin):
    list_display = ('token', 'user', 'scope')

class ProviderRegistrationAdmin(ReadOnlyAdmin):
    list_display = ('source_name', 'status_', 'submitted_at', 'submitted_by', 'direct_source')
    list_filter = ('direct_source', 'status',)
    readonly_fields = ('submitted_at', 'submitted_by',)

    def status_(self, obj):
        return ProviderRegistration.STATUS[obj.status].title()

class SiteBannerAdmin(admin.ModelAdmin):
    list_display = ('title', 'color', 'icon', 'active')
    list_editable = ('active',)
    ordering = ('-active', '-last_modified_at')
    readonly_fields = ('created_at', 'created_by', 'last_modified_at', 'last_modified_by')

    def save_model(self, request, obj, form, change):
        if not change:
            obj.created_by = request.user
        obj.last_modified_by = request.user
        super().save_model(request, obj, form, change)

class SourceConfigFilter(admin.SimpleListFilter):
    title = 'Source Config'
    parameter_name = 'source_config'

    def lookups(self, request, model_admin):
        # TODO make this into a cool hierarchy deal
        # return SourceConfig.objects.select_related('source').values_list('
        return SourceConfig.objects.order_by('label').values_list('id', 'label')

    def queryset(self, request, queryset):
        if self.value():
            return queryset.filter(source_config=self.value())

class HarvestLogAdmin(admin.ModelAdmin):
    list_display = ('id', 'source', 'label', 'share_version', 'status_', 'start_date_', 'end_date_', 'harvest_log_actions', )
    list_filter = ('status', SourceConfigFilter, )
    list_select_related = ('source_config__source', )
    readonly_fields = ('harvest_log_actions',)
    actions = ('restart_tasks', )

    STATUS_COLORS = {
        HarvestLog.STATUS.created: 'blue',
        HarvestLog.STATUS.started: 'cyan',
        HarvestLog.STATUS.failed: 'red',
        HarvestLog.STATUS.succeeded: 'green',
        HarvestLog.STATUS.rescheduled: 'goldenrod',
        HarvestLog.STATUS.forced: 'maroon',
        HarvestLog.STATUS.skipped: 'orange',
        HarvestLog.STATUS.retried: 'darkseagreen',
        HarvestLog.STATUS.cancelled: 'grey',
    }

    def source(self, obj):
        return obj.source_config.source.long_title

    def label(self, obj):
        return obj.source_config.label

    def start_date_(self, obj):
        return obj.start_date.isoformat()

    def end_date_(self, obj):
        return obj.end_date.isoformat()

    def status_(self, obj):
        return format_html(
            '{}',
            self.STATUS_COLORS[obj.status],
            HarvestLog.STATUS[obj.status].title(),
        )

    def restart_tasks(self, request, queryset):
        queryset.update(status=HarvestLog.STATUS.created)
    restart_tasks.short_description = 'Re-enqueue'

    def harvest_log_actions(self, obj):
        url = furl(reverse('admin:source-config-harvest', args=[obj.source_config_id]))
        url.args['start'] = self.start_date_(obj)
        url.args['end'] = self.end_date_(obj)
        url.args['superfluous'] = True
        return format_html('Restart', url.url)
    harvest_log_actions.short_description = 'Actions'

class HarvestForm(forms.Form):
    start = forms.DateField(widget=AdminDateWidget())
    end = forms.DateField(widget=AdminDateWidget())
    superfluous = forms.BooleanField(required=False)

    def clean(self):
        super().clean()
        if self.cleaned_data['start'] > self.cleaned_data['end']:
            raise forms.ValidationError('Start date cannot be after end date.')

class SourceConfigAdmin(admin.ModelAdmin):
    list_display = ('label', 'source_', 'version', 'enabled', 'source_config_actions')
    list_select_related = ('source',)
    readonly_fields = ('source_config_actions',)
    search_fields = ['label', 'source__name', 'source__long_title']

    def source_(self, obj):
        return obj.source.long_title

    def enabled(self, obj):
        return not obj.disabled
    enabled.boolean = True

    def get_urls(self):
        return [
            url(
                r'^(?P<config_id>.+)/harvest/$',
                self.admin_site.admin_view(self.harvest),
                name='source-config-harvest'
            )
        ] + super().get_urls()

    def source_config_actions(self, obj):
        if obj.harvester_id is None:
            return ''
        return format_html(
            'Harvest',
            reverse('admin:source-config-harvest', args=[obj.pk]),
        )
    source_config_actions.short_description = 'Actions'

    def harvest(self, request, config_id):
        config = self.get_object(request, config_id)
        if config.harvester_id is None:
            raise ValueError('You need a harvester to harvest.')

        if request.method == 'POST':
            form = HarvestForm(request.POST)
            if form.is_valid():
                for log in HarvestScheduler(config).range(form.cleaned_data['start'], form.cleaned_data['end']):
                    tasks.harvest.apply_async((), {'log_id': log.id, 'superfluous': form.cleaned_data['superfluous']})

                self.message_user(request, 'Started harvesting {}!'.format(config.label))
                url = reverse('admin:share_harvestlog_changelist', current_app=self.admin_site.name)
                return HttpResponseRedirect(url)
        else:
            initial = {'start': config.earliest_date, 'end': timezone.now().date()}
            for field in HarvestForm.base_fields.keys():
                if field in request.GET:
                    initial[field] = request.GET[field]
            form = HarvestForm(initial=initial)

        context = self.admin_site.each_context(request)
        context['opts'] = self.model._meta
        context['form'] = form
        context['source_config'] = config
        context['title'] = 'Harvest {}'.format(config.label)
        return TemplateResponse(request, 'admin/harvest.html', context)

class SourceAdminInline(admin.StackedInline):
    model = Source

class ShareUserAdmin(admin.ModelAdmin):
    inlines = (SourceAdminInline,)

class SourceAddForm(forms.ModelForm):

    title = forms.CharField(min_length=3, max_length=255, help_text='What this source will be displayed as to the end user. Must be unique.')
    url = forms.URLField(min_length=11, max_length=255, help_text='The home page or canonical URL for this source, make sure it is unique!.\nThe reverse DNS notation prepended with "sources." will be used to create a user for this source. IE share.osf.io -> sources.io.osf.share')

    class Meta:
        model = Source
        exclude = ('access_token', )
        fields = ('title', 'url', 'icon', )

    def validate_unique(self):
        # Forces validation checks on every field
        try:
            self.instance.validate_unique()
        except forms.ValidationError as e:
            # Translate field names because I'm a bad person
            if 'long_title' in e.error_dict:
                e.error_dict['title'] = e.error_dict.pop('long_title')
            if 'name' in e.error_dict:
                e.error_dict['url'] = e.error_dict.pop('name')
            if 'home_page' in e.error_dict:
                e.error_dict['url'] = e.error_dict.pop('home_page')
            self._update_errors(e)

    def _post_clean(self):
        if not self._errors:
            self.instance.home_page = self.cleaned_data['url'].lower().strip('/')
            self.instance.long_title = self.cleaned_data.pop('title')
            self.instance.name = '.'.join(reversed(self.instance.home_page.split('//')[1].split('.')))
        return super()._post_clean()

    def save(self, commit=True):
        instance = super().save(commit=False)

        user = ShareUser.objects.create_user(username='sources.' + instance.name, save=commit)
        user.set_unusable_password()
        instance.user = user

        if commit:
            instance.save()

        return instance

class SourceAdmin(admin.ModelAdmin):
    search_fields = ('name', 'long_title')
    readonly_fields = ('access_token', )

    def add_view(self, *args, **kwargs):
        self.form = SourceAddForm
        return super().add_view(*args, **kwargs)

    def save_model(self, request, obj, form, change):
        if obj.user and not (obj.user_id or obj.user.id):
            obj.user.save()
            obj.user_id = obj.user.id  # Django is weird
        obj.save()

    def access_token(self, obj):
        tokens = obj.user.accesstoken_set.all()
        if tokens:
            return tokens[0].token
        return None

admin.site.unregister(AccessToken)
admin.site.register(AccessToken, AccessTokenAdmin)

admin.site.register(ChangeSet, ChangeSetAdmin)
admin.site.register(HarvestLog, HarvestLogAdmin)
admin.site.register(NormalizedData, NormalizedDataAdmin)
admin.site.register(ProviderRegistration, ProviderRegistrationAdmin)
admin.site.register(RawDatum, RawDatumAdmin)
admin.site.register(SiteBanner, SiteBannerAdmin)

admin.site.register(Harvester)
admin.site.register(ShareUser, ShareUserAdmin)
admin.site.register(Source, SourceAdmin)
admin.site.register(SourceConfig, SourceConfigAdmin)
admin.site.register(Transformer)
<|endoftext|>


<commit_before>from __future__ import unicode_literals
import sys
import json
import time
import hashlib
import logging
import traceback
from copy import deepcopy
from functools import wraps
from threading import local, RLock
from collections import defaultdict

import cherrypy

from ws4py.websocket import WebSocket
from ws4py.server.cherrypyserver import WebSocketPlugin, WebSocketTool

from sideboard import Caller

import sideboard.lib
from sideboard.config import config

class threadlocal(object):
    _threadlocal = local()

    @classmethod
    def get(cls, key, default=None):
        return getattr(cls._threadlocal, key, default)

    @classmethod
    def set(cls, key, val):
        return setattr(cls._threadlocal, key, val)

    @classmethod
    def clear(cls):
        cls._threadlocal.__dict__.clear()

    @classmethod
    def get_client(cls):
        return cls.get("client") or cls.get("message", {}).get("client")

    @classmethod
    def reset(cls, **kwargs):
        cls.clear()
        for key, val in kwargs.items():
            cls.set(key, val)

def _normalize_topics(*topics):
    """
    Converts a list of types, strings, or whatever else into a list of strings.

    >>> _normalize_topics()
    []

    >>> _normalize_topics('')
    []

    >>> _normalize_topics('   ')
    []

    >>> _normalize_topics(None)
    []

    >>> _normalize_topics('topic-one', 'topic-two')
    ['topic-one', 'topic-two']

    >>> _normalize_topics('repeated-topic', 'repeated-topic')
    ['repeated-topic']

    >>> _normalize_topics('topic-left', None, 'topic-right', None)
    ['topic-left', 'topic-right']

    >>> _normalize_topics('', 'topic-left', '', 'topic-right')
    ['topic-left', 'topic-right']

    >>> _normalize_topics('   ', '   topic-padded-left', '   ', 'topic-padded-right   ', '   ')
    ['topic-padded-left', 'topic-padded-right']

    >>> _normalize_topics(type)
    ['type']

    >>> _normalize_topics(dict)
    ['dict']

    >>> _normalize_topics(dict(foo="bar"))
    ["{'foo': 'bar'}"]
    """
    normalized_topics = []
    for topic in topics:
        if topic is not None:
            if isinstance(topic, type):
                normalized_topics.append(topic.__name__)
            elif isinstance(topic, basestring):
                topic = topic.strip()
                if topic != '':
                    normalized_topics.append(topic)
            else:
                normalized_topics.append(str(topic))
    return list(set(normalized_topics))

def notify(topics, trigger="manual", delay=0, originating_client=None):
    broadcaster.delayed(delay, trigger, _normalize_topics(*sideboard.lib.listify(topics)),
                        originating_client or threadlocal.get_client())

def notifies(*args, **kwargs):
    """
    Adds a notifies attribute to the decorated function. The notifies
    attribute specifies a list of the topics which the function notifies.

    >>> @notifies('topic-one', 'topic-two')
    ... def fn():
    ...     pass
    >>> getattr(fn, 'notifies')
    ['topic-one', 'topic-two']

    >>> @notifies(dict)
    ... def fn_dict():
    ...     pass
    >>> getattr(fn_dict, 'notifies')
    ['dict']
    """
    delay = kwargs.pop("delay", 0)
    topics = _normalize_topics(*args)

    def decorated_func(func):
        @wraps(func)
        def notifier_func(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            finally:
                notify(topics, trigger=func.__name__, delay=delay)

        notifier_func.notifies = topics
        return notifier_func

    return decorated_func

def subscribes(*args):
    """
    Adds a subscribes attribute to the decorated function. The subscribes
    attribute specifies a list of the topics to which the function subscribes.

    >>> @subscribes('topic-one', 'topic-two')
    ... def fn():
    ...     pass
    >>> getattr(fn, 'subscribes')
    ['topic-one', 'topic-two']

    >>> @subscribes(dict)
    ... def fn_dict():
    ...     pass
    >>> getattr(fn_dict, 'subscribes')
    ['dict']
    """
    topics = _normalize_topics(*args)

    def decorated_func(func):
        func.subscribes = topics
        return func

    return decorated_func

def _fingerprint(x):
    """
    Calculates the md5 sum of the given argument.

    If _fingerprint is passed a string, it calculates the md5 sum of the
    string. If _fingerprint is passed anything else, then the json encoding
    of the argument is used to calculate the md5 sum.

    >>> _fingerprint(None)
    '37a6259cc0c1dae299a7866489dff0bd'

    >>> _fingerprint('test')
    '098f6bcd4621d373cade4e832627b4f6'

    >>> _fingerprint({'key':'value'})
    'a7353f7cddce808de0032747a0b7be50'

    >>> _fingerprint(dict(key='value'))
    'a7353f7cddce808de0032747a0b7be50'

    >>> _fingerprint({'a':1, 'b':2})
    '608de49a4600dbb5b173492759792e4a'

    >>> _fingerprint({'b':2, 'a':1})
    '608de49a4600dbb5b173492759792e4a'

    >>> _fingerprint({'a':{'x':3, 'y':4}, 'b':2})
    '2c22e445e9278c66dd7ea78b757defe6'

    >>> _fingerprint({'b':2, 'a':{'y':4, 'x':3}})
    '2c22e445e9278c66dd7ea78b757defe6'
    """
    md5 = hashlib.md5()
    if isinstance(x, basestring):
        md5.update(x)
    else:
        md5.update(json.dumps(x, cls=sideboard.lib.serializer, sort_keys=True, separators=(',', ':')))
    return md5.hexdigest()

def get_params(params):
    if params is None:
        return [], {}
    elif isinstance(params, dict):
        return [], params
    elif isinstance(params, (list, tuple)):
        return params, {}
    else:
        return [params], {}

class WebSocketDispatcher(WebSocket):
    ERROR = object()
    username = None
    subscriptions = defaultdict(lambda: defaultdict(lambda: defaultdict(set)))

    debug = True
    services = {}
    log = logging.getLogger(__name__)

    def __init__(self, *args, **kwargs):
        WebSocket.__init__(self, *args, **kwargs)
        self.send_lock = RLock()
        self.client_locks, self.cached_queries, self.cached_fingerprints = \
            defaultdict(RLock), defaultdict(dict), defaultdict(dict)
        self.username = self.check_authentication()

    @classmethod
    def check_authentication(cls):
        return cherrypy.session["username"]

    @classmethod
    def get_all_subscribed(cls):
        websockets = set()
        for model, subscriptions in cls.subscriptions.items():
            for websocket, clients in subscriptions.items():
                websockets.add(websocket)
        return websockets

    @classmethod
    def get_all(cls):
        return websocket_plugin.pool.keys()[:]

    @classmethod
    def close_all(cls):
        for ws in cls.get_all():
            try:
                ws.close()
            except:
                pass

    def send(self, **message):
        message = {k: v for k, v in message.items() if v is not None}
        if "data" in message and "client" in message:
            fingerprint = _fingerprint(message["data"])
            client, callback = message["client"], message.get("callback")
            repeat_send = callback in self.cached_fingerprints[client]
            cached_fingerprint = self.cached_fingerprints[client].get(callback)
            self.cached_fingerprints[client][callback] = fingerprint
            if cached_fingerprint == fingerprint and repeat_send:
                return

        self.log.debug("sending {}", message)
        message = json.dumps(message, cls=sideboard.lib.serializer, separators=(",", ":"), sort_keys=True)
        with self.send_lock:
            WebSocket.send(self, message)

    def closed(self, code, reason=""):
        self.log.info("closing: code={!r} reason={!r}", code, reason)
        self.unsubscribe_all()
        WebSocket.closed(self, code, reason)
        notify("Site")

    @classmethod
    def broadcast(cls, trigger, models, originating_client=None):
        triggered = set()
        for model in models:
            for websocket, clients in cls.subscriptions[model].items():
                for client, callbacks in clients.copy().items():
                    if client != originating_client:
                        for callback in callbacks:
                            triggered.add((websocket, client, callback))

        for websocket, client, callback in triggered:
            try:
                websocket.trigger(trigger, client, callback)
            except:
                cls.log.warn("ignoring unexpected trigger error", exc_info=True)

    def get_method(self, action):
        service, method = action.split(".")
        return service, getattr(self.services[service], method)

    def unsubscribe(self, client):
        self.client_locks.pop(client, None)
        self.cached_queries.pop(client, None)
        self.cached_fingerprints.pop(client, None)
        for clients in self.subscriptions.values():
            clients[self].pop(client, None)

    def unsubscribe_all(self):
        for clients in self.subscriptions.values():
            clients.pop(self, None)

    def update_subscriptions(self, client, callback, models):
        for clients in self.subscriptions.values():
            clients[self][client].discard(callback)

        for model in models:
            self.subscriptions[model][self][client].add(callback)

    def trigger(self, trigger, client, callback):
        if callback in self.cached_queries[client]:
            function, args, kwargs = self.cached_queries[client][callback]
            result = function(*args, **kwargs)
            self.send(trigger=trigger, client=client, callback=callback,
                      data=result)

    def update_triggers(self, client, callback, function, result,
                        delta_time=None, *args, **kwargs):
        if hasattr(function, "subscribes") and client is not None:
            self.cached_queries[client][callback] = (function, args, kwargs)
            self.update_subscriptions(client, callback, function.subscribes)
        if client is not None and callback is None and result is not self.ERROR:
            self.send(trigger="subscribe", client=client, data=result,
                      _time=delta_time)

    def internal_action(self, action, client, callback):
        if action == "unsubscribe":
            self.unsubscribe(client)
        elif action is not None:
            self.log.warn("unknown action {!r}", action)

    def received_message(self, message):
        responder.defer(self, deepcopy(message))

    def handle_message(self, message):
        self.log.debug("received {}", message.data)
        callback, client, delta_time, result = None, None, None, None
        try:
            message = json.loads(str(message.data))
            threadlocal.reset(websocket=self, message=message, username=self.username)
            action, callback, client, method = message.get("action"), message.get("callback"), message.get("client"), message.get("method")
            with self.client_locks[client] if client else RLock():
                self.internal_action(action, client, callback)
                if method:
                    service, function = self.get_method(method)
                    args, kwargs = get_params(message.get("params"))
                    result = self.ERROR
                    try:
                        before_time = time.time()
                        result = function(*args, **kwargs)
                        delta_time = (time.time() - before_time) if self.debug else None
                    finally:
                        self.update_triggers(client, callback, function, result, delta_time,
                                             *args, **kwargs)
        except:
            self.log.error("unexpected websocket dispatch error", exc_info=True)
            exc_class, exc, tb = sys.exc_info()
            str_content = str(exc) or 'Unexpected Error.'
            message = (str_content + '\n' + traceback.format_exc()) if self.debug else str_content
            self.send(error=message, callback=callback, client=client)
        else:
            if callback is not None:
                self.send(data=result, callback=callback, client=client, _time=delta_time)

    def __repr__(self):
        return '<%s username=%s>' % (self.__class__.__name__, self.username)

class WebSocketRoot(object):
    @cherrypy.expose
    def default(self):
        pass

class WebSocketChecker(WebSocketTool):
    def __init__(self):
        cherrypy.Tool.__init__(self, "before_handler", self.upgrade)

    def upgrade(self, **kwargs):
        try:
            kwargs["handler_cls"].check_authentication()
        except:
            raise cherrypy.HTTPError(401, "You must be logged in to establish a websocket connection.")
        else:
            return WebSocketTool.upgrade(self, **kwargs)

cherrypy.tools.websockets = WebSocketChecker()

websocket_plugin = WebSocketPlugin(cherrypy.engine)
websocket_plugin.subscribe()

broadcaster = Caller(WebSocketDispatcher.broadcast)
responder = Caller(WebSocketDispatcher.handle_message, threads=config['ws_thread_pool'])

cherrypy.engine.subscribe("stop", WebSocketDispatcher.close_all)
for _task in [broadcaster, responder]:
    cherrypy.engine.subscribe("start", _task.start, priority=99)
    cherrypy.engine.subscribe("stop", _task.stop)
<commit_msg>corrected websocket importing Caller<commit_after>from __future__ import unicode_literals
import sys
import json
import time
import hashlib
import logging
import traceback
from copy import deepcopy
from functools import wraps
from threading import local, RLock
from collections import defaultdict

import cherrypy

from ws4py.websocket import WebSocket
from ws4py.server.cherrypyserver import WebSocketPlugin, WebSocketTool

from sideboard.threads import Caller

import sideboard.lib
from sideboard.config import config

class threadlocal(object):
    _threadlocal = local()

    @classmethod
    def get(cls, key, default=None):
        return getattr(cls._threadlocal, key, default)

    @classmethod
    def set(cls, key, val):
        return setattr(cls._threadlocal, key, val)

    @classmethod
    def clear(cls):
        cls._threadlocal.__dict__.clear()

    @classmethod
    def get_client(cls):
        return cls.get("client") or cls.get("message", {}).get("client")

    @classmethod
    def reset(cls, **kwargs):
        cls.clear()
        for key, val in kwargs.items():
            cls.set(key, val)

def _normalize_topics(*topics):
    """
    Converts a list of types, strings, or whatever else into a list of strings.

    >>> _normalize_topics()
    []

    >>> _normalize_topics('')
    []

    >>> _normalize_topics('   ')
    []

    >>> _normalize_topics(None)
    []

    >>> _normalize_topics('topic-one', 'topic-two')
    ['topic-one', 'topic-two']

    >>> _normalize_topics('repeated-topic', 'repeated-topic')
    ['repeated-topic']

    >>> _normalize_topics('topic-left', None, 'topic-right', None)
    ['topic-left', 'topic-right']

    >>> _normalize_topics('', 'topic-left', '', 'topic-right')
    ['topic-left', 'topic-right']

    >>> _normalize_topics('   ', '   topic-padded-left', '   ', 'topic-padded-right   ', '   ')
    ['topic-padded-left', 'topic-padded-right']

    >>> _normalize_topics(type)
    ['type']

    >>> _normalize_topics(dict)
    ['dict']

    >>> _normalize_topics(dict(foo="bar"))
    ["{'foo': 'bar'}"]
    """
    normalized_topics = []
    for topic in topics:
        if topic is not None:
            if isinstance(topic, type):
                normalized_topics.append(topic.__name__)
            elif isinstance(topic, basestring):
                topic = topic.strip()
                if topic != '':
                    normalized_topics.append(topic)
            else:
                normalized_topics.append(str(topic))
    return list(set(normalized_topics))

def notify(topics, trigger="manual", delay=0, originating_client=None):
    broadcaster.delayed(delay, trigger, _normalize_topics(*sideboard.lib.listify(topics)),
                        originating_client or threadlocal.get_client())

def notifies(*args, **kwargs):
    """
    Adds a notifies attribute to the decorated function. The notifies
    attribute specifies a list of the topics which the function notifies.

    >>> @notifies('topic-one', 'topic-two')
    ... def fn():
    ...     pass
    >>> getattr(fn, 'notifies')
    ['topic-one', 'topic-two']

    >>> @notifies(dict)
    ... def fn_dict():
    ...     pass
    >>> getattr(fn_dict, 'notifies')
    ['dict']
    """
    delay = kwargs.pop("delay", 0)
    topics = _normalize_topics(*args)

    def decorated_func(func):
        @wraps(func)
        def notifier_func(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            finally:
                notify(topics, trigger=func.__name__, delay=delay)

        notifier_func.notifies = topics
        return notifier_func

    return decorated_func

def subscribes(*args):
    """
    Adds a subscribes attribute to the decorated function. The subscribes
    attribute specifies a list of the topics to which the function subscribes.

    >>> @subscribes('topic-one', 'topic-two')
    ... def fn():
    ...     pass
    >>> getattr(fn, 'subscribes')
    ['topic-one', 'topic-two']

    >>> @subscribes(dict)
    ... def fn_dict():
    ...     pass
    >>> getattr(fn_dict, 'subscribes')
    ['dict']
    """
    topics = _normalize_topics(*args)

    def decorated_func(func):
        func.subscribes = topics
        return func

    return decorated_func

def _fingerprint(x):
    """
    Calculates the md5 sum of the given argument.

    If _fingerprint is passed a string, it calculates the md5 sum of the
    string. If _fingerprint is passed anything else, then the json encoding
    of the argument is used to calculate the md5 sum.

    >>> _fingerprint(None)
    '37a6259cc0c1dae299a7866489dff0bd'

    >>> _fingerprint('test')
    '098f6bcd4621d373cade4e832627b4f6'

    >>> _fingerprint({'key':'value'})
    'a7353f7cddce808de0032747a0b7be50'

    >>> _fingerprint(dict(key='value'))
    'a7353f7cddce808de0032747a0b7be50'

    >>> _fingerprint({'a':1, 'b':2})
    '608de49a4600dbb5b173492759792e4a'

    >>> _fingerprint({'b':2, 'a':1})
    '608de49a4600dbb5b173492759792e4a'

    >>> _fingerprint({'a':{'x':3, 'y':4}, 'b':2})
    '2c22e445e9278c66dd7ea78b757defe6'

    >>> _fingerprint({'b':2, 'a':{'y':4, 'x':3}})
    '2c22e445e9278c66dd7ea78b757defe6'
    """
    md5 = hashlib.md5()
    if isinstance(x, basestring):
        md5.update(x)
    else:
        md5.update(json.dumps(x, cls=sideboard.lib.serializer, sort_keys=True, separators=(',', ':')))
    return md5.hexdigest()

def get_params(params):
    if params is None:
        return [], {}
    elif isinstance(params, dict):
        return [], params
    elif isinstance(params, (list, tuple)):
        return params, {}
    else:
        return [params], {}

class WebSocketDispatcher(WebSocket):
    ERROR = object()
    username = None
    subscriptions = defaultdict(lambda: defaultdict(lambda: defaultdict(set)))

    debug = True
    services = {}
    log = logging.getLogger(__name__)

    def __init__(self, *args, **kwargs):
        WebSocket.__init__(self, *args, **kwargs)
        self.send_lock = RLock()
        self.client_locks, self.cached_queries, self.cached_fingerprints = \
            defaultdict(RLock), defaultdict(dict), defaultdict(dict)
        self.username = self.check_authentication()

    @classmethod
    def check_authentication(cls):
        return cherrypy.session["username"]

    @classmethod
    def get_all_subscribed(cls):
        websockets = set()
        for model, subscriptions in cls.subscriptions.items():
            for websocket, clients in subscriptions.items():
                websockets.add(websocket)
        return websockets

    @classmethod
    def get_all(cls):
        return websocket_plugin.pool.keys()[:]

    @classmethod
    def close_all(cls):
        for ws in cls.get_all():
            try:
                ws.close()
            except:
                pass

    def send(self, **message):
        message = {k: v for k, v in message.items() if v is not None}
        if "data" in message and "client" in message:
            fingerprint = _fingerprint(message["data"])
            client, callback = message["client"], message.get("callback")
            repeat_send = callback in self.cached_fingerprints[client]
            cached_fingerprint = self.cached_fingerprints[client].get(callback)
            self.cached_fingerprints[client][callback] = fingerprint
            if cached_fingerprint == fingerprint and repeat_send:
                return

        self.log.debug("sending {}", message)
        message = json.dumps(message, cls=sideboard.lib.serializer, separators=(",", ":"), sort_keys=True)
        with self.send_lock:
            WebSocket.send(self, message)

    def closed(self, code, reason=""):
        self.log.info("closing: code={!r} reason={!r}", code, reason)
        self.unsubscribe_all()
        WebSocket.closed(self, code, reason)
        notify("Site")

    @classmethod
    def broadcast(cls, trigger, models, originating_client=None):
        triggered = set()
        for model in models:
            for websocket, clients in cls.subscriptions[model].items():
                for client, callbacks in clients.copy().items():
                    if client != originating_client:
                        for callback in callbacks:
                            triggered.add((websocket, client, callback))

        for websocket, client, callback in triggered:
            try:
                websocket.trigger(trigger, client, callback)
            except:
                cls.log.warn("ignoring unexpected trigger error", exc_info=True)

    def get_method(self, action):
        service, method = action.split(".")
        return service, getattr(self.services[service], method)

    def unsubscribe(self, client):
        self.client_locks.pop(client, None)
        self.cached_queries.pop(client, None)
        self.cached_fingerprints.pop(client, None)
        for clients in self.subscriptions.values():
            clients[self].pop(client, None)

    def unsubscribe_all(self):
        for clients in self.subscriptions.values():
            clients.pop(self, None)

    def update_subscriptions(self, client, callback, models):
        for clients in self.subscriptions.values():
            clients[self][client].discard(callback)

        for model in models:
            self.subscriptions[model][self][client].add(callback)

    def trigger(self, trigger, client, callback):
        if callback in self.cached_queries[client]:
            function, args, kwargs = self.cached_queries[client][callback]
            result = function(*args, **kwargs)
            self.send(trigger=trigger, client=client, callback=callback,
                      data=result)

    def update_triggers(self, client, callback, function, result,
                        delta_time=None, *args, **kwargs):
        if hasattr(function, "subscribes") and client is not None:
            self.cached_queries[client][callback] = (function, args, kwargs)
            self.update_subscriptions(client, callback, function.subscribes)
        if client is not None and callback is None and result is not self.ERROR:
            self.send(trigger="subscribe", client=client, data=result,
                      _time=delta_time)

    def internal_action(self, action, client, callback):
        if action == "unsubscribe":
            self.unsubscribe(client)
        elif action is not None:
            self.log.warn("unknown action {!r}", action)

    def received_message(self, message):
        responder.defer(self, deepcopy(message))

    def handle_message(self, message):
        self.log.debug("received {}", message.data)
        callback, client, delta_time, result = None, None, None, None
        try:
            message = json.loads(str(message.data))
            threadlocal.reset(websocket=self, message=message, username=self.username)
            action, callback, client, method = message.get("action"), message.get("callback"), message.get("client"), message.get("method")
            with self.client_locks[client] if client else RLock():
                self.internal_action(action, client, callback)
                if method:
                    service, function = self.get_method(method)
                    args, kwargs = get_params(message.get("params"))
                    result = self.ERROR
                    try:
                        before_time = time.time()
                        result = function(*args, **kwargs)
                        delta_time = (time.time() - before_time) if self.debug else None
                    finally:
                        self.update_triggers(client, callback, function, result, delta_time,
                                             *args, **kwargs)
        except:
            self.log.error("unexpected websocket dispatch error", exc_info=True)
            exc_class, exc, tb = sys.exc_info()
            str_content = str(exc) or 'Unexpected Error.'
            message = (str_content + '\n' + traceback.format_exc()) if self.debug else str_content
            self.send(error=message, callback=callback, client=client)
        else:
            if callback is not None:
                self.send(data=result, callback=callback, client=client, _time=delta_time)

    def __repr__(self):
        return '<%s username=%s>' % (self.__class__.__name__, self.username)

class WebSocketRoot(object):
    @cherrypy.expose
    def default(self):
        pass

class WebSocketChecker(WebSocketTool):
    def __init__(self):
        cherrypy.Tool.__init__(self, "before_handler", self.upgrade)

    def upgrade(self, **kwargs):
        try:
            kwargs["handler_cls"].check_authentication()
        except:
            raise cherrypy.HTTPError(401, "You must be logged in to establish a websocket connection.")
        else:
            return WebSocketTool.upgrade(self, **kwargs)

cherrypy.tools.websockets = WebSocketChecker()

websocket_plugin = WebSocketPlugin(cherrypy.engine)
websocket_plugin.subscribe()

broadcaster = Caller(WebSocketDispatcher.broadcast)
responder = Caller(WebSocketDispatcher.handle_message, threads=config['ws_thread_pool'])

cherrypy.engine.subscribe("stop", WebSocketDispatcher.close_all)
for _task in [broadcaster, responder]:
    cherrypy.engine.subscribe("start", _task.start, priority=99)
    cherrypy.engine.subscribe("stop", _task.stop)
<|endoftext|>


<commit_before>"""
Akshay Tambe
Pachter and Doudna groups
"""

from sircel.utils import IO_utils, Plot_utils
from sircel import Split_reads
import argparse
import os
import sys
import json
import time
import shutil
import subprocess
from pathlib import Path
import gc

from scipy.sparse import coo_matrix
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import pairwise_distances
from scipy.spatial.distance import *

def run_all(args):
	print('\nInspecting and pre-processing inputs')

	if(args['output_dir'][-1] == '/'):
		args['output_dir'] = args['output_dir'][0:-1]
	if not os.path.exists(args['output_dir']):
		os.makedirs(args['output_dir'])
	if not os.path.exists(args['output_dir'] + '/plots'):
		os.makedirs(args['output_dir'] + '/plots')
	with (Path(__file__).parent / 'params.json').open() as r:
		kallisto = json.load(r)['kallisto']
	assert kallisto

	split_args = {}
	check_pipeline_input(args, kallisto)

	if args['dropseq']:
		args['barcode_start']	= 0
		args['barcode_end'] 		= 12
		args['umi_start']			= 12
		args['umi_end'] 			= 20
		if args['kmer_size'] == None:
			args['kmer_size'] = 8
		print('Unzipping files (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = \
			IO_utils.unzip(args['barcodes'].split(','))
		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped

	elif args['10xgenomics']:
		args['barcode_start']	= 0
		args['barcode_end'] 		= 26
		args['umi_start'] 		= 26
		args['umi_end'] 			= 34
		if args['kmer_size'] == None:
			args['kmer_size'] = 20
		print('Unzipping and merging files (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = IO_utils.merge_barcodefiles_10x(
			args['barcodes'].split(','),
			args['umis'].split(','))

		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped

	else:
		if args['kmer_size'] == None:
			args['kmer_size'] = 8
		print('Unzipping (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = \
			IO_utils.unzip(args['barcodes'].split(','))
		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped

	check_split_input(args)
	output_files, elapsed_time = Split_reads.run_all(args)
	output_files['args'] = args
	print('Done idetifying barcodes and splitting reads.\n' + \
		'\tTime elapsed: %0.002f seconds\n' % elapsed_time)

	#print(args['kallisto_idx'])
	if(args['kallisto_idx'] != None):
		print('Running kallisto')
		kallisto_dir = '%s/kallisto_outputs' % args['output_dir']
		if not os.path.exists(kallisto_dir):
			os.makedirs(kallisto_dir)
		output_files['kallisto'] = run_kallisto(
			args,
			kallisto,
			kallisto_dir,
			output_files)
		print('Getting transcript compatibility counts')
		output_files['tcc'] = write_transcript_compatability_counts(
			args,
			output_files,
			kallisto_dir)

	print('Removing temp files')
	os.unlink(reads_unzipped)
	os.unlink(barcodes_unzipped)
	output_files['run_outputs'] = \
		'%s/run_outputs.json' % args['output_dir']
	with open(output_files['run_outputs'], 'w') as writer:
		writer.write(json.dumps(output_files, indent=3))

	print('Done.')
	return output_files

def run_kallisto(args, kallisto, kallisto_dir, output_files):
	kallisto_start_time = time.time()

	kallisto_cmd = [
		kallisto, 'pseudo',
		'-b',			output_files['split']['batch'],
		'-i', 		args['kallisto_idx'],
		'-o',			kallisto_dir,
		'-t',			str(args['threads']),
		'--umi']

	kallisto = subprocess.check_call(kallisto_cmd)
	kallisto_output = {
		'tsv' :				'%s/matrix.tsv' % kallisto_dir,
		'run_info'	:		'%s/run_info.json' % kallisto_dir,
		'equiv_classes' :	'%s/matrix.ec' % kallisto_dir,
		'cells' : 			'%s/matrix.cells' % kallisto_dir}
	current_time = time.time()
	elapsed_time = current_time - kallisto_start_time
	print('Time elapsed %0.2f' % elapsed_time)

	return kallisto_output

def write_transcript_compatability_counts(args, input_files, kallisto_dir):
	"""
	Modifed from Vasilis Ntranos scRNA-seq-TCC-prep
		https://github.com/pachterlab/scRNA-Seq-TCC-prep
	"""

	import numpy as np
	import pickle

	for fname in input_files['kallisto'].values():
		assert(os.path.exists(fname)), \
			print('kallisto output file not found: %s' % fname)

	output_files = {
		'tcc_coo':		'%s/tcc_matrix.dat' % kallisto_dir,
		'tcc_csr':		'%s/tcc_matrix_csr.dat' % kallisto_dir,
		'l1_dist' : 	'%s/pairwise_l1_distance.npy' % kallisto_dir,
		'nonzero_eq' :	'%s/nonzero_equiv_classes.npy' % kallisto_dir,
		'tcc_norm_t' :	'%s/tcc_normalized_transposed.npy' % kallisto_dir}

	print('\tLoading kallisto matrices')
	#matrix.ec file
	equiv_class_fname =  input_files['kallisto']['equiv_classes']
	tsv_fname = input_files['kallisto']['tsv']

	tsv = np.genfromtxt(tsv_fname, delimiter='\t' , dtype=int)
	tsv_rows, tsv_cols, tsv_data = tsv.T
	nonzero_equiv_classes = np.unique(tsv_rows)

	map_rows = {
		val:ind for ind,val in enumerate(nonzero_equiv_classes)}
	map_cols = {
		val:ind for ind,val in enumerate(np.unique(tsv_cols))}

	print('\tPreparing TCC matrix')
	#use scipy coo_matrix for sparse data
	tcc_matrix = coo_matrix(
		(tsv_data.astype(float),
		( [map_rows[r] for r in tsv_rows], [map_cols[c] for c in tsv_cols])))
	tcc_matrix_csr = tcc_matrix.tocsr()
	#save coo_matrix and csr_matrix using pickle (?)
	with open(output_files['tcc_coo'], 'wb') as outf:
		pickle.dump(tcc_matrix, outf, pickle.HIGHEST_PROTOCOL)
	with open(output_files['tcc_csr'], 'wb') as outf:
		pickle.dump(tcc_matrix_csr, outf, pickle.HIGHEST_PROTOCOL)

	print('\tNormalizing TCCs')
	tcc_normalized = normalize(tcc_matrix_csr, norm='l1', axis=0)
	tcc_normalized_transposed = tcc_normalized.transpose()

	print('\tComputing L1 distance')
	l1_distance = pairwise_distances(
		tcc_normalized_transposed,
		metric=get_l1_distance,
		n_jobs=args['threads'])

	print('\tSaving outputs as .npy matrices')
	np.save(output_files['l1_dist'], l1_distance)
	np.save(output_files['nonzero_eq'], nonzero_equiv_classes)
	np.save(output_files['tcc_norm_t'], tcc_normalized_transposed)
	return output_files

def get_l1_distance(p,q):
    return cityblock(p,q).sum()

def check_pipeline_input(args, kallisto):

	assert not (args['10xgenomics'] and args['dropseq']), \
		'10xgenomics and dropseq options are mutually exclusive'
	assert Path(kallisto).is_file() or shutil.which(kallisto), \
		'Cannot find kallisto executable %s' % kallisto
	assert os.path.exists(args['reads']), \
		'Cannot find reads file %s' % args['reads']
	assert os.path.exists(args['barcodes']), \
		'Cannot find barcodes file %s' % args['barcodes']

	if args['kallisto_idx'] is not None:
		assert(os.path.exists(args['kallisto_idx'])), \
			'Cannot find kallisto index %s' % \
			args['kallisto_idx']
	if args['10xgenomics']:
		assert(os.path.exists(args['umis'])), \
			'Cannot find reads file %s' % args['umis']

def check_split_input(args):
	"""
	More here
	"""
	assert args['barcode_start'] < args['barcode_end'], \
		'Barcode end position %i is <= start position %i' % \
		(args['barcode_end'], args['barcode_start'])
	assert args['umi_start'] <= args['umi_end'], \
		'UMI end position %i is less than start position %i' % \
		(args['umi_end'], args['umi_start'])
	assert os.path.exists(args['reads']), \
		'Cannot find reads file %s' % args['reads']
	assert os.path.exists(args['barcodes']), \
		'Cannot find barcodes file %s' % args['barcodes']
	assert args['kmer_size'] > 0, \
		'Kmer size must be positve. %i' % args['kmer_size']

def get_args(args=None):
	if args is None:
		args = sys.argv[1:]

	parser = argparse.ArgumentParser(
		description = 'This script splits reads for dropseq data',
		formatter_class = argparse.ArgumentDefaultsHelpFormatter)

	parser.add_argument('--dropseq',
		help='Use barcode / umi positions from Macosko et al Dropseq dataset',
		required=False,
		action='store_true')
	#parser.add_argument('--10xgenomics',
	#	help='Use barcode / umi positions from 10xGenomics v2 chemistry data',
	#	required=False,
	#	action='store_true')

	parser.add_argument('--barcodes',
		type=str,
		help='Cell barcodes file name or comma separarted list.',
		required=True)
	#parser.add_argument('--umis',
	#	type=str,
	#	help='Umis (R1) file name or comma separarted list. Only required for 10x genomics',
	#	required=False)
	parser.add_argument('--reads',
		type=str,
		help='RNAseq reads file name or comma separarted list.',
		required=True)
	parser.add_argument('--output_dir',
		type=str,
		help='Directory where outputs are written',
		required=True)

	parser.add_argument('--kmer_size',
		type=int,
		help='Size of kmers for making barcode De Bruijn graph.',
		default=None)
	parser.add_argument('--breadth',
		type=int,
		help='How many nodes search.',
		default=10000)
	parser.add_argument('--depth',
		type=int,
		help='How many paths to search for at each node.',
		default=10)
	parser.add_argument('--threads',
		type=int,
		help='Number of threads to use.',
		default=32)
	parser.add_argument('--kallisto_idx',
		type=str,
		help='Path to kallisto index. Optional.' + \
			'If not provided, script terminates after splitting barcodes',
		default=None)
	parser.add_argument('--num_cells',
		type=int,
		help='Estimated number of cells.',
		default=None)

	parser.add_argument(
		'--barcode_start',
		type=int,
		default=None)
	parser.add_argument(
		'--barcode_end',
		type=int,
		default=None)
	parser.add_argument(
		'--umi_start',
		type=int,
		default=0)
	parser.add_argument(
		'--umi_end',
		type=int,
		default=0)

	#set this to false only for reviewer expts
	parser.add_argument('--split_levenshtein',
		type = bool,
		help = argparse.SUPPRESS,
		default = True)

	return vars(parser.parse_args(args))

<commit_msg>Naive_pipeline.py<commit_after>"""
Akshay Tambe
Pachter and Doudna groups
"""

from sircel.utils import IO_utils, Plot_utils
from sircel import Split_reads
import argparse
import os
import sys
import json
import time
import shutil
import subprocess
from pathlib import Path
import gc

from scipy.sparse import coo_matrix
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import pairwise_distances
from scipy.spatial.distance import *

def run_all(args):
	print('\nInspecting and pre-processing inputs')

	if(args['output_dir'][-1] == '/'):
		args['output_dir'] = args['output_dir'][0:-1]
	if not os.path.exists(args['output_dir']):
		os.makedirs(args['output_dir'])
	if not os.path.exists(args['output_dir'] + '/plots'):
		os.makedirs(args['output_dir'] + '/plots')
	with (Path(__file__).parent / 'params.json').open() as r:
		kallisto = json.load(r)['kallisto']
	assert kallisto

	split_args = {}
	check_pipeline_input(args, kallisto)

	if args['dropseq']:
		args['barcode_start']	= 0
		args['barcode_end'] 		= 12
		args['umi_start']			= 12
		args['umi_end'] 			= 20
		if args['kmer_size'] == None:
			args['kmer_size'] = 8
		print('Unzipping files (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = \
			IO_utils.unzip(args['barcodes'].split(','))
		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped
	"""
	elif args['10xgenomics']:
		args['barcode_start']	= 0
		args['barcode_end'] 		= 26
		args['umi_start'] 		= 26
		args['umi_end'] 			= 34
		if args['kmer_size'] == None:
			args['kmer_size'] = 20
		print('Unzipping and merging files (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = IO_utils.merge_barcodefiles_10x(
			args['barcodes'].split(','),
			args['umis'].split(','))

		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped
	"""
	else:
		if args['kmer_size'] == None:
			args['kmer_size'] = 8
		print('Unzipping (temporary)')
		reads_unzipped = \
			IO_utils.unzip(args['reads'].split(','))
		barcodes_unzipped = \
			IO_utils.unzip(args['barcodes'].split(','))
		args['reads'] = reads_unzipped
		args['barcodes'] = barcodes_unzipped

	check_split_input(args)
	output_files, elapsed_time = Split_reads.run_all(args)
	output_files['args'] = args
	print('Done idetifying barcodes and splitting reads.\n' + \
		'\tTime elapsed: %0.002f seconds\n' % elapsed_time)

	#print(args['kallisto_idx'])
	if(args['kallisto_idx'] != None):
		print('Running kallisto')
		kallisto_dir = '%s/kallisto_outputs' % args['output_dir']
		if not os.path.exists(kallisto_dir):
			os.makedirs(kallisto_dir)
		output_files['kallisto'] = run_kallisto(
			args,
			kallisto,
			kallisto_dir,
			output_files)
		print('Getting transcript compatibility counts')
		output_files['tcc'] = write_transcript_compatability_counts(
			args,
			output_files,
			kallisto_dir)

	print('Removing temp files')
	os.unlink(reads_unzipped)
	os.unlink(barcodes_unzipped)
	output_files['run_outputs'] = \
		'%s/run_outputs.json' % args['output_dir']
	with open(output_files['run_outputs'], 'w') as writer:
		writer.write(json.dumps(output_files, indent=3))

	print('Done.')
	return output_files

def run_kallisto(args, kallisto, kallisto_dir, output_files):
	kallisto_start_time = time.time()

	kallisto_cmd = [
		kallisto, 'pseudo',
		'-b',			output_files['split']['batch'],
		'-i', 		args['kallisto_idx'],
		'-o',			kallisto_dir,
		'-t',			str(args['threads']),
		'--umi']

	kallisto = subprocess.check_call(kallisto_cmd)
	kallisto_output = {
		'tsv' :				'%s/matrix.tsv' % kallisto_dir,
		'run_info'	:		'%s/run_info.json' % kallisto_dir,
		'equiv_classes' :	'%s/matrix.ec' % kallisto_dir,
		'cells' : 			'%s/matrix.cells' % kallisto_dir}
	current_time = time.time()
	elapsed_time = current_time - kallisto_start_time
	print('Time elapsed %0.2f' % elapsed_time)

	return kallisto_output

def write_transcript_compatability_counts(args, input_files, kallisto_dir):
	"""
	Modifed from Vasilis Ntranos scRNA-seq-TCC-prep
		https://github.com/pachterlab/scRNA-Seq-TCC-prep
	"""

	import numpy as np
	import pickle

	for fname in input_files['kallisto'].values():
		assert(os.path.exists(fname)), \
			print('kallisto output file not found: %s' % fname)

	output_files = {
		'tcc_coo':		'%s/tcc_matrix.dat' % kallisto_dir,
		'tcc_csr':		'%s/tcc_matrix_csr.dat' % kallisto_dir,
		'l1_dist' : 	'%s/pairwise_l1_distance.npy' % kallisto_dir,
		'nonzero_eq' :	'%s/nonzero_equiv_classes.npy' % kallisto_dir,
		'tcc_norm_t' :	'%s/tcc_normalized_transposed.npy' % kallisto_dir}

	print('\tLoading kallisto matrices')
	#matrix.ec file
	equiv_class_fname =  input_files['kallisto']['equiv_classes']
	tsv_fname = input_files['kallisto']['tsv']

	tsv = np.genfromtxt(tsv_fname, delimiter='\t' , dtype=int)
	tsv_rows, tsv_cols, tsv_data = tsv.T
	nonzero_equiv_classes = np.unique(tsv_rows)

	map_rows = {
		val:ind for ind,val in enumerate(nonzero_equiv_classes)}
	map_cols = {
		val:ind for ind,val in enumerate(np.unique(tsv_cols))}

	print('\tPreparing TCC matrix')
	#use scipy coo_matrix for sparse data
	tcc_matrix = coo_matrix(
		(tsv_data.astype(float),
		( [map_rows[r] for r in tsv_rows], [map_cols[c] for c in tsv_cols])))
	tcc_matrix_csr = tcc_matrix.tocsr()
	#save coo_matrix and csr_matrix using pickle (?)
	with open(output_files['tcc_coo'], 'wb') as outf:
		pickle.dump(tcc_matrix, outf, pickle.HIGHEST_PROTOCOL)
	with open(output_files['tcc_csr'], 'wb') as outf:
		pickle.dump(tcc_matrix_csr, outf, pickle.HIGHEST_PROTOCOL)

	print('\tNormalizing TCCs')
	tcc_normalized = normalize(tcc_matrix_csr, norm='l1', axis=0)
	tcc_normalized_transposed = tcc_normalized.transpose()

	print('\tComputing L1 distance')
	l1_distance = pairwise_distances(
		tcc_normalized_transposed,
		metric=get_l1_distance,
		n_jobs=args['threads'])

	print('\tSaving outputs as .npy matrices')
	np.save(output_files['l1_dist'], l1_distance)
	np.save(output_files['nonzero_eq'], nonzero_equiv_classes)
	np.save(output_files['tcc_norm_t'], tcc_normalized_transposed)
	return output_files

def get_l1_distance(p,q):
    return cityblock(p,q).sum()

def check_pipeline_input(args, kallisto):

	#assert not (args['10xgenomics'] and args['dropseq']), \
	#	'10xgenomics and dropseq options are mutually exclusive'
	assert Path(kallisto).is_file() or shutil.which(kallisto), \
		'Cannot find kallisto executable %s' % kallisto
	assert os.path.exists(args['reads']), \
		'Cannot find reads file %s' % args['reads']
	assert os.path.exists(args['barcodes']), \
		'Cannot find barcodes file %s' % args['barcodes']

	if args['kallisto_idx'] is not None:
		assert(os.path.exists(args['kallisto_idx'])), \
			'Cannot find kallisto index %s' % \
			args['kallisto_idx']
	#if args['10xgenomics']:
	#	assert(os.path.exists(args['umis'])), \
	#		'Cannot find reads file %s' % args['umis']

def check_split_input(args):
	"""
	More here
	"""
	assert args['barcode_start'] < args['barcode_end'], \
		'Barcode end position %i is <= start position %i' % \
		(args['barcode_end'], args['barcode_start'])
	assert args['umi_start'] <= args['umi_end'], \
		'UMI end position %i is less than start position %i' % \
		(args['umi_end'], args['umi_start'])
	assert os.path.exists(args['reads']), \
		'Cannot find reads file %s' % args['reads']
	assert os.path.exists(args['barcodes']), \
		'Cannot find barcodes file %s' % args['barcodes']
	assert args['kmer_size'] > 0, \
		'Kmer size must be positve. %i' % args['kmer_size']

def get_args(args=None):
	if args is None:
		args = sys.argv[1:]

	parser = argparse.ArgumentParser(
		description = 'This script splits reads for dropseq data',
		formatter_class = argparse.ArgumentDefaultsHelpFormatter)

	parser.add_argument('--dropseq',
		help='Use barcode / umi positions from Macosko et al Dropseq dataset',
		required=False,
		action='store_true')
	#parser.add_argument('--10xgenomics',
	#	help='Use barcode / umi positions from 10xGenomics v2 chemistry data',
	#	required=False,
	#	action='store_true')

	parser.add_argument('--barcodes',
		type=str,
		help='Cell barcodes file name or comma separarted list.',
		required=True)
	#parser.add_argument('--umis',
	#	type=str,
	#	help='Umis (R1) file name or comma separarted list. Only required for 10x genomics',
	#	required=False)
	parser.add_argument('--reads',
		type=str,
		help='RNAseq reads file name or comma separarted list.',
		required=True)
	parser.add_argument('--output_dir',
		type=str,
		help='Directory where outputs are written',
		required=True)

	parser.add_argument('--kmer_size',
		type=int,
		help='Size of kmers for making barcode De Bruijn graph.',
		default=None)
	parser.add_argument('--breadth',
		type=int,
		help='How many nodes search.',
		default=10000)
	parser.add_argument('--depth',
		type=int,
		help='How many paths to search for at each node.',
		default=10)
	parser.add_argument('--threads',
		type=int,
		help='Number of threads to use.',
		default=32)
	parser.add_argument('--kallisto_idx',
		type=str,
		help='Path to kallisto index. Optional.' + \
			'If not provided, script terminates after splitting barcodes',
		default=None)
	parser.add_argument('--num_cells',
		type=int,
		help='Estimated number of cells.',
		default=None)

	parser.add_argument(
		'--barcode_start',
		type=int,
		default=None)
	parser.add_argument(
		'--barcode_end',
		type=int,
		default=None)
	parser.add_argument(
		'--umi_start',
		type=int,
		default=0)
	parser.add_argument(
		'--umi_end',
		type=int,
		default=0)

	#set this to false only for reviewer expts
	parser.add_argument('--split_levenshtein',
		type = bool,
		help = argparse.SUPPRESS,
		default = True)

	return vars(parser.parse_args(args))

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
u"""SRW execution template.

:copyright: Copyright (c) 2015 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function

from pykern import pkio
from pykern import pkjinja
from . import template_common
import h5py
import numpy as np
import os
import re

#: How long before killing WARP process
MAX_SECONDS = 60 * 60

_MODE_TEXT = {
    '0': '0',
    '1': '1 (real part)',
    '2': '1 (imaginary part)',
}

_PARTICLE_ARG_PATH = {
    'x' : 'position/x',
    'y' : 'position/y',
    'z' : 'position/z',
    'ux' : 'momentum/x',
    'uy' : 'momentum/y',
    'uz' : 'momentum/z',
}

def background_percent_complete(data, persistent_files_dir, is_running):
    simulation_id = data['models']['simulation']['simulationId']
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    if len(files) < 2:
        return {
            'percent_complete': 0,
            'frame_count': 0,
            'total_frames': 0,
        }
    file_index = len(files) - 1
    # look at 2nd to last file if running, last one may be incomplete
    if is_running:
        file_index -= 1
    field = 'E'
    coordinate = 'r'
    Lplasma_lab = float(data['models']['electronPlasma']['length']) / 1e3
    zmmin = float(data['models']['simulationGrid']['zMin']) / 1e6
    dfile = h5py.File(files[file_index], 'r')
    dset = dfile['fields/{}/{}'.format(field, coordinate)]
    dz = dset.attrs['dx']
    zmin = dset.attrs['xmin']
    percent_complete = (zmin - zmmin) / (Lplasma_lab - zmmin)
    if percent_complete < 0:
        percent_complete = 0.0
    elif percent_complete > 1.0:
        percent_complete = 1.0
    return {
        'percent_complete': percent_complete * 100,
        'frame_count': file_index + 1,
        #TODO(pjm): this isn't a great calculation ...
        'total_frames': int((file_index + 1) / percent_complete + 0.3),
    }

def extract_field_report(field, coordinate, mode, dfile, iteration):
    if field == 'rho' :
        dset = dfile['fields/rho']
        coordinate = ''
    else:
        dset = dfile['fields/{}/{}'.format(field, coordinate)]
    F = np.flipud(np.array(dset[mode,:,:]).T)
    Nr, Nz = F.shape[0], F.shape[1]
    dz = dset.attrs['dx']
    dr = dset.attrs['dy']
    zmin = dset.attrs['xmin']
    extent = np.array([zmin-0.5*dz, zmin+0.5*dz+dz*Nz, 0., (Nr+1)*dr])
    return {
        'x_range': [extent[0], extent[1], len(F[0])],
        'y_range': [extent[2], extent[3], len(F)],
        'x_label': 'z [m]',
        'y_label': 'r [m]',
        'title': "{} {} in the mode {} at {}".format(
            field, coordinate, _MODE_TEXT[str(mode)], _iteration_title(dfile, iteration)),
        'z_matrix': np.flipud(F).tolist(),
    }

def fixup_old_data(data):
    if 'laserPreviewReport' not in data['models']:
        data['models']['laserPreviewReport'] = {}
    if 'particleAnimation' not in data['models']:
        data['models']['particleAnimation'] = {
            'x': 'z',
            'y': 'x',
            'histogramBins': 100,
        }
    if 'simulationStatus' not in data['models']:
        data['models']['simulationStatus'] = {
            'startTime': 0,
            'state': 'initial',
        }
    if 'startTime' not in data['models']['simulationStatus']:
        data['models']['simulationStatus']['startTime'] = 0
    if 'histogramBins' not in data['models']['particleAnimation']:
        data['models']['particleAnimation']['histogramBins'] = 100
    if 'framesPerSecond' not in data['models']['fieldAnimation']:
        data['models']['fieldAnimation']['framesPerSecond'] = 20
        data['models']['particleAnimation']['framesPerSecond'] = 20
    if 'rScale' not in data['models']['simulationGrid']:
        grid = data['models']['simulationGrid']
        grid['rScale'] = 4
        grid['rLength'] = '20.324980154380'
        grid['rMin'] = 0
        grid['rMax'] = '20.324980154380'
        grid['rCellsPerSpotSize'] = 8
        grid['rCount'] = 32
        grid['zScale'] = 2
        grid['zLength'] = '20.324980154631'
        grid['zMin'] = '-20.324980154631'
        grid['zMax'] = '1.60'
        grid['zCellsPerWavelength'] = 8
        grid['zCount'] = 214
        del grid['xMin']
        del grid['xMax']
        del grid['xCount']
        del grid['zLambda']
    if 'rParticlesPerCell' not in data['models']['simulationGrid']:
        data['models']['simulationGrid']['rParticlesPerCell'] = 1
        data['models']['simulationGrid']['zParticlesPerCell'] = 2
    if 'field' not in data['models']['laserPreviewReport']:
        laserPreview = data['models']['laserPreviewReport']
        laserPreview['field'] = 'E'
        laserPreview['coordinate'] = 'r'
        laserPreview['mode'] = '1'

def generate_parameters_file(data, schema, persistent_files_dir=None):
    _validate_data(data, schema)
    v = template_common.flatten_data(data['models'], {})
    if persistent_files_dir:
        if not persistent_files_dir.check():
            pkio.mkdir_parent(persistent_files_dir)
        v['outputDir'] = '"{}"'.format(persistent_files_dir)
    else:
        v['outputDir'] = None
    v['enablePlasma'] = 1
    return pkjinja.render_resource('warp.py', v)

def get_simulation_frame(persistent_files_dir, data):
    frame_index = int(data['frame_index'])
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    filename = files[frame_index]
    iteration = int(re.search(r'data(\d+)', filename).group(1))
    dfile = h5py.File(filename, "r")
    args = data['animation_args'].split('_')

    if data['model_name'] == 'fieldAnimation':
        return _field_animation(args, dfile, iteration, len(files))
    if data['model_name'] == 'particleAnimation':
        return _particle_animation(args, dfile, iteration, len(files))
    raise RuntimeError('{}: unknown simulation frame model'.format(data['model_name']))

def prepare_aux_files(wd):
    pass

def remove_last_frame(persistent_files_dir):
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    if len(files) > 0:
        os.remove(files[-1])

def run_all_text():
    return '''
doit = True
while(doit):
    step(10)
    doit = ( w3d.zmmin + top.zgrid < Lplasma )
'''

def _field_animation(args, dfile, iteration, frame_count):
    field = args[0]
    coordinate = args[1]
    mode = int(args[2])
    res = extract_field_report(field, coordinate, mode, dfile, iteration)
    res['frameCount'] = frame_count
    return res

def _h5_file_list(path_to_dir):
    if not os.path.isdir(path_to_dir):
        return []
    all_files = os.listdir(path_to_dir)

    # Select the hdf5 files
    h5_files = []
    for filename in all_files :
        if filename[-3:] == '.h5' :
            h5_files.append( os.path.join( path_to_dir, filename) )
    # Sort them
    h5_files.sort()
    return h5_files

def _iteration_title(dfile, iteration):
    return '{:.1f} fs (iteration {})'.format(
        iteration * 1e15 * float(dfile.attrs['timeStepUnitSI']), iteration)

def _particle_animation(args, dfile, iteration, frame_count):
    xarg = args[0]
    yarg = args[1]
    histogramBins = args[2]
    x = dfile['particles/electrons/{}'.format(_PARTICLE_ARG_PATH[xarg])][:]
    y = dfile['particles/electrons/{}'.format(_PARTICLE_ARG_PATH[yarg])][:]
    hist, edges = np.histogramdd([x, y], int(histogramBins))
    xunits = ' [m]' if len(xarg) == 1 else ''
    yunits = ' [m]' if len(yarg) == 1 else ''
    return {
        'x_range': [float(edges[0][0]), float(edges[0][-1]), len(hist)],
        'y_range': [float(edges[1][0]), float(edges[1][-1]), len(hist[0])],
        'x_label': '{}{}'.format(xarg, xunits),
        'y_label': '{}{}'.format(yarg, yunits),
        'title': 't = {}'.format(_iteration_title(dfile, iteration)),
        'z_matrix': hist.T.tolist(),
        'frameCount': frame_count,
    }

def _validate_data(data, schema):
    # ensure enums match, convert ints/floats, apply scaling
    enum_info = template_common.parse_enums(schema['enum'])
    for k in data['models']:
        if k in schema['model']:
            template_common.validate_model(data['models'][k], schema['model'][k], enum_info)
<commit_msg>stub for new_simulation()<commit_after># -*- coding: utf-8 -*-
u"""SRW execution template.

:copyright: Copyright (c) 2015 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function

from pykern import pkio
from pykern import pkjinja
from . import template_common
import h5py
import numpy as np
import os
import re

#: How long before killing WARP process
MAX_SECONDS = 60 * 60

_MODE_TEXT = {
    '0': '0',
    '1': '1 (real part)',
    '2': '1 (imaginary part)',
}

_PARTICLE_ARG_PATH = {
    'x' : 'position/x',
    'y' : 'position/y',
    'z' : 'position/z',
    'ux' : 'momentum/x',
    'uy' : 'momentum/y',
    'uz' : 'momentum/z',
}

def background_percent_complete(data, persistent_files_dir, is_running):
    simulation_id = data['models']['simulation']['simulationId']
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    if len(files) < 2:
        return {
            'percent_complete': 0,
            'frame_count': 0,
            'total_frames': 0,
        }
    file_index = len(files) - 1
    # look at 2nd to last file if running, last one may be incomplete
    if is_running:
        file_index -= 1
    field = 'E'
    coordinate = 'r'
    Lplasma_lab = float(data['models']['electronPlasma']['length']) / 1e3
    zmmin = float(data['models']['simulationGrid']['zMin']) / 1e6
    dfile = h5py.File(files[file_index], 'r')
    dset = dfile['fields/{}/{}'.format(field, coordinate)]
    dz = dset.attrs['dx']
    zmin = dset.attrs['xmin']
    percent_complete = (zmin - zmmin) / (Lplasma_lab - zmmin)
    if percent_complete < 0:
        percent_complete = 0.0
    elif percent_complete > 1.0:
        percent_complete = 1.0
    return {
        'percent_complete': percent_complete * 100,
        'frame_count': file_index + 1,
        #TODO(pjm): this isn't a great calculation ...
        'total_frames': int((file_index + 1) / percent_complete + 0.3),
    }

def extract_field_report(field, coordinate, mode, dfile, iteration):
    if field == 'rho' :
        dset = dfile['fields/rho']
        coordinate = ''
    else:
        dset = dfile['fields/{}/{}'.format(field, coordinate)]
    F = np.flipud(np.array(dset[mode,:,:]).T)
    Nr, Nz = F.shape[0], F.shape[1]
    dz = dset.attrs['dx']
    dr = dset.attrs['dy']
    zmin = dset.attrs['xmin']
    extent = np.array([zmin-0.5*dz, zmin+0.5*dz+dz*Nz, 0., (Nr+1)*dr])
    return {
        'x_range': [extent[0], extent[1], len(F[0])],
        'y_range': [extent[2], extent[3], len(F)],
        'x_label': 'z [m]',
        'y_label': 'r [m]',
        'title': "{} {} in the mode {} at {}".format(
            field, coordinate, _MODE_TEXT[str(mode)], _iteration_title(dfile, iteration)),
        'z_matrix': np.flipud(F).tolist(),
    }

def fixup_old_data(data):
    if 'laserPreviewReport' not in data['models']:
        data['models']['laserPreviewReport'] = {}
    if 'particleAnimation' not in data['models']:
        data['models']['particleAnimation'] = {
            'x': 'z',
            'y': 'x',
            'histogramBins': 100,
        }
    if 'simulationStatus' not in data['models']:
        data['models']['simulationStatus'] = {
            'startTime': 0,
            'state': 'initial',
        }
    if 'startTime' not in data['models']['simulationStatus']:
        data['models']['simulationStatus']['startTime'] = 0
    if 'histogramBins' not in data['models']['particleAnimation']:
        data['models']['particleAnimation']['histogramBins'] = 100
    if 'framesPerSecond' not in data['models']['fieldAnimation']:
        data['models']['fieldAnimation']['framesPerSecond'] = 20
        data['models']['particleAnimation']['framesPerSecond'] = 20
    if 'rScale' not in data['models']['simulationGrid']:
        grid = data['models']['simulationGrid']
        grid['rScale'] = 4
        grid['rLength'] = '20.324980154380'
        grid['rMin'] = 0
        grid['rMax'] = '20.324980154380'
        grid['rCellsPerSpotSize'] = 8
        grid['rCount'] = 32
        grid['zScale'] = 2
        grid['zLength'] = '20.324980154631'
        grid['zMin'] = '-20.324980154631'
        grid['zMax'] = '1.60'
        grid['zCellsPerWavelength'] = 8
        grid['zCount'] = 214
        del grid['xMin']
        del grid['xMax']
        del grid['xCount']
        del grid['zLambda']
    if 'rParticlesPerCell' not in data['models']['simulationGrid']:
        data['models']['simulationGrid']['rParticlesPerCell'] = 1
        data['models']['simulationGrid']['zParticlesPerCell'] = 2
    if 'field' not in data['models']['laserPreviewReport']:
        laserPreview = data['models']['laserPreviewReport']
        laserPreview['field'] = 'E'
        laserPreview['coordinate'] = 'r'
        laserPreview['mode'] = '1'

def generate_parameters_file(data, schema, persistent_files_dir=None):
    _validate_data(data, schema)
    v = template_common.flatten_data(data['models'], {})
    if persistent_files_dir:
        if not persistent_files_dir.check():
            pkio.mkdir_parent(persistent_files_dir)
        v['outputDir'] = '"{}"'.format(persistent_files_dir)
    else:
        v['outputDir'] = None
    v['enablePlasma'] = 1
    return pkjinja.render_resource('warp.py', v)

def get_simulation_frame(persistent_files_dir, data):
    frame_index = int(data['frame_index'])
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    filename = files[frame_index]
    iteration = int(re.search(r'data(\d+)', filename).group(1))
    dfile = h5py.File(filename, "r")
    args = data['animation_args'].split('_')

    if data['model_name'] == 'fieldAnimation':
        return _field_animation(args, dfile, iteration, len(files))
    if data['model_name'] == 'particleAnimation':
        return _particle_animation(args, dfile, iteration, len(files))
    raise RuntimeError('{}: unknown simulation frame model'.format(data['model_name']))

def new_simulation(data, new_simulation_data):
    pass

def prepare_aux_files(wd):
    pass

def remove_last_frame(persistent_files_dir):
    files = _h5_file_list(str(persistent_files_dir.join('diags/hdf5')))
    if len(files) > 0:
        os.remove(files[-1])

def run_all_text():
    return '''
doit = True
while(doit):
    step(10)
    doit = ( w3d.zmmin + top.zgrid < Lplasma )
'''

def _field_animation(args, dfile, iteration, frame_count):
    field = args[0]
    coordinate = args[1]
    mode = int(args[2])
    res = extract_field_report(field, coordinate, mode, dfile, iteration)
    res['frameCount'] = frame_count
    return res

def _h5_file_list(path_to_dir):
    if not os.path.isdir(path_to_dir):
        return []
    all_files = os.listdir(path_to_dir)

    # Select the hdf5 files
    h5_files = []
    for filename in all_files :
        if filename[-3:] == '.h5' :
            h5_files.append( os.path.join( path_to_dir, filename) )
    # Sort them
    h5_files.sort()
    return h5_files

def _iteration_title(dfile, iteration):
    return '{:.1f} fs (iteration {})'.format(
        iteration * 1e15 * float(dfile.attrs['timeStepUnitSI']), iteration)

def _particle_animation(args, dfile, iteration, frame_count):
    xarg = args[0]
    yarg = args[1]
    histogramBins = args[2]
    x = dfile['particles/electrons/{}'.format(_PARTICLE_ARG_PATH[xarg])][:]
    y = dfile['particles/electrons/{}'.format(_PARTICLE_ARG_PATH[yarg])][:]
    hist, edges = np.histogramdd([x, y], int(histogramBins))
    xunits = ' [m]' if len(xarg) == 1 else ''
    yunits = ' [m]' if len(yarg) == 1 else ''
    return {
        'x_range': [float(edges[0][0]), float(edges[0][-1]), len(hist)],
        'y_range': [float(edges[1][0]), float(edges[1][-1]), len(hist[0])],
        'x_label': '{}{}'.format(xarg, xunits),
        'y_label': '{}{}'.format(yarg, yunits),
        'title': 't = {}'.format(_iteration_title(dfile, iteration)),
        'z_matrix': hist.T.tolist(),
        'frameCount': frame_count,
    }

def _validate_data(data, schema):
    # ensure enums match, convert ints/floats, apply scaling
    enum_info = template_common.parse_enums(schema['enum'])
    for k in data['models']:
        if k in schema['model']:
            template_common.validate_model(data['models'][k], schema['model'][k], enum_info)
<|endoftext|>