<commit_before>from django.contrib.auth.decorators import login_required
from django.core.exceptions import ObjectDoesNotExist
from django.core.urlresolvers import reverse
from django.db.models import Q, Count
from django.http import HttpResponse, HttpResponseServerError
from django.shortcuts import render_to_response, redirect
from django.template import RequestContext
from django.utils import simplejson

import json
import math

from foodmarks.fm.constants import *
from foodmarks.fm.forms import *
from foodmarks.fm.models import *

PAGE_SIZE = 50

class JsonResponse(HttpResponse):
    '''
    this class makes it easier to quickly turn content into json and return it as a response
    '''

    def __init__(self, content, mimetype="application/json", *args, **kwargs):
        if content:
            content = simplejson.dumps(content)
        else:
            content = simplejson.dumps({})
        super(JsonResponse, self).__init__(content, mimetype, *args, **kwargs)

def index(request):
    ctx = RequestContext(request)
    ctx['recipes'] = Recipe.objects.order_by('-time_created')[0:PAGE_SIZE]
    return render_to_response('index.html', ctx)

@login_required(login_url="/accounts/login/")
def bookmarklet(request):
    ctx = RequestContext(request)
    ctx['add'] = True

    recipe = None
    saved = _save_recipe(request, c, recipe=recipe)

    if saved:
        return redirect(reverse(my_recipes), permanent=True)

    return render_to_response('bookmarklet.html', ctx)

@login_required(login_url="/accounts/login/")
def add_recipe(request):
    ctx = RequestContext(request)
    ctx['add'] = True

    recipe = None
    if request.method == 'GET':
        recipe_id = request.GET.get('recipe', None)
        if recipe_id is not None:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                pass

    saved = _save_recipe(request, c, recipe=recipe)

    if saved:
        return redirect(reverse(my_recipes), permanent=True)
    else:
        return render_to_response('edit_recipe.html', ctx)

def _save_recipe(request, c, ribbon=None, recipe=None):
    if request.method == 'GET' and not recipe:
        url = request.GET.get('url', None)
        title = request.GET.get('title', None)
        if url:
            try:
                recipe = Recipe.objects.get(link=url)
            except ObjectDoesNotExist:
                pass

    if request.user.is_staff:
        ctx['users'] = User.objects.all()
        ctx['user_id'] = request.user.id
        specified_user = request.user
        if request.POST:
            ctx['user_id'] = int(request.POST.get('user', request.user.id))
            try:
                specified_user = User.objects.get(id=ctx['user_id'])
            except ObjectDoesNotExist:
                pass

    ctx['known_values_to_keys'] = known_values_to_keys
    if ribbon and not recipe:
        recipe = ribbon.recipe

    if not recipe and request.method == 'POST':
        recipe_id = request.POST.get('recipe-id', None)
        if recipe_id:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                pass
        else:
            link = request.POST.get('re-link', None)
            if link:
                try:
                    recipe = Recipe.objects.get(link=link)
                except ObjectDoesNotExist:
                    pass

    if recipe and not ribbon:
        try:
            ribbon = Ribbon.objects.get(recipe=recipe, user=request.user)
        except ObjectDoesNotExist:
            pass

    recipe_form = RecipeForm(request.POST or None, prefix="re",
                             instance=recipe)
    ribbon_form = RibbonForm(request.POST or None, prefix="ri",
                             instance=ribbon)

    saved = False
    if recipe_form.is_valid() and ribbon_form.is_valid():
        recipe = recipe_form.save(commit=False)

        '''
        if recipe.link:
            other = Recipe.objects.filter(link=recipe.link)
            if other and other.id != recipe.id:
                recipe = RecipeForm(
                    request.POST, prefix="re", instance=other).save(
                    commit=False)
                    '''

        recipe.save()
        ribbon = ribbon_form.save(commit=False)
        ribbon.recipe = recipe
        if request.user.is_staff:
            ribbon.user = specified_user
        else:
            ribbon.user = request.user
        ribbon.save()

        if request.POST:
            tags = json.loads(request.POST['tag-json'])
            for key in tags:
                for value in tags[key]:
                    if value == '':
                        value = None
                    if tags[key][value].get('id', False):
                        tag = Tag.objects.get(id=tags[key][value]['id'])
                        if tags[key][value].get('deleted', False):
                            tag.delete()
                    elif not tags[key][value].get('deleted', False):
                        Tag(ribbon=ribbon, key=key, value=value).save()
        saved = True
    else:
        tags = {}
        if ribbon:
            actual_tags = ribbon.tag_set.all()
            for actual_tag in actual_tags:
                if not actual_tag.key in tags:
                    tags[actual_tag.key] = {}
                tags[actual_tag.key][actual_tag.value] = \
                    {'id': actual_tag.id}
        elif recipe and request.user.get_profile().copy_tags:
            actual_tags = Tag.objects.filter(ribbon__recipe=recipe)
            for actual_tag in actual_tags:
                if not actual_tag.key in tags:
                    tags[actual_tag.key] = {}
                tags[actual_tag.key][actual_tag.value] = {}

    ctx['recipe_form'] = recipe_form
    ctx['ribbon_form'] = ribbon_form
    ctx['ribbon'] = ribbon
    ctx['recipe'] = recipe
    ctx['tags'] = tags

    if request.method == 'GET':
        url = request.GET.get('url', None)
        title = request.GET.get('title', None)
        if not recipe:
            if url:
                ctx['recipe_form'].initial['link'] = url
            if title:
                ctx['recipe_form'].initial['title'] = title

    return saved

@login_required(login_url="/accounts/login/")
def edit_recipe(request, ribbon_id):
    ctx = RequestContext(request)
    ctx['edit'] = True
    try:
        ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
    except ObjectDoesNotExist:
        return redirect(reverse(my_recipes), permanent=True)

    saved = _save_recipe(request, c, ribbon)
    if saved:
        ctx['message'] = 'Recipe successfully saved.'
    return render_to_response('edit_recipe.html', ctx)

def _paginate_content(request, ctx, key="ribbons"):
    if request.method == 'GET':
        page = request.GET.get('page', 1)
        if page == 'all':
            ctx[page] = 'all'
            return
        if not page:
            page = 1
        page = int(page)
        page_size = int(request.GET.get('page_size', PAGE_SIZE))
    else:
        page = 1
        page_size = PAGE_SIZE

    ctx['num_pages'] = int(math.ceil(ctx[key].count() / float(page_size)))

    ctx[key] = ctx[key][(page - 1) * page_size: page * page_size]
    if key == 'ribbons':
        ctx['ribbons'] = ctx['ribbons'].select_related('recipe')

    ctx['page_range'] = xrange(1, ctx['num_pages'] + 1)
    ctx['page'] = page

@login_required(login_url="/accounts/login/")
def my_recipes(request):
    ctx = RequestContext(request)
    ctx['ribbons'] = Ribbon.objects.filter(user=request.user)
    _paginate_content(request, ctx, 'ribbons')
    return render_to_response('my_recipes.html', ctx)

def recipe(request, recipe_id):
    ctx = RequestContext(request)

    try:
        recipe = Recipe.objects.get(id=recipe_id)
    except ObjectDoesNotExist:
        pass

    if request.user.is_authenticated():
        try:
            ribbon = Ribbon.objects.get(recipe=recipe, user=request.user)
            ctx['ribbon'] = ribbon
        except ObjectDoesNotExist:
            pass

    ctx['recipe'] = recipe

    return render_to_response('recipe.html', ctx)

def _get_key_order(val):
    key = val[0]
    try:
        index = ordered_known_keys.index(key)
        return '{0: 3d}'.format(index)
    except ValueError:
        return '{0: 3d}{1}'.format(len(ordered_known_keys), key)

def search_recipes(request):
    ctx = RequestContext(request)

    if request.method == 'GET':
        all_ribbons = request.GET.get('all', False) or \
            not request.user.is_authenticated()
        if all_ribbons:
            ribbons = Ribbon.objects.all()
        else:
            ribbons = Ribbon.objects.filter(user=request.user)

        query_string = request.GET.get('q', '')
        if query_string != '':
            recipe_match_q = Q()

            query_tokens = query_string.split(' ')
            for token in query_tokens:
                if token:
                    recipe_match_q |= Q(recipe__title__icontains=token)
            ctx['q'] = query_string
            ribbons = ribbons.filter(recipe_match_q)

        ctx['own_ribbons'] = not all_ribbons
        selected_tags = request.GET.getlist('tag')
        if selected_tags:
            split_selected_tags = map(lambda s:s.split(':'), selected_tags)

            for split_selected_tag in split_selected_tags:
                ribbons = ribbons.filter(tag__key=split_selected_tag[0], tag__value=split_selected_tag[1])

        ctx['ribbons'] = ribbons
        _paginate_content(request, ctx, 'ribbons')

        tags = Tag.objects.filter(ribbon__in=ribbons).values('key', 'value').annotate(count=Count('value')).order_by('value')
        tag_map = {}
        priority = 0
        for tag in tags:
            if not tag['key'] in tag_map:
                tag_map[tag['key']] = {}
            value = {'value': tag['value'], 'count': tag['count'],
                     'priority': priority}
            if tag['key'] + ':' + tag['value'] in selected_tags:
                value['selected'] = True
            tag_map[tag['key']][tag['value']] = value
            priority += 1

        if all_ribbons:
            ctx['recipes'] = Recipe.objects.filter(ribbon__in=ribbons)
            _paginate_content(request, ctx, 'recipes')

    tags_ordered = []
    for key, values in tag_map.items():
        tags_ordered.append((key, sorted(values.values(),
                                         key=lambda a: a['priority']),))
    ctx['search_tags'] = sorted(tags_ordered, key=_get_key_order)

    return render_to_response('search.html', ctx)

@login_required(login_url="/accounts/login/")
def delete_ribbon(request, ribbon_id):
    try:
        ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
    except ObjectDoesNotExist:
        return HttpResponse('ERROR', mimetype="application/json")
    recipe = ribbon.recipe
    ribbon.delete()
    if not recipe.ribbon_set.exists():
        recipe.delete();
    return HttpResponse('OK', mimetype="application/json")

@login_required(login_url="/accounts/login/")
def recipe_box(request):
    ctx = RequestContext(request)
    ctx['ribbons'] = Ribbon.objects.filter(
        user=request.user, is_boxed=True).select_related(
        'recipe').order_by('-time_created')
    return render_to_response('recipe_box.html', ctx)

@login_required(login_url="/accounts/login/")
def action(request):
    if not request.method == 'POST':
        return JsonResponse({'status':'OK'})
    action = request.POST.get('action', None)
    if action == 'deleteRibbon':
        recipe_id = request.POST.get('recipeId', None)
        ribbon_id = request.POST.get('ribbonId', None)
        try:
            ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
        except ObjectDoesNotExist:
            return HttpResponse('ERROR', mimetype="application/json")
        recipe = ribbon.recipe
        ribbon.delete()
        if not recipe.ribbon_set.exists():
            recipe.delete()
        return HttpResponse('OK', mimetype="application/json")
    elif action == 'changeBoxStatus':
        new_status = request.POST.get('newStatus', None) == 'true'
        recipe_id = request.POST.get('recipeId', None)
        ribbon_id = request.POST.get('ribbonId', None)
        if recipe_id:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                return JsonResponse({'status': 'FAIL'})
        else:
            return JsonResponse({'status': 'FAIL'})
        copy_tags = False
        try:
            ribbon = Ribbon.objects.get(id=ribbon_id)
            if ribbon.recipe != recipe:
                return JsonResponse({'status': 'FAIL'})
        except ObjectDoesNotExist:
            ribbon = Ribbon(recipe=recipe, user=request.user)
            copy_tags = request.user.get_profile().copy_tags
        ribbon.is_boxed = new_status
        ribbon.save()
        if copy_tags:
            tags = Tag.objects.filter(ribbon__recipe=recipe)
            for tag in tags:
                Tag(key=tag.key, value=tag.value, ribbon=ribbon).save()

        return JsonResponse({'status': 'OK', 'ribbonId': ribbon.id})

    else:
        return JsonResponse({'status':'OK'})

@login_required(login_url="/accounts/login/")
def get_tag_category(request):
    if not request.method == 'GET':
        return JsonResponse({'status': 'OK'})
    value = request.GET.get('value', None)
    if not value:
        return JsonResponse({'status': 'OK', 'categories': ['']})
    seen_categories = set()
    categories = []

    users_tags = Tag.objects.filter(value=value, ribbon__user=request.user).values('key').annotate(count=Count('key')).order_by('-count')
    for tag in users_tags:
        categories.append(tag['key'])
        seen_categories.add(tag['key'])

    known_key = known_values_to_keys.get(value, None)
    if known_key and not known_key in seen_categories:
        categories.append(known_key)
        seen_categories.add(known_key)

    other_tags = Tag.objects.filter(value=value).values('key').annotate(count=Count('key')).order_by('-count')
    for tag in other_tags:
        if not tag['key'] in seen_categories:
            categories.append(tag['key'])
            seen_categories.add(tag['key'])

    return JsonResponse({'status': 'OK', 'categories': categories})
<commit_msg>hotfix<commit_after>from django.contrib.auth.decorators import login_required
from django.core.exceptions import ObjectDoesNotExist
from django.core.urlresolvers import reverse
from django.db.models import Q, Count
from django.http import HttpResponse, HttpResponseServerError
from django.shortcuts import render_to_response, redirect
from django.template import RequestContext
from django.utils import simplejson

import json
import math

from foodmarks.fm.constants import *
from foodmarks.fm.forms import *
from foodmarks.fm.models import *

PAGE_SIZE = 50

class JsonResponse(HttpResponse):
    '''
    this class makes it easier to quickly turn content into json and return it as a response
    '''

    def __init__(self, content, mimetype="application/json", *args, **kwargs):
        if content:
            content = simplejson.dumps(content)
        else:
            content = simplejson.dumps({})
        super(JsonResponse, self).__init__(content, mimetype, *args, **kwargs)

def index(request):
    ctx = RequestContext(request)
    ctx['recipes'] = Recipe.objects.order_by('-time_created')[0:PAGE_SIZE]
    return render_to_response('index.html', ctx)

@login_required(login_url="/accounts/login/")
def bookmarklet(request):
    ctx = RequestContext(request)
    ctx['add'] = True

    recipe = None
    saved = _save_recipe(request, ctx, recipe=recipe)

    if saved:
        return redirect(reverse(my_recipes), permanent=True)

    return render_to_response('bookmarklet.html', ctx)

@login_required(login_url="/accounts/login/")
def add_recipe(request):
    ctx = RequestContext(request)
    ctx['add'] = True

    recipe = None
    if request.method == 'GET':
        recipe_id = request.GET.get('recipe', None)
        if recipe_id is not None:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                pass

    saved = _save_recipe(request, ctx, recipe=recipe)

    if saved:
        return redirect(reverse(my_recipes), permanent=True)
    else:
        return render_to_response('edit_recipe.html', ctx)

def _save_recipe(request, ctx, ribbon=None, recipe=None):
    if request.method == 'GET' and not recipe:
        url = request.GET.get('url', None)
        title = request.GET.get('title', None)
        if url:
            try:
                recipe = Recipe.objects.get(link=url)
            except ObjectDoesNotExist:
                pass

    if request.user.is_staff:
        ctx['users'] = User.objects.all()
        ctx['user_id'] = request.user.id
        specified_user = request.user
        if request.POST:
            ctx['user_id'] = int(request.POST.get('user', request.user.id))
            try:
                specified_user = User.objects.get(id=ctx['user_id'])
            except ObjectDoesNotExist:
                pass

    ctx['known_values_to_keys'] = known_values_to_keys
    if ribbon and not recipe:
        recipe = ribbon.recipe

    if not recipe and request.method == 'POST':
        recipe_id = request.POST.get('recipe-id', None)
        if recipe_id:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                pass
        else:
            link = request.POST.get('re-link', None)
            if link:
                try:
                    recipe = Recipe.objects.get(link=link)
                except ObjectDoesNotExist:
                    pass

    if recipe and not ribbon:
        try:
            ribbon = Ribbon.objects.get(recipe=recipe, user=request.user)
        except ObjectDoesNotExist:
            pass

    recipe_form = RecipeForm(request.POST or None, prefix="re",
                             instance=recipe)
    ribbon_form = RibbonForm(request.POST or None, prefix="ri",
                             instance=ribbon)

    saved = False
    if recipe_form.is_valid() and ribbon_form.is_valid():
        recipe = recipe_form.save(commit=False)

        '''
        if recipe.link:
            other = Recipe.objects.filter(link=recipe.link)
            if other and other.id != recipe.id:
                recipe = RecipeForm(
                    request.POST, prefix="re", instance=other).save(
                    commit=False)
                    '''

        recipe.save()
        ribbon = ribbon_form.save(commit=False)
        ribbon.recipe = recipe
        if request.user.is_staff:
            ribbon.user = specified_user
        else:
            ribbon.user = request.user
        ribbon.save()

        if request.POST:
            tags = json.loads(request.POST['tag-json'])
            for key in tags:
                for value in tags[key]:
                    if value == '':
                        value = None
                    if tags[key][value].get('id', False):
                        tag = Tag.objects.get(id=tags[key][value]['id'])
                        if tags[key][value].get('deleted', False):
                            tag.delete()
                    elif not tags[key][value].get('deleted', False):
                        Tag(ribbon=ribbon, key=key, value=value).save()
        saved = True
    else:
        tags = {}
        if ribbon:
            actual_tags = ribbon.tag_set.all()
            for actual_tag in actual_tags:
                if not actual_tag.key in tags:
                    tags[actual_tag.key] = {}
                tags[actual_tag.key][actual_tag.value] = \
                    {'id': actual_tag.id}
        elif recipe and request.user.get_profile().copy_tags:
            actual_tags = Tag.objects.filter(ribbon__recipe=recipe)
            for actual_tag in actual_tags:
                if not actual_tag.key in tags:
                    tags[actual_tag.key] = {}
                tags[actual_tag.key][actual_tag.value] = {}

    ctx['recipe_form'] = recipe_form
    ctx['ribbon_form'] = ribbon_form
    ctx['ribbon'] = ribbon
    ctx['recipe'] = recipe
    ctx['tags'] = tags

    if request.method == 'GET':
        url = request.GET.get('url', None)
        title = request.GET.get('title', None)
        if not recipe:
            if url:
                ctx['recipe_form'].initial['link'] = url
            if title:
                ctx['recipe_form'].initial['title'] = title

    return saved

@login_required(login_url="/accounts/login/")
def edit_recipe(request, ribbon_id):
    ctx = RequestContext(request)
    ctx['edit'] = True
    try:
        ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
    except ObjectDoesNotExist:
        return redirect(reverse(my_recipes), permanent=True)

    saved = _save_recipe(request, ctx, ribbon)
    if saved:
        ctx['message'] = 'Recipe successfully saved.'
    return render_to_response('edit_recipe.html', ctx)

def _paginate_content(request, ctx, key="ribbons"):
    if request.method == 'GET':
        page = request.GET.get('page', 1)
        if page == 'all':
            ctx[page] = 'all'
            return
        if not page:
            page = 1
        page = int(page)
        page_size = int(request.GET.get('page_size', PAGE_SIZE))
    else:
        page = 1
        page_size = PAGE_SIZE

    ctx['num_pages'] = int(math.ceil(ctx[key].count() / float(page_size)))

    ctx[key] = ctx[key][(page - 1) * page_size: page * page_size]
    if key == 'ribbons':
        ctx['ribbons'] = ctx['ribbons'].select_related('recipe')

    ctx['page_range'] = xrange(1, ctx['num_pages'] + 1)
    ctx['page'] = page

@login_required(login_url="/accounts/login/")
def my_recipes(request):
    ctx = RequestContext(request)
    ctx['ribbons'] = Ribbon.objects.filter(user=request.user)
    _paginate_content(request, ctx, 'ribbons')
    return render_to_response('my_recipes.html', ctx)

def recipe(request, recipe_id):
    ctx = RequestContext(request)

    try:
        recipe = Recipe.objects.get(id=recipe_id)
    except ObjectDoesNotExist:
        pass

    if request.user.is_authenticated():
        try:
            ribbon = Ribbon.objects.get(recipe=recipe, user=request.user)
            ctx['ribbon'] = ribbon
        except ObjectDoesNotExist:
            pass

    ctx['recipe'] = recipe

    return render_to_response('recipe.html', ctx)

def _get_key_order(val):
    key = val[0]
    try:
        index = ordered_known_keys.index(key)
        return '{0: 3d}'.format(index)
    except ValueError:
        return '{0: 3d}{1}'.format(len(ordered_known_keys), key)

def search_recipes(request):
    ctx = RequestContext(request)

    if request.method == 'GET':
        all_ribbons = request.GET.get('all', False) or \
            not request.user.is_authenticated()
        if all_ribbons:
            ribbons = Ribbon.objects.all()
        else:
            ribbons = Ribbon.objects.filter(user=request.user)

        query_string = request.GET.get('q', '')
        if query_string != '':
            recipe_match_q = Q()

            query_tokens = query_string.split(' ')
            for token in query_tokens:
                if token:
                    recipe_match_q |= Q(recipe__title__icontains=token)
            ctx['q'] = query_string
            ribbons = ribbons.filter(recipe_match_q)

        ctx['own_ribbons'] = not all_ribbons
        selected_tags = request.GET.getlist('tag')
        if selected_tags:
            split_selected_tags = map(lambda s:s.split(':'), selected_tags)

            for split_selected_tag in split_selected_tags:
                ribbons = ribbons.filter(tag__key=split_selected_tag[0], tag__value=split_selected_tag[1])

        ctx['ribbons'] = ribbons
        _paginate_content(request, ctx, 'ribbons')

        tags = Tag.objects.filter(ribbon__in=ribbons).values('key', 'value').annotate(count=Count('value')).order_by('value')
        tag_map = {}
        priority = 0
        for tag in tags:
            if not tag['key'] in tag_map:
                tag_map[tag['key']] = {}
            value = {'value': tag['value'], 'count': tag['count'],
                     'priority': priority}
            if tag['key'] + ':' + tag['value'] in selected_tags:
                value['selected'] = True
            tag_map[tag['key']][tag['value']] = value
            priority += 1

        if all_ribbons:
            ctx['recipes'] = Recipe.objects.filter(ribbon__in=ribbons)
            _paginate_content(request, ctx, 'recipes')

    tags_ordered = []
    for key, values in tag_map.items():
        tags_ordered.append((key, sorted(values.values(),
                                         key=lambda a: a['priority']),))
    ctx['search_tags'] = sorted(tags_ordered, key=_get_key_order)

    return render_to_response('search.html', ctx)

@login_required(login_url="/accounts/login/")
def delete_ribbon(request, ribbon_id):
    try:
        ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
    except ObjectDoesNotExist:
        return HttpResponse('ERROR', mimetype="application/json")
    recipe = ribbon.recipe
    ribbon.delete()
    if not recipe.ribbon_set.exists():
        recipe.delete();
    return HttpResponse('OK', mimetype="application/json")

@login_required(login_url="/accounts/login/")
def recipe_box(request):
    ctx = RequestContext(request)
    ctx['ribbons'] = Ribbon.objects.filter(
        user=request.user, is_boxed=True).select_related(
        'recipe').order_by('-time_created')
    return render_to_response('recipe_box.html', ctx)

@login_required(login_url="/accounts/login/")
def action(request):
    if not request.method == 'POST':
        return JsonResponse({'status':'OK'})
    action = request.POST.get('action', None)
    if action == 'deleteRibbon':
        recipe_id = request.POST.get('recipeId', None)
        ribbon_id = request.POST.get('ribbonId', None)
        try:
            ribbon = Ribbon.objects.get(id=ribbon_id, user=request.user)
        except ObjectDoesNotExist:
            return HttpResponse('ERROR', mimetype="application/json")
        recipe = ribbon.recipe
        ribbon.delete()
        if not recipe.ribbon_set.exists():
            recipe.delete()
        return HttpResponse('OK', mimetype="application/json")
    elif action == 'changeBoxStatus':
        new_status = request.POST.get('newStatus', None) == 'true'
        recipe_id = request.POST.get('recipeId', None)
        ribbon_id = request.POST.get('ribbonId', None)
        if recipe_id:
            try:
                recipe = Recipe.objects.get(id=recipe_id)
            except ObjectDoesNotExist:
                return JsonResponse({'status': 'FAIL'})
        else:
            return JsonResponse({'status': 'FAIL'})
        copy_tags = False
        try:
            ribbon = Ribbon.objects.get(id=ribbon_id)
            if ribbon.recipe != recipe:
                return JsonResponse({'status': 'FAIL'})
        except ObjectDoesNotExist:
            ribbon = Ribbon(recipe=recipe, user=request.user)
            copy_tags = request.user.get_profile().copy_tags
        ribbon.is_boxed = new_status
        ribbon.save()
        if copy_tags:
            tags = Tag.objects.filter(ribbon__recipe=recipe)
            for tag in tags:
                Tag(key=tag.key, value=tag.value, ribbon=ribbon).save()

        return JsonResponse({'status': 'OK', 'ribbonId': ribbon.id})

    else:
        return JsonResponse({'status':'OK'})

@login_required(login_url="/accounts/login/")
def get_tag_category(request):
    if not request.method == 'GET':
        return JsonResponse({'status': 'OK'})
    value = request.GET.get('value', None)
    if not value:
        return JsonResponse({'status': 'OK', 'categories': ['']})
    seen_categories = set()
    categories = []

    users_tags = Tag.objects.filter(value=value, ribbon__user=request.user).values('key').annotate(count=Count('key')).order_by('-count')
    for tag in users_tags:
        categories.append(tag['key'])
        seen_categories.add(tag['key'])

    known_key = known_values_to_keys.get(value, None)
    if known_key and not known_key in seen_categories:
        categories.append(known_key)
        seen_categories.add(known_key)

    other_tags = Tag.objects.filter(value=value).values('key').annotate(count=Count('key')).order_by('-count')
    for tag in other_tags:
        if not tag['key'] in seen_categories:
            categories.append(tag['key'])
            seen_categories.add(tag['key'])

    return JsonResponse({'status': 'OK', 'categories': categories})
<|endoftext|>


<commit_before>"""Subroutines / ingredients which are not specific to any particular recipe

TODO: This is implementation specific and should be moved outside this module.

"""
import logging
import sys
from lofarpipe.support.clusterdesc import ClusterDesc, get_compute_nodes
import itertools
from lofarpipe.support.lofarnode import LOFARnodeTCP
from lofarpipe.support.baserecipe import BaseRecipe
from lofarpipe.support.remotecommand import RemoteCommandRecipeMixIn
from lofarpipe.support.pipelinelogging import log_time

def nodes_available(config):
    """ Obtain list of available nodes """
    clusterdesc = ClusterDesc(config.get('cluster', "clusterdesc"))
    if clusterdesc.subclusters:
        available_nodes = dict(
            (cl.name, itertools.cycle(get_compute_nodes(cl)))
            for cl in clusterdesc.subclusters
            )
    else:
        available_nodes = {
            clusterdesc.name: get_compute_nodes(clusterdesc)
            }
    return list(itertools.chain(*available_nodes.values()))

class TrapNode(LOFARnodeTCP):
    """
    This class should be inherited to create a Trap Node recipe. The
    trapstep method should be used from implementing your recipe node logic.
    """
    def run(self, *args, **kwargs):
        with log_time(self.logger):
            # capture all logging and sent it to the master
            logdrain = logging.getLogger()
            logdrain.level = self.logger.level
            logdrain.handlers = self.logger.handlers
            self.logger.handlers = []
            #self.logger = logdrain

            try:
                self.trapstep(*args, **kwargs)
            except Exception,e:
                self.logger.error(e)
                return 1
        return 0

    def trapstep(self):
        raise NotImplementedError()

class TrapMaster(BaseRecipe, RemoteCommandRecipeMixIn):
    """
    This class should be inherited to create a Trap Master recipe. The
    trapstep method should be used from implementing your recipe node logic.
    """
    def go(self, *args, **kwargs):
        super(TrapMaster, self).go()
        with log_time(self.logger):
            #capture all logging, not only trap namespace
            #logdrain = logging.getLogger()
            #logdrain.level = self.logger.level
            #logdrain.handlers = self.logger.handlers
            #self.logger = logdrain

            # call the actual do-er
            self.trapstep(*args, **kwargs)

        if self.error.isSet():
            self.logger.warn("Failed null_detections process detected")
            return 1
        else:
            return 0

    def trapstep(self):
        raise NotImplementedError()

def node_run(name, nodeclass):
    """this function should be called from every Trap Node Recipe so it can
    be called on a node.
    """
    if name == "__main__":
        jobid, jobhost, jobport = sys.argv[1:4]
        sys.exit(nodeclass(jobid, jobhost, jobport).run_with_stored_arguments())<commit_msg>we don't want to do this here<commit_after>"""Subroutines / ingredients which are not specific to any particular recipe

TODO: This is implementation specific and should be moved outside this module.

"""
import logging
import sys
from lofarpipe.support.clusterdesc import ClusterDesc, get_compute_nodes
import itertools
from lofarpipe.support.lofarnode import LOFARnodeTCP
from lofarpipe.support.baserecipe import BaseRecipe
from lofarpipe.support.remotecommand import RemoteCommandRecipeMixIn
from lofarpipe.support.pipelinelogging import log_time

def nodes_available(config):
    """ Obtain list of available nodes """
    clusterdesc = ClusterDesc(config.get('cluster', "clusterdesc"))
    if clusterdesc.subclusters:
        available_nodes = dict(
            (cl.name, itertools.cycle(get_compute_nodes(cl)))
            for cl in clusterdesc.subclusters
            )
    else:
        available_nodes = {
            clusterdesc.name: get_compute_nodes(clusterdesc)
            }
    return list(itertools.chain(*available_nodes.values()))

class TrapNode(LOFARnodeTCP):
    """
    This class should be inherited to create a Trap Node recipe. The
    trapstep method should be used from implementing your recipe node logic.
    """
    def run(self, *args, **kwargs):
        with log_time(self.logger):
            # capture all logging and sent it to the master
            logdrain = logging.getLogger()
            logdrain.level = self.logger.level
            logdrain.handlers = self.logger.handlers
            self.logger.handlers = []
            #self.logger = logdrain

            try:
                self.trapstep(*args, **kwargs)
            except Exception,e:
                self.logger.error(e)
                return 1
        return 0

    def trapstep(self):
        raise NotImplementedError()

class TrapMaster(BaseRecipe, RemoteCommandRecipeMixIn):
    """
    This class should be inherited to create a Trap Master recipe. The
    trapstep method should be used from implementing your recipe node logic.
    """
    def go(self, *args, **kwargs):
        super(TrapMaster, self).go()
        # call the actual do-er
        self.trapstep(*args, **kwargs)
        if self.error.isSet():
            self.logger.warn("Failed null_detections process detected")
            return 1
        else:
            return 0

    def trapstep(self):
        raise NotImplementedError()

def node_run(name, nodeclass):
    """this function should be called from every Trap Node Recipe so it can
    be called on a node.
    """
    if name == "__main__":
        jobid, jobhost, jobport = sys.argv[1:4]
        sys.exit(nodeclass(jobid, jobhost, jobport).run_with_stored_arguments())<|endoftext|>


<commit_before>from django import forms
from django.forms.widgets import CheckboxSelectMultiple
from django.utils.translation import ugettext_lazy as _

from django.conf import settings
from django_comments_xtd.conf import settings as comments_settings
from django_comments_xtd.forms import CommentForm
from django_comments_xtd.models import TmpXtdComment

from .models import Job
from cms.forms import ContentManageableModelForm

class JobForm(ContentManageableModelForm):
    class Meta:
        model = Job
        fields = (
            'category',
            'job_types',
            'company',
            'city',
            'region',
            'country',
            'description',
            'requirements',
            'contact',
            'email',
            'url',
            'telecommuting',
            'agencies',
        )
        widgets = {
            'job_types': CheckboxSelectMultiple(),
        }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fields['job_types'].help_text = None
        self.fields['telecommuting'].label = 'Is telecommuting allowed?'
        self.fields['agencies'].label = 'Is job on behalf of an agency?'

class JobCommentForm(CommentForm):
    reply_to = forms.IntegerField(required=True, initial=0, widget=forms.HiddenInput())

    def __init__(self, *args, **kwargs):
        comment = kwargs.pop("comment", None)
        if comment:
            initial = kwargs.pop("initial", {})
            initial.update({"reply_to": comment.pk})
            kwargs["initial"] = initial
        super(JobCommentForm, self).__init__(*args, **kwargs)
        self.fields['name'] = forms.CharField(
            widget=forms.TextInput(attrs={'placeholder':_('name')}))
        self.fields['email'] = forms.EmailField(
            label=_("Email"), help_text=_("Required for comment verification"),
            widget=forms.TextInput(attrs={'placeholder':_('email')})
            )
        self.fields['url'] = forms.URLField(
            required=False,
            widget=forms.TextInput(attrs={'placeholder':_('website')}))
        self.fields['comment'] = forms.CharField(
            widget=forms.Textarea(attrs={'placeholder':_('comment')}),
            max_length=comments_settings.COMMENT_MAX_LENGTH)

    def get_comment_model(self):
        return TmpXtdComment

    def get_comment_create_data(self):
        data = super(JobCommentForm, self).get_comment_create_data()
        data.update({'thread_id': 0, 'level': 0, 'order': 1,
                     'parent_id': self.cleaned_data['reply_to'],
                     'followup': True})
        if comments_settings.COMMENTS_XTD_CONFIRM_EMAIL:
            # comment must be verified before getting approved
            data['is_public'] = False
        return data
<commit_msg>update form label for agency checkbox to indicate it is ok for agencies to contact the job submitter<commit_after>from django import forms
from django.forms.widgets import CheckboxSelectMultiple
from django.utils.translation import ugettext_lazy as _

from django.conf import settings
from django_comments_xtd.conf import settings as comments_settings
from django_comments_xtd.forms import CommentForm
from django_comments_xtd.models import TmpXtdComment

from .models import Job
from cms.forms import ContentManageableModelForm

class JobForm(ContentManageableModelForm):
    class Meta:
        model = Job
        fields = (
            'category',
            'job_types',
            'company',
            'city',
            'region',
            'country',
            'description',
            'requirements',
            'contact',
            'email',
            'url',
            'telecommuting',
            'agencies',
        )
        widgets = {
            'job_types': CheckboxSelectMultiple(),
        }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fields['job_types'].help_text = None
        self.fields['telecommuting'].label = 'Is telecommuting allowed?'
        self.fields['agencies'].label = 'Agencies are OK to contact?'

class JobCommentForm(CommentForm):
    reply_to = forms.IntegerField(required=True, initial=0, widget=forms.HiddenInput())

    def __init__(self, *args, **kwargs):
        comment = kwargs.pop("comment", None)
        if comment:
            initial = kwargs.pop("initial", {})
            initial.update({"reply_to": comment.pk})
            kwargs["initial"] = initial
        super(JobCommentForm, self).__init__(*args, **kwargs)
        self.fields['name'] = forms.CharField(
            widget=forms.TextInput(attrs={'placeholder':_('name')}))
        self.fields['email'] = forms.EmailField(
            label=_("Email"), help_text=_("Required for comment verification"),
            widget=forms.TextInput(attrs={'placeholder':_('email')})
            )
        self.fields['url'] = forms.URLField(
            required=False,
            widget=forms.TextInput(attrs={'placeholder':_('website')}))
        self.fields['comment'] = forms.CharField(
            widget=forms.Textarea(attrs={'placeholder':_('comment')}),
            max_length=comments_settings.COMMENT_MAX_LENGTH)

    def get_comment_model(self):
        return TmpXtdComment

    def get_comment_create_data(self):
        data = super(JobCommentForm, self).get_comment_create_data()
        data.update({'thread_id': 0, 'level': 0, 'order': 1,
                     'parent_id': self.cleaned_data['reply_to'],
                     'followup': True})
        if comments_settings.COMMENTS_XTD_CONFIRM_EMAIL:
            # comment must be verified before getting approved
            data['is_public'] = False
        return data
<|endoftext|>


<commit_before>import functools
import json
import warnings
import weakref

from cached_property import cached_property
from docker import utils as docker_utils, auth as docker_auth

import fabricio

from fabricio.utils import Options

class Image(object):

    container_options_mapping = (
        ('user', 'user'),
        ('publish', 'ports'),
        ('env', 'env'),
        ('volume', 'volumes'),
        ('link', 'links'),
        ('add-host', 'hosts'),
        ('net', 'network'),
        ('restart', 'restart_policy'),
        ('stop-signal', 'stop_signal'),
    )

    def __init__(self, name=None, tag=None, registry=None):
        if name:
            _registry, _name, _tag = self.parse_image_name(name)
            self.name = _name
            self.tag = tag or _tag or 'latest'  # TODO 'latest' is unnecessary
            self.registry = registry or _registry
        else:
            self.name = name
            self.tag = tag
            self.registry = registry
        self.field_names = {}
        self.container = None

    def __str__(self):
        if self.container is not None:
            return self.id
        return repr(self)

    def __repr__(self):
        if self.registry:
            return '{registry}/{name}:{tag}'.format(
                registry=self.registry,
                name=self.name,
                tag=self.tag,
            )
        return '{name}:{tag}'.format(name=self.name, tag=self.tag)

    def __get__(self, container, owner_cls):
        if container is None:
            return self
        field_name = self.get_field_name(owner_cls)
        image = container.__dict__.get(field_name)
        if image is None:
            image = container.__dict__[field_name] = self.__class__(
                name=self.name,
                tag=self.tag,
                registry=self.registry,
            )
        image.container = weakref.proxy(container)  # TODO maybe this is unnecessary
        return image

    def __set__(self, container, image):
        field_name = self.get_field_name(type(container))
        container.__dict__[field_name] = (
            image[:]  # actually this means copy of image
            if isinstance(image, Image) else
            self.__class__(name=image)
        )

    def __getitem__(self, item):
        if isinstance(item, slice):
            registry, tag = item.start, item.stop
        else:
            registry, tag = None, item
        return self.__class__(
            name=self.name,
            tag=tag or self.tag,
            registry=registry or self.registry,
        )

    def get_field_name(self, owner_cls):
        field_name = self.field_names.get(owner_cls)
        if field_name is None:
            for attr in dir(owner_cls):
                if getattr(owner_cls, attr) is self:
                    if field_name is not None:
                        raise ValueError(
                            'Same instance of Image used for more than one '
                            'attributes of class {cls}'.format(
                                cls=owner_cls.__name__,
                            )
                        )
                    self.field_names[owner_cls] = field_name = attr
        return field_name

    @staticmethod
    def parse_image_name(image):
        repository, tag = docker_utils.parse_repository_tag(image)
        registry, name = docker_auth.resolve_repository_name(repository)
        if registry == docker_auth.INDEX_NAME:
            registry = ''
        return registry, name, tag

    @classmethod
    def make_container_options(cls, temporary=None, name=None, options=()):
        options = dict(options)
        container_options = Options()
        for remap, option in cls.container_options_mapping:
            container_options[remap] = options.pop(option, None)
        container_options.update(
            (
                ('name', name),
                ('rm', temporary),
                ('tty', temporary),
                ('detach', temporary is not None and not temporary),
            ),
            **options
        )
        return container_options

    @property
    def info(self):
        command = 'docker inspect --type image {image}'
        info = fabricio.run(command.format(image=self))
        return json.loads(str(info))[0]

    @cached_property
    def id(self):
        if self.container is None:
            return self.info.get('Id')
        try:
            return self.container.info['Image']
        except ReferenceError:
            raise ReferenceError(
                'weakly-referenced property `container` no longer exists, '
                'usually this happens when accessing container\'s image w/o '
                'storing container in variable, e.g. when using method chaining'
            )

    def delete(self, force=False, ignore_errors=False, deferred=False):
        command = 'docker rmi {force}{image}'
        force = force and '--force ' or ''
        callback = functools.partial(
            fabricio.run,
            command.format(image=self, force=force),
            ignore_errors=ignore_errors,
        )
        if deferred:
            return callback
        callback()

    def run(
        self,
        cmd=None,
        temporary=True,
        quiet=True,
        name=None,
        options=(),
        **kwargs
    ):
        if kwargs:
            warnings.warn(
                'Container options must be provided in `options` arg, '
                'kwargs behavior will be removed in v0.4',
                category=RuntimeWarning, stacklevel=2,
            )
            options = dict(options, **kwargs)
        command = 'docker run {options} {image} {cmd}'
        return fabricio.run(
            command.format(
                image=self,
                cmd=cmd or '',
                options=self.make_container_options(
                    temporary=temporary,
                    name=name,
                    options=options,
                ),
            ),
            quiet=quiet,
        )
<commit_msg>removed weakref usage<commit_after>import functools
import json
import warnings

from cached_property import cached_property
from docker import utils as docker_utils, auth as docker_auth

import fabricio

from fabricio.utils import Options

class Image(object):

    container_options_mapping = (
        ('user', 'user'),
        ('publish', 'ports'),
        ('env', 'env'),
        ('volume', 'volumes'),
        ('link', 'links'),
        ('add-host', 'hosts'),
        ('net', 'network'),
        ('restart', 'restart_policy'),
        ('stop-signal', 'stop_signal'),
    )

    def __init__(self, name=None, tag=None, registry=None):
        if name:
            _registry, _name, _tag = self.parse_image_name(name)
            self.name = _name
            self.tag = tag or _tag or 'latest'  # TODO 'latest' is unnecessary
            self.registry = registry or _registry
        else:
            self.name = name
            self.tag = tag
            self.registry = registry
        self.field_names = {}
        self.container = None

    def __str__(self):
        if self.container is not None:
            return self.id
        return repr(self)

    def __repr__(self):
        if self.registry:
            return '{registry}/{name}:{tag}'.format(
                registry=self.registry,
                name=self.name,
                tag=self.tag,
            )
        return '{name}:{tag}'.format(name=self.name, tag=self.tag)

    def __get__(self, container, owner_cls):
        if container is None:
            return self
        field_name = self.get_field_name(owner_cls)
        image = container.__dict__.get(field_name)
        if image is None:
            image = container.__dict__[field_name] = self.__class__(
                name=self.name,
                tag=self.tag,
                registry=self.registry,
            )
        # this cause circular reference between container and image, but it's
        # not a problem due to temporary nature of Fabric runtime
        image.container = container
        return image

    def __set__(self, container, image):
        field_name = self.get_field_name(type(container))
        container.__dict__[field_name] = (
            image[:]  # actually this means copy of image
            if isinstance(image, Image) else
            self.__class__(name=image)
        )

    def __getitem__(self, item):
        if isinstance(item, slice):
            registry, tag = item.start, item.stop
        else:
            registry, tag = None, item
        return self.__class__(
            name=self.name,
            tag=tag or self.tag,
            registry=registry or self.registry,
        )

    def get_field_name(self, owner_cls):
        field_name = self.field_names.get(owner_cls)
        if field_name is None:
            for attr in dir(owner_cls):
                if getattr(owner_cls, attr) is self:
                    if field_name is not None:
                        raise ValueError(
                            'Same instance of Image used for more than one '
                            'attributes of class {cls}'.format(
                                cls=owner_cls.__name__,
                            )
                        )
                    self.field_names[owner_cls] = field_name = attr
        return field_name

    @staticmethod
    def parse_image_name(image):
        repository, tag = docker_utils.parse_repository_tag(image)
        registry, name = docker_auth.resolve_repository_name(repository)
        if registry == docker_auth.INDEX_NAME:
            registry = ''
        return registry, name, tag

    @classmethod
    def make_container_options(cls, temporary=None, name=None, options=()):
        options = dict(options)
        container_options = Options()
        for remap, option in cls.container_options_mapping:
            container_options[remap] = options.pop(option, None)
        container_options.update(
            (
                ('name', name),
                ('rm', temporary),
                ('tty', temporary),
                ('detach', temporary is not None and not temporary),
            ),
            **options
        )
        return container_options

    @property
    def info(self):
        command = 'docker inspect --type image {image}'
        info = fabricio.run(command.format(image=self))
        return json.loads(str(info))[0]

    @cached_property
    def id(self):
        if self.container is None:
            return self.info.get('Id')
        return self.container.info['Image']

    def delete(self, force=False, ignore_errors=False, deferred=False):
        command = 'docker rmi {force}{image}'
        force = force and '--force ' or ''
        callback = functools.partial(
            fabricio.run,
            command.format(image=self, force=force),
            ignore_errors=ignore_errors,
        )
        if deferred:
            return callback
        callback()

    def run(
        self,
        cmd=None,
        temporary=True,
        quiet=True,
        name=None,
        options=(),
        **kwargs
    ):
        if kwargs:
            warnings.warn(
                'Container options must be provided in `options` arg, '
                'kwargs behavior will be removed in v0.4',
                category=RuntimeWarning, stacklevel=2,
            )
            options = dict(options, **kwargs)
        command = 'docker run {options} {image} {cmd}'
        return fabricio.run(
            command.format(
                image=self,
                cmd=cmd or '',
                options=self.make_container_options(
                    temporary=temporary,
                    name=name,
                    options=options,
                ),
            ),
            quiet=quiet,
        )
<|endoftext|>


<commit_before>#! /usr/bin/env python3

"""Flappy Bird, implemented using Pygame."""

import math
import os
from random import randint
from collections import deque

import pygame
from pygame.locals import *

FPS = 60
ANIMATION_SPEED = 0.18  # pixels per millisecond
WIN_WIDTH = 284 * 2     # BG image size: 284x512 px; tiled twice
WIN_HEIGHT = 512

class Bird(pygame.sprite.Sprite):
    """Represents the bird controlled by the player.

    The bird is the 'hero' of this game.  The player can make it climb
    (ascend quickly), otherwise it sinks (descends more slowly).  It must
    pass through the space in between pipes (for every pipe passed, one
    point is scored); if it crashes into a pipe, the game ends.

    Attributes:
    x: The bird's X coordinate.
    y: The bird's Y coordinate.
    msec_to_climb: The number of milliseconds left to climb, where a
        complete climb lasts Bird.CLIMB_DURATION milliseconds.

    Constants:
    WIDTH: The width, in pixels, of the bird's image.
    HEIGHT: The height, in pixels, of the bird's image.
    SINK_SPEED: With which speed, in pixels per millisecond, the bird
        descends in one second while not climbing.
    CLIMB_SPEED: With which speed, in pixels per millisecond, the bird
        ascends in one second while climbing, on average.  See also the
        Bird.update docstring.
    CLIMB_DURATION: The number of milliseconds it takes the bird to
        execute a complete climb.
    """

    WIDTH = HEIGHT = 32
    SINK_SPEED = 0.18
    CLIMB_SPEED = 0.3
    CLIMB_DURATION = 333.3

    def __init__(self, x, y, msec_to_climb, images):
        """Initialise a new Bird instance.

        Arguments:
        x: The bird's initial X coordinate.
        y: The bird's initial Y coordinate.
        msec_to_climb: The number of milliseconds left to climb, where a
            complete climb lasts Bird.CLIMB_DURATION milliseconds.  Use
            this if you want the bird to make a (small?) climb at the
            very beginning of the game.
        images: A tuple containing the images used by this bird.  It
            must contain the following images, in the following order:
            0. image of the bird with its wing pointing upward
            1. image of the bird with its wing pointing downward
        """
        super(Bird, self).__init__()
        self.x, self.y = x, y
        self.msec_to_climb = msec_to_climb
        self._img_wingup, self._img_wingdown = images
        self._mask_wingup = pygame.mask.from_surface(self._img_wingup)
        self._mask_wingdown = pygame.mask.from_surface(self._img_wingdown)

    def update(self, delta_frames=1):
        """Update the bird's position.

        This function uses the cosine function to achieve a smooth climb:
        In the first and last few frames, the bird climbs very little, in the
        middle of the climb, it climbs a lot.
        One complete climb lasts CLIMB_DURATION milliseconds, during which
        the bird ascends with an average speed of CLIMB_SPEED px/ms.
        This Bird's msec_to_climb attribute will automatically be
        decreased accordingly if it was > 0 when this method was called.

        Attributes:
        delta_frames: The number of frames elapsed since this method was
            last called.
        """
        if self.msec_to_climb > 0:
            frac_climb_done = ((Bird.CLIMB_DURATION - self.msec_to_climb) /
                               Bird.CLIMB_DURATION)
            self.y -= (Bird.CLIMB_SPEED * frames_to_msec(delta_frames) *
                       (1 - math.cos(frac_climb_done * math.pi)))
            self.msec_to_climb -= frames_to_msec(delta_frames)
        else:
            self.y += Bird.SINK_SPEED * frames_to_msec(delta_frames)

    @property
    def image(self):
        """Get a Surface containing this bird's image.

        This will decide whether to return an image where the bird's
        visible wing is pointing upward or where it is pointing downward
        based on pygame.time.get_ticks().  This will animate the flapping
        bird, even though pygame doesn't support animated GIFs.
        """
        if pygame.time.get_ticks() % 500 >= 250:
            return self._img_wingup
        else:
            return self._img_wingdown

    @property
    def mask(self):
        """Get a bitmask for usage in collision detection.

        The bitmask excludes all pixels in self.image with a
        transparency greater than 127."""
        if pygame.time.get_ticks() % 500 >= 250:
            return self._mask_wingup
        else:
            return self._mask_wingdown

    @property
    def rect(self):
        """Get the bird's position, width, and height, as a pygame.Rect."""
        return Rect(self.x, self.y, Bird.WIDTH, Bird.HEIGHT)

class PipePair(pygame.sprite.Sprite):
    """Represents an obstacle.

    A PipePair has a top and a bottom pipe, and only between them can
    the bird pass -- if it collides with either part, the game is over.

    Attributes:
    x: The PipePair's X position.  This is a float, to make movement
        smoother.  Note that there is no y attribute, as it will only
        ever be 0.
    image: A pygame.Surface which can be blitted to the display surface
        to display the PipePair.
    mask: A bitmask which excludes all pixels in self.image with a
        transparency greater than 127.  This can be used for collision
        detection.
    top_pieces: The number of pieces, including the end piece, in the
        top pipe.
    bottom_pieces: The number of pieces, including the end piece, in
        the bottom pipe.

    Constants:
    WIDTH: The width, in pixels, of a pipe piece.  Because a pipe is
        only one piece wide, this is also the width of a PipePair's
        image.
    PIECE_HEIGHT: The height, in pixels, of a pipe piece.
    ADD_INTERVAL: The interval, in milliseconds, in between adding new
        pipes.
    ADD_EVENT: Identifies the event in the queue which signifies that a
        new pipe should be added.
    """

    WIDTH = 80
    PIECE_HEIGHT = 32
    ADD_INTERVAL = 3000
    ADD_EVENT = USEREVENT + 1

    def __init__(self, pipe_end_img, pipe_body_img):
        """Initialises a new random PipePair.

        The new PipePair will automatically be assigned an x attribute of
        float(WIN_WIDTH - 1).

        Arguments:
        pipe_end_img: The image to use to represent a pipe's end piece.
        pipe_body_img: The image to use to represent one horizontal slice
            of a pipe's body.
        """
        self.x = float(WIN_WIDTH - 1)
        self.score_counted = False

        self.image = pygame.Surface((PipePair.WIDTH, WIN_HEIGHT), SRCALPHA)
        self.image.convert()   # speeds up blitting
        self.image.fill((0, 0, 0, 0))
        total_pipe_body_pieces = int(
            (WIN_HEIGHT -                  # fill window from top to bottom
             3 * Bird.HEIGHT -             # make room for bird to fit through
             3 * PipePair.PIECE_HEIGHT) /  # 2 end pieces + 1 body piece
            PipePair.PIECE_HEIGHT          # to get number of pipe pieces
        )
        self.bottom_pieces = randint(1, total_pipe_body_pieces)
        self.top_pieces = total_pipe_body_pieces - self.bottom_pieces

        # bottom pipe
        for i in range(1, self.bottom_pieces + 1):
            piece_pos = (0, WIN_HEIGHT - i*PipePair.PIECE_HEIGHT)
            self.image.blit(pipe_body_img, piece_pos)
        bottom_pipe_end_y = WIN_HEIGHT - self.bottom_height_px
        bottom_end_piece_pos = (0, bottom_pipe_end_y - PipePair.PIECE_HEIGHT)
        self.image.blit(pipe_end_img, bottom_end_piece_pos)

        # top pipe
        for i in range(self.top_pieces):
            self.image.blit(pipe_body_img, (0, i * PipePair.PIECE_HEIGHT))
        top_pipe_end_y = self.top_height_px
        self.image.blit(pipe_end_img, (0, top_pipe_end_y))

        # compensate for added end pieces
        self.top_pieces += 1
        self.bottom_pieces += 1

        # for collision detection
        self.mask = pygame.mask.from_surface(self.image)

    @property
    def top_height_px(self):
        """Get the top pipe's height, in pixels."""
        return self.top_pieces * PipePair.PIECE_HEIGHT

    @property
    def bottom_height_px(self):
        """Get the bottom pipe's height, in pixels."""
        return self.bottom_pieces * PipePair.PIECE_HEIGHT

    @property
    def visible(self):
        """Get whether this PipePair on screen, visible to the player."""
        return -PipePair.WIDTH < self.x < WIN_WIDTH

    @property
    def rect(self):
        """Get the Rect which contains this PipePair."""
        return Rect(self.x, 0, PipePair.WIDTH, PipePair.PIECE_HEIGHT)

    def update(self, delta_frames=1):
        """Update the PipePair's position.

        Attributes:
        delta_frames: The number of frames elapsed since this method was
            last called.
        """
        self.x -= ANIMATION_SPEED * frames_to_msec(delta_frames)

    def collides_with(self, bird):
        """Get whether the bird collides with a pipe in this PipePair.

        Arguments:
        bird: The Bird which should be tested for collision with this
            PipePair.
        """
        return pygame.sprite.collide_mask(self, bird)

def load_images():
    """Load all images required by the game and return a dict of them.

    The returned dict has the following keys:
    background: The game's background image.
    bird-wingup: An image of the bird with its wing pointing upward.
        Use this and bird-wingdown to create a flapping bird.
    bird-wingdown: An image of the bird with its wing pointing downward.
        Use this and bird-wingup to create a flapping bird.
    pipe-end: An image of a pipe's end piece (the slightly wider bit).
        Use this and pipe-body to make pipes.
    pipe-body: An image of a slice of a pipe's body.  Use this and
        pipe-body to make pipes.
    """

    def load_image(img_file_name):
        """Return the loaded pygame image with the specified file name.

        This function looks for images in the game's images folder
        (./images/).  All images are converted before being returned to
        speed up blitting.

        Arguments:
        img_file_name: The file name (including its extension, e.g.
            '.png') of the required image, without a file path.
        """
        file_name = os.path.join('.', 'images', img_file_name)
        img = pygame.image.load(file_name)
        # converting all images before use speeds up blitting
        img.convert()
        return img

    return {'background': load_image('background.png'),
            'pipe-end': load_image('pipe_end.png'),
            'pipe-body': load_image('pipe_body.png'),
            # images for animating the flapping bird -- animated GIFs are
            # not supported in pygame
            'bird-wingup': load_image('bird_wing_up.png'),
            'bird-wingdown': load_image('bird_wing_down.png'),}

def frames_to_msec(frames, fps=FPS):
    """Convert frames to milliseconds at the specified framerate.

    Arguments:
    frames: How many frames to convert to milliseconds.
    fps: The framerate to use for conversion.  Default: FPS.
    """
    return 1000.0 * frames / fps

def msec_to_frames(milliseconds, fps=FPS):
    """Convert milliseconds to frames at the specified framerate.

    Arguments:
    milliseconds: How many milliseconds to convert to frames.
    fps: The framerate to use for conversion.  Default: FPS.
    """
    return fps * milliseconds / 1000.0

def main():
    """The application's entry point.

    If someone executes this module (instead of importing it, for
    example), this function is called.
    """

    pygame.init()

    display_surface = pygame.display.set_mode((WIN_WIDTH, WIN_HEIGHT))
    pygame.display.set_caption('Pygame Flappy Bird')

    clock = pygame.time.Clock()
    score_font = pygame.font.SysFont(None, 32, bold=True)  # default font
    images = load_images()

    # the bird stays in the same x position, so bird.x is a constant
    # center bird on screen
    bird = Bird(50, int(WIN_HEIGHT/2 - Bird.HEIGHT/2), 2,
                (images['bird-wingup'], images['bird-wingdown']))

    pipes = deque()

    frame_clock = 0  # this counter is only incremented if the game isn't paused
    score = 0
    done = paused = False
    while not done:
        clock.tick(FPS)

        # Handle this 'manually'.  If we used pygame.time.set_timer(),
        # pipe addition would be messed up when paused.
        if not (paused or frame_clock % msec_to_frames(PipePair.ADD_INTERVAL)):
            pygame.event.post(pygame.event.Event(PipePair.ADD_EVENT))

        for e in pygame.event.get():
            if e.type == QUIT or (e.type == KEYUP and e.key == K_ESCAPE):
                done = True
                break
            elif e.type == KEYUP and e.key in (K_PAUSE, K_p):
                paused = not paused
            elif e.type == MOUSEBUTTONUP or (e.type == KEYUP and
                    e.key in (K_UP, K_RETURN, K_SPACE)):
                bird.msec_to_climb = Bird.CLIMB_DURATION
            elif e.type == PipePair.ADD_EVENT:
                pp = PipePair(images['pipe-end'], images['pipe-body'])
                pipes.append(pp)

        if paused:
            continue  # don't draw anything

        # check for collisions
        pipe_collision = any(p.collides_with(bird) for p in pipes)
        if pipe_collision or 0 >= bird.y or bird.y >= WIN_HEIGHT - Bird.HEIGHT:
            done = True

        for x in (0, WIN_WIDTH / 2):
            display_surface.blit(images['background'], (x, 0))

        while pipes and not pipes[0].visible:
            pipes.popleft()

        for p in pipes:
            p.update()
            display_surface.blit(p.image, p.rect)

        bird.update()
        display_surface.blit(bird.image, bird.rect)

        # update and display score
        for p in pipes:
            if p.x + PipePair.WIDTH < bird.x and not p.score_counted:
                score += 1
                p.score_counted = True

        score_surface = score_font.render(str(score), True, (255, 255, 255))
        score_x = WIN_WIDTH/2 - score_surface.get_width()/2
        display_surface.blit(score_surface, (score_x, PipePair.PIECE_HEIGHT))

        pygame.display.flip()
        frame_clock += 1
    print('Game over! Score: %i' % score)
    pygame.quit()

if __name__ == '__main__':
    # If this module had been imported, __name__ would be 'flappybird'.
    # It was executed (e.g. by double-clicking the file), so call main.
    main()
<commit_msg>Removed pipe addition event; minor cleanup<commit_after>#! /usr/bin/env python3

"""Flappy Bird, implemented using Pygame."""

import math
import os
from random import randint
from collections import deque

import pygame
from pygame.locals import *

FPS = 60
ANIMATION_SPEED = 0.18  # pixels per millisecond
WIN_WIDTH = 284 * 2     # BG image size: 284x512 px; tiled twice
WIN_HEIGHT = 512

class Bird(pygame.sprite.Sprite):
    """Represents the bird controlled by the player.

    The bird is the 'hero' of this game.  The player can make it climb
    (ascend quickly), otherwise it sinks (descends more slowly).  It must
    pass through the space in between pipes (for every pipe passed, one
    point is scored); if it crashes into a pipe, the game ends.

    Attributes:
    x: The bird's X coordinate.
    y: The bird's Y coordinate.
    msec_to_climb: The number of milliseconds left to climb, where a
        complete climb lasts Bird.CLIMB_DURATION milliseconds.

    Constants:
    WIDTH: The width, in pixels, of the bird's image.
    HEIGHT: The height, in pixels, of the bird's image.
    SINK_SPEED: With which speed, in pixels per millisecond, the bird
        descends in one second while not climbing.
    CLIMB_SPEED: With which speed, in pixels per millisecond, the bird
        ascends in one second while climbing, on average.  See also the
        Bird.update docstring.
    CLIMB_DURATION: The number of milliseconds it takes the bird to
        execute a complete climb.
    """

    WIDTH = HEIGHT = 32
    SINK_SPEED = 0.18
    CLIMB_SPEED = 0.3
    CLIMB_DURATION = 333.3

    def __init__(self, x, y, msec_to_climb, images):
        """Initialise a new Bird instance.

        Arguments:
        x: The bird's initial X coordinate.
        y: The bird's initial Y coordinate.
        msec_to_climb: The number of milliseconds left to climb, where a
            complete climb lasts Bird.CLIMB_DURATION milliseconds.  Use
            this if you want the bird to make a (small?) climb at the
            very beginning of the game.
        images: A tuple containing the images used by this bird.  It
            must contain the following images, in the following order:
                0. image of the bird with its wing pointing upward
                1. image of the bird with its wing pointing downward
        """
        super(Bird, self).__init__()
        self.x, self.y = x, y
        self.msec_to_climb = msec_to_climb
        self._img_wingup, self._img_wingdown = images
        self._mask_wingup = pygame.mask.from_surface(self._img_wingup)
        self._mask_wingdown = pygame.mask.from_surface(self._img_wingdown)

    def update(self, delta_frames=1):
        """Update the bird's position.

        This function uses the cosine function to achieve a smooth climb:
        In the first and last few frames, the bird climbs very little, in the
        middle of the climb, it climbs a lot.
        One complete climb lasts CLIMB_DURATION milliseconds, during which
        the bird ascends with an average speed of CLIMB_SPEED px/ms.
        This Bird's msec_to_climb attribute will automatically be
        decreased accordingly if it was > 0 when this method was called.

        Arguments:
        delta_frames: The number of frames elapsed since this method was
            last called.
        """
        if self.msec_to_climb > 0:
            frac_climb_done = 1 - self.msec_to_climb/Bird.CLIMB_DURATION
            self.y -= (Bird.CLIMB_SPEED * frames_to_msec(delta_frames) *
                       (1 - math.cos(frac_climb_done * math.pi)))
            self.msec_to_climb -= frames_to_msec(delta_frames)
        else:
            self.y += Bird.SINK_SPEED * frames_to_msec(delta_frames)

    @property
    def image(self):
        """Get a Surface containing this bird's image.

        This will decide whether to return an image where the bird's
        visible wing is pointing upward or where it is pointing downward
        based on pygame.time.get_ticks().  This will animate the flapping
        bird, even though pygame doesn't support animated GIFs.
        """
        if pygame.time.get_ticks() % 500 >= 250:
            return self._img_wingup
        else:
            return self._img_wingdown

    @property
    def mask(self):
        """Get a bitmask for use in collision detection.

        The bitmask excludes all pixels in self.image with a
        transparency greater than 127."""
        if pygame.time.get_ticks() % 500 >= 250:
            return self._mask_wingup
        else:
            return self._mask_wingdown

    @property
    def rect(self):
        """Get the bird's position, width, and height, as a pygame.Rect."""
        return Rect(self.x, self.y, Bird.WIDTH, Bird.HEIGHT)

class PipePair(pygame.sprite.Sprite):
    """Represents an obstacle.

    A PipePair has a top and a bottom pipe, and only between them can
    the bird pass -- if it collides with either part, the game is over.

    Attributes:
    x: The PipePair's X position.  This is a float, to make movement
        smoother.  Note that there is no y attribute, as it will only
        ever be 0.
    image: A pygame.Surface which can be blitted to the display surface
        to display the PipePair.
    mask: A bitmask which excludes all pixels in self.image with a
        transparency greater than 127.  This can be used for collision
        detection.
    top_pieces: The number of pieces, including the end piece, in the
        top pipe.
    bottom_pieces: The number of pieces, including the end piece, in
        the bottom pipe.

    Constants:
    WIDTH: The width, in pixels, of a pipe piece.  Because a pipe is
        only one piece wide, this is also the width of a PipePair's
        image.
    PIECE_HEIGHT: The height, in pixels, of a pipe piece.
    ADD_INTERVAL: The interval, in milliseconds, in between adding new
        pipes.
    """

    WIDTH = 80
    PIECE_HEIGHT = 32
    ADD_INTERVAL = 3000

    def __init__(self, pipe_end_img, pipe_body_img):
        """Initialises a new random PipePair.

        The new PipePair will automatically be assigned an x attribute of
        float(WIN_WIDTH - 1).

        Arguments:
        pipe_end_img: The image to use to represent a pipe's end piece.
        pipe_body_img: The image to use to represent one horizontal slice
            of a pipe's body.
        """
        self.x = float(WIN_WIDTH - 1)
        self.score_counted = False

        self.image = pygame.Surface((PipePair.WIDTH, WIN_HEIGHT), SRCALPHA)
        self.image.convert()   # speeds up blitting
        self.image.fill((0, 0, 0, 0))
        total_pipe_body_pieces = int(
            (WIN_HEIGHT -                  # fill window from top to bottom
             3 * Bird.HEIGHT -             # make room for bird to fit through
             3 * PipePair.PIECE_HEIGHT) /  # 2 end pieces + 1 body piece
            PipePair.PIECE_HEIGHT          # to get number of pipe pieces
        )
        self.bottom_pieces = randint(1, total_pipe_body_pieces)
        self.top_pieces = total_pipe_body_pieces - self.bottom_pieces

        # bottom pipe
        for i in range(1, self.bottom_pieces + 1):
            piece_pos = (0, WIN_HEIGHT - i*PipePair.PIECE_HEIGHT)
            self.image.blit(pipe_body_img, piece_pos)
        bottom_pipe_end_y = WIN_HEIGHT - self.bottom_height_px
        bottom_end_piece_pos = (0, bottom_pipe_end_y - PipePair.PIECE_HEIGHT)
        self.image.blit(pipe_end_img, bottom_end_piece_pos)

        # top pipe
        for i in range(self.top_pieces):
            self.image.blit(pipe_body_img, (0, i * PipePair.PIECE_HEIGHT))
        top_pipe_end_y = self.top_height_px
        self.image.blit(pipe_end_img, (0, top_pipe_end_y))

        # compensate for added end pieces
        self.top_pieces += 1
        self.bottom_pieces += 1

        # for collision detection
        self.mask = pygame.mask.from_surface(self.image)

    @property
    def top_height_px(self):
        """Get the top pipe's height, in pixels."""
        return self.top_pieces * PipePair.PIECE_HEIGHT

    @property
    def bottom_height_px(self):
        """Get the bottom pipe's height, in pixels."""
        return self.bottom_pieces * PipePair.PIECE_HEIGHT

    @property
    def visible(self):
        """Get whether this PipePair on screen, visible to the player."""
        return -PipePair.WIDTH < self.x < WIN_WIDTH

    @property
    def rect(self):
        """Get the Rect which contains this PipePair."""
        return Rect(self.x, 0, PipePair.WIDTH, PipePair.PIECE_HEIGHT)

    def update(self, delta_frames=1):
        """Update the PipePair's position.

        Arguments:
        delta_frames: The number of frames elapsed since this method was
            last called.
        """
        self.x -= ANIMATION_SPEED * frames_to_msec(delta_frames)

    def collides_with(self, bird):
        """Get whether the bird collides with a pipe in this PipePair.

        Arguments:
        bird: The Bird which should be tested for collision with this
            PipePair.
        """
        return pygame.sprite.collide_mask(self, bird)

def load_images():
    """Load all images required by the game and return a dict of them.

    The returned dict has the following keys:
    background: The game's background image.
    bird-wingup: An image of the bird with its wing pointing upward.
        Use this and bird-wingdown to create a flapping bird.
    bird-wingdown: An image of the bird with its wing pointing downward.
        Use this and bird-wingup to create a flapping bird.
    pipe-end: An image of a pipe's end piece (the slightly wider bit).
        Use this and pipe-body to make pipes.
    pipe-body: An image of a slice of a pipe's body.  Use this and
        pipe-body to make pipes.
    """

    def load_image(img_file_name):
        """Return the loaded pygame image with the specified file name.

        This function looks for images in the game's images folder
        (./images/).  All images are converted before being returned to
        speed up blitting.

        Arguments:
        img_file_name: The file name (including its extension, e.g.
            '.png') of the required image, without a file path.
        """
        file_name = os.path.join('.', 'images', img_file_name)
        img = pygame.image.load(file_name)
        img.convert()
        return img

    return {'background': load_image('background.png'),
            'pipe-end': load_image('pipe_end.png'),
            'pipe-body': load_image('pipe_body.png'),
            # images for animating the flapping bird -- animated GIFs are
            # not supported in pygame
            'bird-wingup': load_image('bird_wing_up.png'),
            'bird-wingdown': load_image('bird_wing_down.png')}

def frames_to_msec(frames, fps=FPS):
    """Convert frames to milliseconds at the specified framerate.

    Arguments:
    frames: How many frames to convert to milliseconds.
    fps: The framerate to use for conversion.  Default: FPS.
    """
    return 1000.0 * frames / fps

def msec_to_frames(milliseconds, fps=FPS):
    """Convert milliseconds to frames at the specified framerate.

    Arguments:
    milliseconds: How many milliseconds to convert to frames.
    fps: The framerate to use for conversion.  Default: FPS.
    """
    return fps * milliseconds / 1000.0

def main():
    """The application's entry point.

    If someone executes this module (instead of importing it, for
    example), this function is called.
    """

    pygame.init()

    display_surface = pygame.display.set_mode((WIN_WIDTH, WIN_HEIGHT))
    pygame.display.set_caption('Pygame Flappy Bird')

    clock = pygame.time.Clock()
    score_font = pygame.font.SysFont(None, 32, bold=True)  # default font
    images = load_images()

    # the bird stays in the same x position, so bird.x is a constant
    # center bird on screen
    bird = Bird(50, int(WIN_HEIGHT/2 - Bird.HEIGHT/2), 2,
                (images['bird-wingup'], images['bird-wingdown']))

    pipes = deque()

    frame_clock = 0  # this counter is only incremented if the game isn't paused
    score = 0
    done = paused = False
    while not done:
        clock.tick(FPS)

        # Handle this 'manually'.  If we used pygame.time.set_timer(),
        # pipe addition would be messed up when paused.
        if not (paused or frame_clock % msec_to_frames(PipePair.ADD_INTERVAL)):
            pp = PipePair(images['pipe-end'], images['pipe-body'])
            pipes.append(pp)

        for e in pygame.event.get():
            if e.type == QUIT or (e.type == KEYUP and e.key == K_ESCAPE):
                done = True
                break
            elif e.type == KEYUP and e.key in (K_PAUSE, K_p):
                paused = not paused
            elif e.type == MOUSEBUTTONUP or (e.type == KEYUP and
                    e.key in (K_UP, K_RETURN, K_SPACE)):
                bird.msec_to_climb = Bird.CLIMB_DURATION

        if paused:
            continue  # don't draw anything

        # check for collisions
        pipe_collision = any(p.collides_with(bird) for p in pipes)
        if pipe_collision or 0 >= bird.y or bird.y >= WIN_HEIGHT - Bird.HEIGHT:
            done = True

        for x in (0, WIN_WIDTH / 2):
            display_surface.blit(images['background'], (x, 0))

        while pipes and not pipes[0].visible:
            pipes.popleft()

        for p in pipes:
            p.update()
            display_surface.blit(p.image, p.rect)

        bird.update()
        display_surface.blit(bird.image, bird.rect)

        # update and display score
        for p in pipes:
            if p.x + PipePair.WIDTH < bird.x and not p.score_counted:
                score += 1
                p.score_counted = True

        score_surface = score_font.render(str(score), True, (255, 255, 255))
        score_x = WIN_WIDTH/2 - score_surface.get_width()/2
        display_surface.blit(score_surface, (score_x, PipePair.PIECE_HEIGHT))

        pygame.display.flip()
        frame_clock += 1
    print('Game over! Score: %i' % score)
    pygame.quit()

if __name__ == '__main__':
    # If this module had been imported, __name__ would be 'flappybird'.
    # It was executed (e.g. by double-clicking the file), so call main.
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import argparse
import re

import pandas as pd

from isochrones.cluster import StarClusterModel, StarCatalog
from isochrones import get_ichrone
from isochrones.priors import FehPrior

try:
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
except ImportError:
    rank = 0

parser = argparse.ArgumentParser()
parser.add_argument('starfile', type=str, help='File with star cluster data.')
parser.add_argument('--bands', nargs='*', help='bands to use (vals, uncs must be in table)')
parser.add_argument('--props', nargs='*', help='properties to use (valus, uncs must be in table)')
parser.add_argument('--models', type=str, default='mist')
parser.add_argument('--max_distance', type=float, default=1000)
parser.add_argument('--mineep', type=int, default=202)
parser.add_argument('--maxeep', type=int, default=800)
parser.add_argument('--maxAV', type=float, default=0.1)
parser.add_argument('--overwrite', '-o', action='store_true')
parser.add_argument('--nlive', type=int, default=1000)

args = parser.parse_args()

if rank==0:
    ic = get_ichrone(args.models)

    if re.search('.h5', args.starfile) or re.search('.hdf', args.starfile):
        stars = pd.read_hdf(args.starfile)
    else:
        stars = pd.read_csv(args.starfile)

    cat = StarCatalog(stars)

    model = StarClusterModel(ic, cat, eep_bounds=(args.mineep, args.maxeep),
                             max_distance=args.max_distance, max_AV=args.maxAV, name=args.name)
    model.set_prior('feh', FehPrior(halo_fraction=0.5))

else:
    model = None

model = comm.bcast(model, root=0)
model.fit(overwrite=args.overwrite, n_live_points=args.nlive, init_MPI=False)
<commit_msg>include props<commit_after>#!/usr/bin/env python
import argparse
import re

import pandas as pd

from isochrones.cluster import StarClusterModel, StarCatalog
from isochrones import get_ichrone
from isochrones.priors import FehPrior

try:
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
except ImportError:
    rank = 0

parser = argparse.ArgumentParser()
parser.add_argument('starfile', type=str, help='File with star cluster data.')
parser.add_argument('--bands', nargs='*', help='bands to use (vals, uncs must be in table)')
parser.add_argument('--props', nargs='*', help='properties to use (valus, uncs must be in table)')
parser.add_argument('--models', type=str, default='mist')
parser.add_argument('--max_distance', type=float, default=1000)
parser.add_argument('--mineep', type=int, default=202)
parser.add_argument('--maxeep', type=int, default=800)
parser.add_argument('--maxAV', type=float, default=0.1)
parser.add_argument('--overwrite', '-o', action='store_true')
parser.add_argument('--nlive', type=int, default=1000)

args = parser.parse_args()

if rank==0:
    ic = get_ichrone(args.models)

    if re.search('.h5', args.starfile) or re.search('.hdf', args.starfile):
        stars = pd.read_hdf(args.starfile)
    else:
        stars = pd.read_csv(args.starfile)

    cat = StarCatalog(stars, props=args.props)

    model = StarClusterModel(ic, cat, eep_bounds=(args.mineep, args.maxeep),
                             max_distance=args.max_distance, max_AV=args.maxAV, name=args.name)
    model.set_prior('feh', FehPrior(halo_fraction=0.5))

else:
    model = None

model = comm.bcast(model, root=0)
model.fit(overwrite=args.overwrite, n_live_points=args.nlive, init_MPI=False)
<|endoftext|>