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

from django.db import models
from django.contrib import admin
from django.contrib.admin import SimpleListFilter
from django_mptt_admin.admin import DjangoMpttAdmin
from django.contrib.contenttypes import generic
from django.forms import TextInput, Textarea

from django.core import urlresolvers
from django.utils.html import format_html

from promrep.forms import PostInlineForm

from models import Person, Office, Praenomen, PostAssertion, \
    Group, RoleType, DateType, SecondarySource, Gens, \
    PostAssertionNote, Tribe, Location

admin.site.register(DateType)
admin.site.register(RoleType)

class PostAssertionNoteInline(admin.StackedInline):
    model = PostAssertion.notes.through
    extra = 0

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-open',)

    verbose_name = 'Post Person Note'
    verbose_name_plural = 'Post Person Notes'

    raw_id_fields = ('postassertionnote', )

    related_lookup_fields = {
        'fk': ['postassertionnote', ],
    }

class PostAssertionAdmin(admin.ModelAdmin):
    list_display = ('id', 'person', 'office',
                    'date_start', 'date_end', 'secondary_source', 'review_flag',
                    'created_by', 'created', 'modified')

    list_filter = ('role', 'office', 'secondary_source', )

    readonly_fields = ('id', )

    fieldsets = [
            ('Database Info',
                {'fields': [('id', 'review_flag'), ]}),
            ('', {'fields':
                    [
                        'person',
                        'office',
                        'secondary_source',
                        'location',
                        ('role', 'uncertain'),
                        ('original_text', 'office_xref'),
                    ],
                }
                ),
            ('Dates', {'fields': [
                        ('date_display_text'),
                        ('date_source_text', 'date_secondary_source'),
                        ('date_start', 'date_start_uncertain'),
                        ('date_end', 'date_end_uncertain')
                     ]})
            ]

    raw_id_fields = ('group', 'person', 'office', 'location')

    related_lookup_fields = {
         'fk': ['group', 'person', 'office', 'location', ],
    }

    inlines = (PostAssertionNoteInline, )

admin.site.register(PostAssertion, PostAssertionAdmin)

class NoteAdmin(admin.ModelAdmin):
    list_display = ('id', 'note_type', 'secondary_source', 'text', 'extra_info', 'created', 'modified')
    list_display_links = ('note_type', 'secondary_source', 'text', 'extra_info')
    readonly_fields = ('id', 'created', 'modified')

    search_fields = ['id', 'note_type', 'text']
    fields = ('id', ['secondary_source', 'note_type'], 'text', 'extra_info', )

    show_change_link = True

admin.site.register(PostAssertionNote, NoteAdmin)

class LocationAdmin(admin.ModelAdmin):
    list_display = ('id', 'name', 'location_type', 'created', 'modified')
    list_display_links = ('id', 'name', 'location_type', )
    readonly_fields = ('id', 'created', 'modified')

    search_fields = ['id', 'name', ]
    fields = ('id', ['name', 'location_type'], 'description' )

    show_change_link = True

admin.site.register(Location, LocationAdmin)

class PersonInline(admin.StackedInline):
    model = PostAssertion
    form = PostInlineForm

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-closed',)

    verbose_name_plural = 'Persons on this Group'
    verbose_name = 'Person:'

    show_change_link = True

    fields = (
        ['id', 'review_flag', 'position'] ,
        ['person',],
        ['office', 'role',],
        ['uncertain'],
        ['secondary_source', 'original_text', 'office_xref'],
        ['date_display_text', ],
        ['date_source_text', 'date_secondary_source', ],
        ['date_start', 'date_start_uncertain', 'date_end', 'date_end_uncertain'],
        'notes',
        'edit_link'
    )

    sortable_field_name = 'position'

    readonly_fields = ('id', )

    raw_id_fields = ('person', 'notes')
    related_lookup_fields = {
        'fk': ['person', ],
        'm2m': ['notes', ]
    }

    extra = 0

class PostAssertionInline(admin.StackedInline):
    model = PostAssertion
    form = PostInlineForm

    extra = 0

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-closed',)

    verbose_name = 'Post Assertion'
    verbose_name_plural = 'Person Post Assertions'

    show_change_link = True

    ordering = ('-date_start', '-date_end', )

    readonly_fields = ('id', )

    fields = (['id', 'review_flag', ] ,
            ['office', 'role'],
            ['uncertain', ],
            ['secondary_source', ],
            ['group',],
            ['original_text', 'office_xref'],
            'date_display_text',
            ['date_source_text', 'date_secondary_source', ],
            ['date_start', 'date_start_uncertain', 'date_end', 'date_end_uncertain'],
            'notes',
            'edit_link',
            )

    raw_id_fields = ('notes', 'group', )

    related_lookup_fields = {
        'fk': ['group', ],
        'm2m': ['notes', ],
    }

class REUpdatedListFilter(SimpleListFilter):
    title = 'RE Updated'
    parameter_name = 're_updated'

    def lookups(self, request, model_admin):
        return (
            ('yes', 'Yes'),
            ('no', 'No'),
        )

    def queryset(self, request, queryset):
        if self.value() == 'yes':
            return queryset.exclude(re_number_old__exact='')

        if self.value() == 'no':
            return queryset.filter(re_number_old__exact='')

class PersonAdmin(admin.ModelAdmin):

    fieldsets = [
            ('Database',
                 {'fields': [('id', 'review_flag')]},),
            ('General Info',
                {
                'classes': ('grp-collapse grp-open',),
                'fields': [
                    ('sex',),
                    ('praenomen', 'praenomen_uncertain'),
                    ('nomen', 'nomen_uncertain'),
                    ('filiation', 'filiation_uncertain'),
                    ('cognomen', 'cognomen_uncertain'),
                    ('other_names', ),
                    ('gens', 'gens_uncertain',),
                    ('tribe', 'tribe_uncertain'),
                    ('origin', ),
                ]}),
            ('RE',
                {'classes': ('grp-collapse grp-open',),
                'fields': [
                    ('re_number', 're_number_old', ),
                ]}
            ),
        ('Dates', {
            'classes': ('grp-collapse grp-open',),
            'fields': [
                ('date_display_text'),
                ('date_source_text', 'date_secondary_source'),
                ('date_first', 'date_first_type'),
                ('date_last', 'date_last_type'),
                ('era_from', 'era_to'),
        ]}),
        ('Patrician', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('patrician', 'patrician_uncertain'),
                              ('patrician_notes')]}),
        ('Novus', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('novus', 'novus_uncertain'),
                              ('novus_notes')]}),
        ('Nobilis', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('nobilis', 'nobilis_uncertain'),
                                ('novus_notes')]}),
        ('Eques', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('eques', 'eques_uncertain'),
                              ('eques_notes')]}),
        ]

    readonly_fields = ('id', )

    list_display = (
        'id',
        'url_to_edit_person',
        'review_flag',
        'updated_by',
        'modified',
        'created_by',
        'created',
        )

    search_fields = ['id', 'nomen', 'cognomen', 'praenomen__abbrev',
                    'praenomen__name', 'other_names', 're_number', ]

    list_filter = ('postassertion__role', 'nomen', 'postassertion__office',
                   'review_flag', REUpdatedListFilter, 'patrician', 'novus',
                   'nobilis', 'eques', )

    inlines = (PostAssertionInline, )
    exclude = ('assertions', )

admin.site.register(Person, PersonAdmin)

class PostInline(admin.TabularInline):
    model = Group
    classes = ('grp-collapse grp-open',)

    readonly_fields = ('id', 'get_persons', 'related_label')
    fields = ('id', 'related_label', 'get_persons')
    extra = 0

class OfficeAdmin(DjangoMpttAdmin):
    readonly_fields = ('id', )
    mptt_indent_field = "name"

    list_display = (
        'id',
        'name',
        'description',
        )

admin.site.register(Office, OfficeAdmin)

class GroupAdmin(admin.ModelAdmin):

    search_fields = ['id', 'postassertion__person__nomen', 'postassertion__person__cognomen', 'postassertion__office']

    list_display = (
        'id',
        'date_year',
        'date_info',
        'notes',
        'modified',
        'created',
    )

    readonly_fields = ('id', )
    list_display_links = ('id', 'notes', )

    fieldsets = [ ('Database Info', {'fields': ['id']}),
                    ( '' ,
                        {
                        'fields': [
                                ('date_year', 'date_info', ),
                                ( 'notes', ),
                                ],
                        }
                    ),
            ]

    inlines = [PersonInline, ]
    exclude = ('persons',  )

admin.site.register(Group, GroupAdmin)

class GensAdmin(admin.ModelAdmin):

    list_display = ('id', 'name', )
    readonly_fields = ('id', )
    list_display_links = ('id', 'name', )

admin.site.register(Gens, GensAdmin)

class TribeAdmin(admin.ModelAdmin):

    list_display = ('id', 'abbrev', 'name', )
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev', 'name', )

admin.site.register(Tribe, TribeAdmin)

class PraenomenAdmin(admin.ModelAdmin):

    list_display = ('id', 'abbrev', 'name')
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev', 'name')

admin.site.register(Praenomen, PraenomenAdmin)

class SecondarySourceAdmin(admin.ModelAdmin):

    list_display = ('id', 'name', 'abbrev_name', 'biblio')
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev_name', 'name', 'biblio')

admin.site.register(SecondarySource, SecondarySourceAdmin)

<commit_msg>displaying group information in the PostAssertionAdmin page;<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

from django.db import models
from django.contrib import admin
from django.contrib.admin import SimpleListFilter
from django_mptt_admin.admin import DjangoMpttAdmin
from django.contrib.contenttypes import generic
from django.forms import TextInput, Textarea

from django.core import urlresolvers
from django.utils.html import format_html

from promrep.forms import PostInlineForm

from models import Person, Office, Praenomen, PostAssertion, \
    Group, RoleType, DateType, SecondarySource, Gens, \
    PostAssertionNote, Tribe, Location

admin.site.register(DateType)
admin.site.register(RoleType)

class PostAssertionNoteInline(admin.StackedInline):
    model = PostAssertion.notes.through
    extra = 0

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-open',)

    verbose_name = 'Post Person Note'
    verbose_name_plural = 'Post Person Notes'

    raw_id_fields = ('postassertionnote', )

    related_lookup_fields = {
        'fk': ['postassertionnote', ],
    }

class PostAssertionAdmin(admin.ModelAdmin):
    list_display = ('id', 'person', 'office',
                    'date_start', 'date_end', 'secondary_source',
                    'review_flag',
                    'created_by', 'created', 'modified')

    list_filter = ('role', 'office', 'secondary_source', )

    readonly_fields = ('id', )

    fieldsets = [
            ('Database Info',
                {'fields': [('id', 'review_flag'), ]}),
            ('', {'fields':
                    [
                        'person',
                        'office',
                        'location',
                        'secondary_source',
                        ('role', 'uncertain'),
                        'group',
                        ('original_text', 'office_xref'),
                    ],
                }
                ),
            ('Dates', {'fields': [
                        ('date_display_text'),
                        ('date_source_text', 'date_secondary_source'),
                        ('date_start', 'date_start_uncertain'),
                        ('date_end', 'date_end_uncertain')
                     ]})
            ]

    raw_id_fields = ('group', 'person', 'office', 'location')

    related_lookup_fields = {
         'fk': ['group', 'person', 'office', 'location', ],
    }

    inlines = (PostAssertionNoteInline, )

admin.site.register(PostAssertion, PostAssertionAdmin)

class NoteAdmin(admin.ModelAdmin):
    list_display = ('id', 'note_type', 'secondary_source', 'text', 'extra_info', 'created', 'modified')
    list_display_links = ('note_type', 'secondary_source', 'text', 'extra_info')
    readonly_fields = ('id', 'created', 'modified')

    search_fields = ['id', 'note_type', 'text']
    fields = ('id', ['secondary_source', 'note_type'], 'text', 'extra_info', )

    show_change_link = True

admin.site.register(PostAssertionNote, NoteAdmin)

class LocationAdmin(admin.ModelAdmin):
    list_display = ('id', 'name', 'location_type', 'created', 'modified')
    list_display_links = ('id', 'name', 'location_type', )
    readonly_fields = ('id', 'created', 'modified')

    search_fields = ['id', 'name', ]
    fields = ('id', ['name', 'location_type'], 'description' )

    show_change_link = True

admin.site.register(Location, LocationAdmin)

class PersonInline(admin.StackedInline):
    model = PostAssertion
    form = PostInlineForm

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-closed',)

    verbose_name_plural = 'Persons on this Group'
    verbose_name = 'Person:'

    show_change_link = True

    fields = (
        ['id', 'review_flag', 'position'] ,
        ['person',],
        ['office', 'role',],
        ['uncertain'],
        ['secondary_source', 'original_text', 'office_xref'],
        ['date_display_text', ],
        ['date_source_text', 'date_secondary_source', ],
        ['date_start', 'date_start_uncertain', 'date_end', 'date_end_uncertain'],
        'notes',
        'edit_link'
    )

    sortable_field_name = 'position'

    readonly_fields = ('id', )

    raw_id_fields = ('person', 'notes')
    related_lookup_fields = {
        'fk': ['person', ],
        'm2m': ['notes', ]
    }

    extra = 0

class PostAssertionInline(admin.StackedInline):
    model = PostAssertion
    form = PostInlineForm

    extra = 0

    classes = ('grp-collapse grp-open',)
    inline_classes = ('grp-collapse grp-closed',)

    verbose_name = 'Post Assertion'
    verbose_name_plural = 'Person Post Assertions'

    show_change_link = True

    ordering = ('-date_start', '-date_end', )

    readonly_fields = ('id', )

    fields = (['id', 'review_flag', ] ,
            ['office', 'role'],
            ['uncertain', ],
            ['secondary_source', ],
            ['group',],
            ['original_text', 'office_xref'],
            'date_display_text',
            ['date_source_text', 'date_secondary_source', ],
            ['date_start', 'date_start_uncertain', 'date_end', 'date_end_uncertain'],
            'notes',
            'edit_link',
            )

    raw_id_fields = ('notes', 'group', )

    related_lookup_fields = {
        'fk': ['group', ],
        'm2m': ['notes', ],
    }

class REUpdatedListFilter(SimpleListFilter):
    title = 'RE Updated'
    parameter_name = 're_updated'

    def lookups(self, request, model_admin):
        return (
            ('yes', 'Yes'),
            ('no', 'No'),
        )

    def queryset(self, request, queryset):
        if self.value() == 'yes':
            return queryset.exclude(re_number_old__exact='')

        if self.value() == 'no':
            return queryset.filter(re_number_old__exact='')

class PersonAdmin(admin.ModelAdmin):

    fieldsets = [
            ('Database',
                 {'fields': [('id', 'review_flag')]},),
            ('General Info',
                {
                'classes': ('grp-collapse grp-open',),
                'fields': [
                    ('sex',),
                    ('praenomen', 'praenomen_uncertain'),
                    ('nomen', 'nomen_uncertain'),
                    ('filiation', 'filiation_uncertain'),
                    ('cognomen', 'cognomen_uncertain'),
                    ('other_names', ),
                    ('gens', 'gens_uncertain',),
                    ('tribe', 'tribe_uncertain'),
                    ('origin', ),
                ]}),
            ('RE',
                {'classes': ('grp-collapse grp-open',),
                'fields': [
                    ('re_number', 're_number_old', ),
                ]}
            ),
        ('Dates', {
            'classes': ('grp-collapse grp-open',),
            'fields': [
                ('date_display_text'),
                ('date_source_text', 'date_secondary_source'),
                ('date_first', 'date_first_type'),
                ('date_last', 'date_last_type'),
                ('era_from', 'era_to'),
        ]}),
        ('Patrician', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('patrician', 'patrician_uncertain'),
                              ('patrician_notes')]}),
        ('Novus', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('novus', 'novus_uncertain'),
                              ('novus_notes')]}),
        ('Nobilis', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('nobilis', 'nobilis_uncertain'),
                                ('novus_notes')]}),
        ('Eques', {
            'classes': ('grp-collapse grp-open',),
            'fields': [('eques', 'eques_uncertain'),
                              ('eques_notes')]}),
        ]

    readonly_fields = ('id', )

    list_display = (
        'id',
        'url_to_edit_person',
        'review_flag',
        'updated_by',
        'modified',
        'created_by',
        'created',
        )

    search_fields = ['id', 'nomen', 'cognomen', 'praenomen__abbrev',
                    'praenomen__name', 'other_names', 're_number', ]

    list_filter = ('postassertion__role', 'nomen', 'postassertion__office',
                   'review_flag', REUpdatedListFilter, 'patrician', 'novus',
                   'nobilis', 'eques', )

    inlines = (PostAssertionInline, )
    exclude = ('assertions', )

admin.site.register(Person, PersonAdmin)

class PostInline(admin.TabularInline):
    model = Group
    classes = ('grp-collapse grp-open',)

    readonly_fields = ('id', 'get_persons', 'related_label')
    fields = ('id', 'related_label', 'get_persons')
    extra = 0

class OfficeAdmin(DjangoMpttAdmin):
    readonly_fields = ('id', )
    mptt_indent_field = "name"

    list_display = (
        'id',
        'name',
        'description',
        )

admin.site.register(Office, OfficeAdmin)

class GroupAdmin(admin.ModelAdmin):

    search_fields = ['id', 'postassertion__person__nomen', 'postassertion__person__cognomen', 'postassertion__office']

    list_display = (
        'id',
        'date_year',
        'date_info',
        'notes',
        'modified',
        'created',
    )

    readonly_fields = ('id', )
    list_display_links = ('id', 'notes', )

    fieldsets = [ ('Database Info', {'fields': ['id']}),
                    ( '' ,
                        {
                        'fields': [
                                ('date_year', 'date_info', ),
                                ( 'notes', ),
                                ],
                        }
                    ),
            ]

    inlines = [PersonInline, ]
    exclude = ('persons',  )

admin.site.register(Group, GroupAdmin)

class GensAdmin(admin.ModelAdmin):

    list_display = ('id', 'name', )
    readonly_fields = ('id', )
    list_display_links = ('id', 'name', )

admin.site.register(Gens, GensAdmin)

class TribeAdmin(admin.ModelAdmin):

    list_display = ('id', 'abbrev', 'name', )
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev', 'name', )

admin.site.register(Tribe, TribeAdmin)

class PraenomenAdmin(admin.ModelAdmin):

    list_display = ('id', 'abbrev', 'name')
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev', 'name')

admin.site.register(Praenomen, PraenomenAdmin)

class SecondarySourceAdmin(admin.ModelAdmin):

    list_display = ('id', 'name', 'abbrev_name', 'biblio')
    readonly_fields = ('id', )
    list_display_links = ('id', 'abbrev_name', 'name', 'biblio')

admin.site.register(SecondarySource, SecondarySourceAdmin)

<|endoftext|>


<commit_before>from collections import OrderedDict

from django.conf import settings as s
from django.core.urlresolvers import reverse
from django.http import JsonResponse
from django.views.generic.detail import DetailView
from haystack.generic_views import FacetedSearchView, SearchView
from promrep.forms import PromrepFacetedSearchForm, SenateSearchForm
from promrep.models import (Office, Person, PostAssertion, Province,
                            RelationshipAssertion, StatusAssertion)
from promrep.solr_backends.solr_backend_field_collapsing import \
    GroupedSearchQuerySet

class PromrepFacetedSearchView(FacetedSearchView):
    facet_fields = ['gender', 'nobilis', 'novus', 'tribe',
                    'patrician', 'province', 'offices', 'life_date_types',
                    'eques']

    autocomplete_facets = PromrepFacetedSearchForm.AUTOCOMPLETE_FACETS
    form_class = PromrepFacetedSearchForm
    load_all = True
    print_paginate = 200

    # Person allows us to search people with no assertions
    queryset = GroupedSearchQuerySet().models(
        PostAssertion,
        StatusAssertion,
        Person,
        RelationshipAssertion).group_by('person_id')

    def get_paginate_by(self, queryset):
        """
        Get the number of items to paginate by, or ``None`` for no pagination.
        """
        if self.request.GET.get('printme'):
            return self.print_paginate
        else:
            return self.paginate_by

    def get_queryset(self):
        queryset = super(PromrepFacetedSearchView, self).get_queryset()
        all_facets = self.autocomplete_facets + self.facet_fields

        for facet in all_facets:
            # only return results with a mincount of 1
            queryset = queryset.facet(
                facet, sort='index', limit=-1, mincount=1)

        selected_facets = self.request.GET.getlist('selected_facets')
        for facet in selected_facets:
            if 'offices' in facet:
                queryset = queryset.order_by(
                    'date_start').order_by('date_end')
                break

        return queryset.order_by('era_order')

    def get_context_data(self, **kwargs):  # noqa
        context = super(
            PromrepFacetedSearchView, self).get_context_data(**kwargs)

        if self.request.GET.getlist('selected_facets'):
            context['selected_facets'] = self.request.GET.getlist(
                'selected_facets')

        qs = self.request.GET.copy()
        context['querydict'] = qs.copy()
        self.paginate_by = 200

        if self.request.GET.get('q'):
            qs.pop('q')

            # always remove the page number
            if 'page' in qs:
                qs.pop('page')

            if len(qs):
                url = '?{0}'.format(qs.urlencode())
            else:
                url = reverse('haystack_search')

            context['remove_text_filter'] = url

        context['era_filter'] = self._get_date_url_and_filter(
            'era_from', 'era_to')

        context['date_filter'] = self._get_date_url_and_filter(
            'date_from', 'date_to')

        # used to generate the lists for the autocomplete dictionary
        context['autocomplete_facets'] = self.autocomplete_facets

        for afacet in context['autocomplete_facets']:

            if self.request.GET.get(afacet):
                qs = self.request.GET.copy()
                qs.pop(afacet)

                url = reverse('haystack_search')

                if len(qs):
                    url = '?{0}'.format(qs.urlencode())

                context[afacet] = (url, self.request.GET.get(afacet))

        office_lookups = s.LOOKUPS['offices']

        # hierarchical facets data
        try:
            magisterial = Office.objects.get(id=office_lookups['magisterial'])
            if magisterial:
                context[
                    'magisterial_office_list'
                ] = magisterial.get_descendants()
        except:
            pass

        try:
            promagistracies = Office.objects.get(
                id=office_lookups['promagistracies'])
            if promagistracies:
                context[
                    'promagistracies_office_list'
                ] = promagistracies.get_descendants()
        except:
            pass

        try:
            priesthoods = Office.objects.get(id=office_lookups['priesthoods'])
            if priesthoods:
                context[
                    'priesthoods_office_list'
                ] = priesthoods.get_descendants()
        except:
            pass

        try:
            non_magisterial = Office.objects.get(
                id=office_lookups['non_magisterial'])
            if non_magisterial:
                context[
                    'non_magisterial_office_list'
                ] = non_magisterial.get_descendants()
        except:
            pass

        try:
            distinctions = Office.objects.get(
                id=office_lookups['distinctions'])
            if distinctions:
                context[
                    'distinctions_office_list'
                ] = distinctions.get_descendants()
        except:
            pass

        if 'offices' not in context['facets']['fields']:
            context.update({'facets': self.get_queryset().facet_counts()})

        context['office_fdict'] = dict(context['facets']['fields']['offices'])

        context['province_list'] = Province.objects.all()
        context['province_fdict'] = dict(
            context['facets']['fields']['province'])

        return context

    def _get_date_url_and_filter(self, field_from, field_to):
        if (field_from or field_to) in self.request.GET:
            field_text = None

            qs = self.request.GET.copy()

            if qs.get(field_from) and qs.get(field_to):
                field_text = '{} to {}'.format(
                    qs.pop(field_from)[0], qs.pop(field_to)[0])
            elif qs.get(field_to):
                field_text = 'Before {}'.format(qs.pop(field_to)[0])
            elif qs.get(field_from):
                field_text = 'After {}'.format(qs.pop(field_from)[0])

            # always remove the page number
            if 'page' in qs:
                qs.pop('page')

            url = reverse('haystack_search')
            if len(qs):
                url = '?{}'.format(qs.urlencode())

            if field_text:
                return (url, field_text)

        return None

class PersonDetailView(DetailView):
    model = Person
    template_name = 'promrep/persons/detail.html'

    def get_context_data(self, **kwargs):  # noqa
        context = super(
            PersonDetailView, self).get_context_data(**kwargs)

        relationships = OrderedDict()

        relationships_qs = RelationshipAssertion.objects.filter(
            person=self.get_object()
        ).order_by('relationship__order', 'relationship_number')

        for r in relationships_qs:
            if r.relationship not in relationships:
                relationships[r.relationship] = []

            relationships[r.relationship].append(r)

        context['relationships'] = relationships

        return context

def get_relationships_network(request, pk):
    network = {}

    try:
        person = Person.objects.get(pk=pk)
        network = _get_relationships_network(person)
    except Person.DoesNotExist:
        network = {'error': 'Person with id {} does not exit'.format(pk)}

    return JsonResponse(network)

def _get_relationships_network(person):
    nodes = []
    edges = []

    relationships = RelationshipAssertion.objects.filter(person=person)
    for relationship in relationships:
        for person in [relationship.person, relationship.related_person]:
            node = {
                'id': person.id,
                'label': person.__unicode__()
            }

            if node not in nodes:
                nodes.append(node)

        edge = {
            'id': relationship.id,
            'label': relationship.relationship.__unicode__(),
            'source': relationship.person.id,
            'target': relationship.related_person.id,
        }

        if edge not in edges:
            edges.append(edge)

    return {'nodes': nodes, 'edges': edges}

class SenateSearchView(SearchView):
    load_all = True
    form_class = SenateSearchForm
    queryset = GroupedSearchQuerySet().models(
        StatusAssertion).narrow('senator:true').group_by('person_id')
    template_name = 'search/senate.html'

    def get_queryset(self):
        queryset = super(SenateSearchView, self).get_queryset()

        if 'senate_date' in self.request.GET:
            queryset = GroupedSearchQuerySet().models(
                StatusAssertion).narrow('senator:true').group_by('person_id')
        else:
            queryset = queryset.narrow('date:[{0} TO {0}]'.format(
                SenateSearchForm.INITIAL_DATE))

        return queryset.order_by(
            'date_start_uncertain').order_by(
                'date_start').order_by('date_end')

    def get_context_data(self, **kwargs):  # noqa
        context = super(SenateSearchView, self).get_context_data(**kwargs)
        qs = self.request.GET.copy()
        senate_date = SenateSearchForm.INITIAL_DATE_DISPLAY

        context['querydict'] = qs.copy()

        if 'senate_date' in qs:
            senate_date = qs.pop('senate_date')[0]

        context['senate_date'] = senate_date

        return context
<commit_msg>Fixed ordering switch in get_queryset of facetedsearchview<commit_after>from collections import OrderedDict

from django.conf import settings as s
from django.core.urlresolvers import reverse
from django.http import JsonResponse
from django.views.generic.detail import DetailView
from haystack.generic_views import FacetedSearchView, SearchView
from promrep.forms import PromrepFacetedSearchForm, SenateSearchForm
from promrep.models import (Office, Person, PostAssertion, Province,
                            RelationshipAssertion, StatusAssertion)
from promrep.solr_backends.solr_backend_field_collapsing import \
    GroupedSearchQuerySet

class PromrepFacetedSearchView(FacetedSearchView):
    facet_fields = ['gender', 'nobilis', 'novus', 'tribe',
                    'patrician', 'province', 'offices', 'life_date_types',
                    'eques']

    autocomplete_facets = PromrepFacetedSearchForm.AUTOCOMPLETE_FACETS
    form_class = PromrepFacetedSearchForm
    load_all = True
    print_paginate = 200

    # Person allows us to search people with no assertions
    queryset = GroupedSearchQuerySet().models(
        PostAssertion,
        StatusAssertion,
        Person,
        RelationshipAssertion).group_by('person_id')

    def get_paginate_by(self, queryset):
        """
        Get the number of items to paginate by, or ``None`` for no pagination.
        """
        if self.request.GET.get('printme'):
            return self.print_paginate
        else:
            return self.paginate_by

    def get_queryset(self):
        queryset = super(PromrepFacetedSearchView, self).get_queryset()
        all_facets = self.autocomplete_facets + self.facet_fields

        for facet in all_facets:
            # only return results with a mincount of 1
            queryset = queryset.facet(
                facet, sort='index', limit=-1, mincount=1)

        selected_facets = self.request.GET.getlist('selected_facets')
        offices = False
        for facet in selected_facets:
            if 'offices' in facet:
                queryset = queryset.order_by(
                    'date_start', 'date_end')
                offices = True
                break
        if offices:
            return queryset
        return queryset.order_by('era_order')

    def get_context_data(self, **kwargs):  # noqa
        context = super(
            PromrepFacetedSearchView, self).get_context_data(**kwargs)

        if self.request.GET.getlist('selected_facets'):
            context['selected_facets'] = self.request.GET.getlist(
                'selected_facets')

        qs = self.request.GET.copy()
        context['querydict'] = qs.copy()
        self.paginate_by = 200

        if self.request.GET.get('q'):
            qs.pop('q')

            # always remove the page number
            if 'page' in qs:
                qs.pop('page')

            if len(qs):
                url = '?{0}'.format(qs.urlencode())
            else:
                url = reverse('haystack_search')

            context['remove_text_filter'] = url

        context['era_filter'] = self._get_date_url_and_filter(
            'era_from', 'era_to')

        context['date_filter'] = self._get_date_url_and_filter(
            'date_from', 'date_to')

        # used to generate the lists for the autocomplete dictionary
        context['autocomplete_facets'] = self.autocomplete_facets

        for afacet in context['autocomplete_facets']:

            if self.request.GET.get(afacet):
                qs = self.request.GET.copy()
                qs.pop(afacet)

                url = reverse('haystack_search')

                if len(qs):
                    url = '?{0}'.format(qs.urlencode())

                context[afacet] = (url, self.request.GET.get(afacet))

        office_lookups = s.LOOKUPS['offices']

        # hierarchical facets data
        try:
            magisterial = Office.objects.get(id=office_lookups['magisterial'])
            if magisterial:
                context[
                    'magisterial_office_list'
                ] = magisterial.get_descendants()
        except:
            pass

        try:
            promagistracies = Office.objects.get(
                id=office_lookups['promagistracies'])
            if promagistracies:
                context[
                    'promagistracies_office_list'
                ] = promagistracies.get_descendants()
        except:
            pass

        try:
            priesthoods = Office.objects.get(id=office_lookups['priesthoods'])
            if priesthoods:
                context[
                    'priesthoods_office_list'
                ] = priesthoods.get_descendants()
        except:
            pass

        try:
            non_magisterial = Office.objects.get(
                id=office_lookups['non_magisterial'])
            if non_magisterial:
                context[
                    'non_magisterial_office_list'
                ] = non_magisterial.get_descendants()
        except:
            pass

        try:
            distinctions = Office.objects.get(
                id=office_lookups['distinctions'])
            if distinctions:
                context[
                    'distinctions_office_list'
                ] = distinctions.get_descendants()
        except:
            pass

        if 'offices' not in context['facets']['fields']:
            context.update({'facets': self.get_queryset().facet_counts()})

        context['office_fdict'] = dict(context['facets']['fields']['offices'])

        context['province_list'] = Province.objects.all()
        context['province_fdict'] = dict(
            context['facets']['fields']['province'])

        return context

    def _get_date_url_and_filter(self, field_from, field_to):
        if (field_from or field_to) in self.request.GET:
            field_text = None

            qs = self.request.GET.copy()

            if qs.get(field_from) and qs.get(field_to):
                field_text = '{} to {}'.format(
                    qs.pop(field_from)[0], qs.pop(field_to)[0])
            elif qs.get(field_to):
                field_text = 'Before {}'.format(qs.pop(field_to)[0])
            elif qs.get(field_from):
                field_text = 'After {}'.format(qs.pop(field_from)[0])

            # always remove the page number
            if 'page' in qs:
                qs.pop('page')

            url = reverse('haystack_search')
            if len(qs):
                url = '?{}'.format(qs.urlencode())

            if field_text:
                return (url, field_text)

        return None

class PersonDetailView(DetailView):
    model = Person
    template_name = 'promrep/persons/detail.html'

    def get_context_data(self, **kwargs):  # noqa
        context = super(
            PersonDetailView, self).get_context_data(**kwargs)

        relationships = OrderedDict()

        relationships_qs = RelationshipAssertion.objects.filter(
            person=self.get_object()
        ).order_by('relationship__order', 'relationship_number')

        for r in relationships_qs:
            if r.relationship not in relationships:
                relationships[r.relationship] = []

            relationships[r.relationship].append(r)

        context['relationships'] = relationships

        return context

def get_relationships_network(request, pk):
    network = {}

    try:
        person = Person.objects.get(pk=pk)
        network = _get_relationships_network(person)
    except Person.DoesNotExist:
        network = {'error': 'Person with id {} does not exit'.format(pk)}

    return JsonResponse(network)

def _get_relationships_network(person):
    nodes = []
    edges = []

    relationships = RelationshipAssertion.objects.filter(person=person)
    for relationship in relationships:
        for person in [relationship.person, relationship.related_person]:
            node = {
                'id': person.id,
                'label': person.__unicode__()
            }

            if node not in nodes:
                nodes.append(node)

        edge = {
            'id': relationship.id,
            'label': relationship.relationship.__unicode__(),
            'source': relationship.person.id,
            'target': relationship.related_person.id,
        }

        if edge not in edges:
            edges.append(edge)

    return {'nodes': nodes, 'edges': edges}

class SenateSearchView(SearchView):
    load_all = True
    form_class = SenateSearchForm
    queryset = GroupedSearchQuerySet().models(
        StatusAssertion).narrow('senator:true').group_by('person_id')
    template_name = 'search/senate.html'

    def get_queryset(self):
        queryset = super(SenateSearchView, self).get_queryset()

        if 'senate_date' in self.request.GET:
            queryset = GroupedSearchQuerySet().models(
                StatusAssertion).narrow('senator:true').group_by('person_id')
        else:
            queryset = queryset.narrow('date:[{0} TO {0}]'.format(
                SenateSearchForm.INITIAL_DATE))

        return queryset.order_by(
            'date_start_uncertain').order_by(
                'date_start').order_by('date_end')

    def get_context_data(self, **kwargs):  # noqa
        context = super(SenateSearchView, self).get_context_data(**kwargs)
        qs = self.request.GET.copy()
        senate_date = SenateSearchForm.INITIAL_DATE_DISPLAY

        context['querydict'] = qs.copy()

        if 'senate_date' in qs:
            senate_date = qs.pop('senate_date')[0]

        context['senate_date'] = senate_date

        return context
<|endoftext|>


<commit_before>VERSION = '3.2.0.dev'
<commit_msg>release a new version 3.2.0<commit_after>VERSION = '3.2.0'
<|endoftext|>


<commit_before>"""`python-memcached`-compatible wrapper around `_pylibmc`.

The interface is pretty much exactly the same as python-memcached, with some
minor differences. If you should happen to spot any, file a bug!

>>> import pylibmc
>>> mc = pylibmc.Client(["127.0.0.1"])
>>> b = mc.behaviors
>>> ks = list(sorted(k for k in b.keys() if not k.startswith("_")))
>>> ks  # doctest: +NORMALIZE_WHITESPACE
['buffer_requests', 'cache_lookups', 'connect_timeout', 'distribution', 'hash',
 'ketama', 'ketama_hash', 'ketama_weighted', 'no block', 'poll_timeout',
 'rcv_timeout', 'retry_timeout', 'server_failure_limit', 'snd_timeout',
 'socket recv size', 'socket send size', 'sort_hosts', 'support_cas',
 'tcp_nodelay', 'verify_key']
>>> b["hash"]
'default'
>>> b["hash"] = 'fnv1a_32'
>>> mc.behaviors["hash"]
'fnv1a_32'
>>> super(pylibmc.Client, mc).get_behaviors()["hash"]
6
"""

import _pylibmc

__all__ = ["hashers", "distributions", "Client"]
__version__ = _pylibmc.__version__

errors = tuple(e for (n, e) in _pylibmc.exceptions)
# *Cough* Uhm, not the prettiest of things but this unpacks all exception
# objects and sets them on the very module object currently constructed.
import sys
modself = sys.modules[__name__]
for name, exc in _pylibmc.exceptions:
    setattr(modself, name, exc)

hashers, hashers_rvs = {}, {}
distributions, distributions_rvs = {}, {}
# Not the prettiest way of doing things, but works well.
for name in dir(_pylibmc):
    if name.startswith("hash_"):
        key, value = name[5:], getattr(_pylibmc, name)
        hashers[key] = value
        hashers_rvs[value] = key
    elif name.startswith("distribution_"):
        key, value = name[13:].replace("_", " "), getattr(_pylibmc, name)
        distributions[key] = value
        distributions_rvs[value] = key

class BehaviorDict(dict):
    def __init__(self, client, *args, **kwds):
        super(BehaviorDict, self).__init__(*args, **kwds)
        self.client = client

    def __setitem__(self, name, value):
        super(BehaviorDict, self).__setitem__(name, value)
        self.client.set_behaviors({name: value})

    def update(self, *args, **kwds):
        super(BehaviorDict, self).update(*args, **kwds)
        self.client.set_behaviors(self.copy())

class Client(_pylibmc.client):
    def __init__(self, servers, binary=False):
        """Initialize a memcached client instance.

        This connects to the servers in *servers*, which will default to being
        TCP servers. If it looks like a filesystem path, a UNIX socket. If
        prefixed with `udp:`, a UDP connection.

        If *binary* is True, the binary memcached protocol is used.
        """
        addr_tups = []
        for server in servers:
            addr = server
            port = 11211
            if server.startswith("udp:"):
                stype = _pylibmc.server_type_udp
                addr = addr[4:]
                if ":" in server:
                    (addr, port) = addr.split(":", 1)
                    port = int(port)
            elif ":" in server:
                stype = _pylibmc.server_type_tcp
                (addr, port) = server.split(":", 1)
                port = int(port)
            elif "/" in server:
                stype = _pylibmc.server_type_unix
                port = 0
            else:
                stype = _pylibmc.server_type_tcp
            addr_tups.append((stype, addr, port))
        super(Client, self).__init__(servers=addr_tups, binary=binary)

    def get_behaviors(self):
        """Gets the behaviors from the underlying C client instance.

        Reverses the integer constants for `hash` and `distribution` into more
        understandable string values. See *set_behaviors* for info.
        """
        behaviors = super(Client, self).get_behaviors()
        behaviors["hash"] = hashers_rvs[behaviors["hash"]]
        behaviors["distribution"] = distributions_rvs[behaviors["distribution"]]
        return BehaviorDict(self, behaviors)

    def set_behaviors(self, behaviors):
        """Sets the behaviors on the underlying C client instance.

        Takes care of morphing the `hash` key, if specified, into the
        corresponding integer constant (which the C client expects.) If,
        however, an unknown value is specified, it's passed on to the C client
        (where it most surely will error out.)

        This also happens for `distribution`.
        """
        behaviors = behaviors.copy()
        if behaviors.get("hash", None) in hashers:
            behaviors["hash"] = hashers[behaviors["hash"]]
        if behaviors.get("distribution") in distributions:
            behaviors["distribution"] = distributions[behaviors["distribution"]]
        return super(Client, self).set_behaviors(behaviors)

    behaviors = property(get_behaviors, set_behaviors)

from contextlib import contextmanager

class ClientPool(list):
    """Client pooling helper.

    This is mostly useful in threaded environments, because a client isn't
    thread-safe at all. Instead, what you want to do is have each thread use
    its own client, but you don't want to reconnect these all the time.

    The solution is a pool, and this class is a helper for that.

    >>> mc = Client(["127.0.0.1"])
    >>> pool = ClientPool()
    >>> pool.fill(mc, 4)
    >>> with pool.reserve() as mc:
    ...     mc.set("hi", "ho")
    ...     mc.delete("hi")
    ...
    True
    True
    """

    @contextmanager
    def reserve(self):
        """Reserve a client, and put it back into the pool when done."""
        mc = self.pop()
        try:
            yield mc
        finally:
            self.append(mc)

    def fill(self, mc, n_slots):
        """Fill *n_slots* of the pool with clones of *mc*."""
        for i in xrange(n_slots):
            self.append(mc.clone())

class ThreadMappedPool(dict):
    """Much like the *ClientPool*, helps you with pooling.

    In a threaded environment, you'd most likely want to have a client per
    thread. And there'd be no harm in one thread keeping the same client at all
    times. So, why not map threads to clients? That's what this class does.

    If a client is reserved, this class checks for a key based on the current
    thread, and if none exists, clones the master client and inserts that key.

    >>> mc = Client(["127.0.0.1"])
    >>> pool = ThreadMappedPool(mc)
    >>> with pool.reserve() as mc:
    ...     mc.set("hi", "ho")
    ...     mc.delete("hi")
    ...
    True
    True
    """

    def __new__(cls, master):
        return super(ThreadMappedPool, cls).__new__(cls)

    def __init__(self, master):
        self.master = master

    @contextmanager
    def reserve(self):
        """Reserve a client.

        Creates a new client based on the master client if none exists for the
        current thread.
        """
        key = threading.current_thread().ident
        mc = self.pop(key, None)
        if mc is None:
            mc = self.master.clone()
        try:
            yield mc
        finally:
            self[key] = mc

# This makes sure ThreadMappedPool doesn't exist with non-thread Pythons.
try:
    import threading
except ImportError:
    del ThreadMappedPool

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Use `thread.get_ident` instead of more convoluted way.<commit_after>"""`python-memcached`-compatible wrapper around `_pylibmc`.

The interface is pretty much exactly the same as python-memcached, with some
minor differences. If you should happen to spot any, file a bug!

>>> import pylibmc
>>> mc = pylibmc.Client(["127.0.0.1"])
>>> b = mc.behaviors
>>> ks = list(sorted(k for k in b.keys() if not k.startswith("_")))
>>> ks  # doctest: +NORMALIZE_WHITESPACE
['buffer_requests', 'cache_lookups', 'connect_timeout', 'distribution', 'hash',
 'ketama', 'ketama_hash', 'ketama_weighted', 'no block', 'poll_timeout',
 'rcv_timeout', 'retry_timeout', 'server_failure_limit', 'snd_timeout',
 'socket recv size', 'socket send size', 'sort_hosts', 'support_cas',
 'tcp_nodelay', 'verify_key']
>>> b["hash"]
'default'
>>> b["hash"] = 'fnv1a_32'
>>> mc.behaviors["hash"]
'fnv1a_32'
>>> super(pylibmc.Client, mc).get_behaviors()["hash"]
6
"""

import _pylibmc

__all__ = ["hashers", "distributions", "Client"]
__version__ = _pylibmc.__version__

errors = tuple(e for (n, e) in _pylibmc.exceptions)
# *Cough* Uhm, not the prettiest of things but this unpacks all exception
# objects and sets them on the very module object currently constructed.
import sys
modself = sys.modules[__name__]
for name, exc in _pylibmc.exceptions:
    setattr(modself, name, exc)

hashers, hashers_rvs = {}, {}
distributions, distributions_rvs = {}, {}
# Not the prettiest way of doing things, but works well.
for name in dir(_pylibmc):
    if name.startswith("hash_"):
        key, value = name[5:], getattr(_pylibmc, name)
        hashers[key] = value
        hashers_rvs[value] = key
    elif name.startswith("distribution_"):
        key, value = name[13:].replace("_", " "), getattr(_pylibmc, name)
        distributions[key] = value
        distributions_rvs[value] = key

class BehaviorDict(dict):
    def __init__(self, client, *args, **kwds):
        super(BehaviorDict, self).__init__(*args, **kwds)
        self.client = client

    def __setitem__(self, name, value):
        super(BehaviorDict, self).__setitem__(name, value)
        self.client.set_behaviors({name: value})

    def update(self, *args, **kwds):
        super(BehaviorDict, self).update(*args, **kwds)
        self.client.set_behaviors(self.copy())

class Client(_pylibmc.client):
    def __init__(self, servers, binary=False):
        """Initialize a memcached client instance.

        This connects to the servers in *servers*, which will default to being
        TCP servers. If it looks like a filesystem path, a UNIX socket. If
        prefixed with `udp:`, a UDP connection.

        If *binary* is True, the binary memcached protocol is used.
        """
        addr_tups = []
        for server in servers:
            addr = server
            port = 11211
            if server.startswith("udp:"):
                stype = _pylibmc.server_type_udp
                addr = addr[4:]
                if ":" in server:
                    (addr, port) = addr.split(":", 1)
                    port = int(port)
            elif ":" in server:
                stype = _pylibmc.server_type_tcp
                (addr, port) = server.split(":", 1)
                port = int(port)
            elif "/" in server:
                stype = _pylibmc.server_type_unix
                port = 0
            else:
                stype = _pylibmc.server_type_tcp
            addr_tups.append((stype, addr, port))
        super(Client, self).__init__(servers=addr_tups, binary=binary)

    def get_behaviors(self):
        """Gets the behaviors from the underlying C client instance.

        Reverses the integer constants for `hash` and `distribution` into more
        understandable string values. See *set_behaviors* for info.
        """
        behaviors = super(Client, self).get_behaviors()
        behaviors["hash"] = hashers_rvs[behaviors["hash"]]
        behaviors["distribution"] = distributions_rvs[behaviors["distribution"]]
        return BehaviorDict(self, behaviors)

    def set_behaviors(self, behaviors):
        """Sets the behaviors on the underlying C client instance.

        Takes care of morphing the `hash` key, if specified, into the
        corresponding integer constant (which the C client expects.) If,
        however, an unknown value is specified, it's passed on to the C client
        (where it most surely will error out.)

        This also happens for `distribution`.
        """
        behaviors = behaviors.copy()
        if behaviors.get("hash", None) in hashers:
            behaviors["hash"] = hashers[behaviors["hash"]]
        if behaviors.get("distribution") in distributions:
            behaviors["distribution"] = distributions[behaviors["distribution"]]
        return super(Client, self).set_behaviors(behaviors)

    behaviors = property(get_behaviors, set_behaviors)

from contextlib import contextmanager

class ClientPool(list):
    """Client pooling helper.

    This is mostly useful in threaded environments, because a client isn't
    thread-safe at all. Instead, what you want to do is have each thread use
    its own client, but you don't want to reconnect these all the time.

    The solution is a pool, and this class is a helper for that.

    >>> mc = Client(["127.0.0.1"])
    >>> pool = ClientPool()
    >>> pool.fill(mc, 4)
    >>> with pool.reserve() as mc:
    ...     mc.set("hi", "ho")
    ...     mc.delete("hi")
    ...
    True
    True
    """

    @contextmanager
    def reserve(self):
        """Reserve a client, and put it back into the pool when done."""
        mc = self.pop()
        try:
            yield mc
        finally:
            self.append(mc)

    def fill(self, mc, n_slots):
        """Fill *n_slots* of the pool with clones of *mc*."""
        for i in xrange(n_slots):
            self.append(mc.clone())

class ThreadMappedPool(dict):
    """Much like the *ClientPool*, helps you with pooling.

    In a threaded environment, you'd most likely want to have a client per
    thread. And there'd be no harm in one thread keeping the same client at all
    times. So, why not map threads to clients? That's what this class does.

    If a client is reserved, this class checks for a key based on the current
    thread, and if none exists, clones the master client and inserts that key.

    >>> mc = Client(["127.0.0.1"])
    >>> pool = ThreadMappedPool(mc)
    >>> with pool.reserve() as mc:
    ...     mc.set("hi", "ho")
    ...     mc.delete("hi")
    ...
    True
    True
    """

    def __new__(cls, master):
        return super(ThreadMappedPool, cls).__new__(cls)

    def __init__(self, master):
        self.master = master

    @contextmanager
    def reserve(self):
        """Reserve a client.

        Creates a new client based on the master client if none exists for the
        current thread.
        """
        key = thread.get_ident()
        mc = self.pop(key, None)
        if mc is None:
            mc = self.master.clone()
        try:
            yield mc
        finally:
            self[key] = mc

# This makes sure ThreadMappedPool doesn't exist with non-thread Pythons.
try:
    import thread
except ImportError:
    del ThreadMappedPool

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>"""
# Pyqtree

Pyqtree is a pure Python spatial index for GIS or rendering usage.
It stores and quickly retrieves items from a 2x2 rectangular grid area,
and grows in depth and detail as more items are added.
The actual quad tree implementation is adapted from
[Matt Rasmussen's compbio library](https://github.com/mdrasmus/compbio/blob/master/rasmus/quadtree.py)
and extended for geospatial use.

## Platforms

Python 2 and 3.

## Dependencies

Pyqtree is written in pure Python and has no dependencies.

## Installing It

Installing Pyqtree can be done by opening your terminal or commandline and typing:

    pip install pyqtree

Alternatively, you can simply download the "pyqtree.py" file and place
it anywhere Python can import it, such as the Python site-packages folder.

## Example Usage

Start your script by importing the quad tree.

    from pyqtree import Index

Setup the spatial index, giving it a bounding box area to keep track of.
The bounding box being in a four-tuple: (xmin, ymin, xmax, ymax).

    spindex = Index(bbox=(0, 0, 100, 100))

Populate the index with items that you want to be retrieved at a later point,
along with each item's geographic bbox.

    # this example assumes you have a list of items with bbox attribute
    for item in items:
        spindex.insert(item, item.bbox)

Then when you have a region of interest and you wish to retrieve items from that region,
just use the index's intersect method. This quickly gives you a list of the stored items
whose bboxes intersects your region of interests.

    overlapbbox = (51, 51, 86, 86)
    matches = spindex.intersect(overlapbbox)

There are other things that can be done as well, but that's it for the main usage!

## More Information:

- [Home Page](http://github.com/karimbahgat/Pyqtree)
- [API Documentation](http://pythonhosted.org/Pyqtree)

## License:

This code is free to share, use, reuse, and modify according to the MIT license, see LICENSE.txt.

## Credits:

- Karim Bahgat (2015)
- Joschua Gandert (2016)

"""

__version__ = "0.25.0"

#PYTHON VERSION CHECK
import sys
PYTHON3 = int(sys.version[0]) == 3
if PYTHON3:
    xrange = range

def _normalize_rect(rect):
    if len(rect) == 2:
        x1, y1 = rect
        x2, y2 = rect
    else:
        x1, y1, x2, y2 = rect
    if x1 > x2:
        x1, x2 = x2, x1
    if y1 > y2:
        y1, y2 = y2, y1
    return (x1, y1, x2, y2)

def _loopallchildren(parent):
    for child in parent.children:
        if child.children:
            for subchild in _loopallchildren(child):
                yield subchild
        yield child

class _QuadNode(object):
    def __init__(self, item, rect):
        self.item = item
        self.rect = rect

    def __eq__(self, other):
        return self.item == other.item and self.rect == other.rect

    def __hash__(self):
        return hash(self.item)

class _QuadTree(object):
    """
    Internal backend version of the index.
    The index being used behind the scenes. Has all the same methods as the user
    index, but requires more technical arguments when initiating it than the
    user-friendly version.
    """

    def __init__(self, x, y, width, height, max_items, max_depth, _depth=0):
        self.nodes = []
        self.children = []
        self.center = (x, y)
        self.width, self.height = width, height
        self.max_items = max_items
        self.max_depth = max_depth
        self._depth = _depth

    def __iter__(self):
        for child in _loopallchildren(self):
            yield child

    def _insert(self, item, bbox):
        rect = _normalize_rect(bbox)
        if len(self.children) == 0:
            node = _QuadNode(item, rect)
            self.nodes.append(node)

            if len(self.nodes) > self.max_items and self._depth < self.max_depth:
                self._split()
        else:
            self._insert_into_children(item, rect)

    def _remove(self, item, bbox):
        rect = _normalize_rect(bbox)
        if len(self.children) == 0:
            node = _QuadNode(item, rect)
            self.nodes.remove(node)
        else:
            self._remove_from_children(item, rect)

    def _intersect(self, rect, results=None):
        if results is None:
            rect = _normalize_rect(rect)
            results = []
        # search children
        if self.children:
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._intersect(rect, results)
                if rect[3] >= self.center[1]:
                    self.children[1]._intersect(rect, results)
            if rect[2] >= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._intersect(rect, results)
                if rect[3] >= self.center[1]:
                    self.children[3]._intersect(rect, results)
        # search node at this level
        for node in self.nodes:
            if (node.rect[2] >= rect[0] and node.rect[0] <= rect[2] and
                node.rect[3] >= rect[1] and node.rect[1] <= rect[3]):
                results.append(node.item)
        return results

    def _insert_into_children(self, item, rect):
        # if rect spans center then insert here
        if (rect[0] <= self.center[0] and rect[2] >= self.center[0] and
            rect[1] <= self.center[1] and rect[3] >= self.center[1]):
            node = _QuadNode(item, rect)
            self.nodes.append(node)
        else:
            # try to insert into children
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._insert(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[1]._insert(item, rect)
            if rect[2] > self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._insert(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[3]._insert(item, rect)

    def _remove_from_children(self, item, rect):
        # if rect spans center then insert here
        if (rect[0] <= self.center[0] and rect[2] >= self.center[0] and
            rect[1] <= self.center[1] and rect[3] >= self.center[1]):
            node = _QuadNode(item, rect)
            self.nodes.remove(node)
        else:
            # try to remove from children
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._remove(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[1]._remove(item, rect)
            if rect[2] > self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._remove(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[3]._remove(item, rect)

    def _split(self):
        quartwidth = self.width / 4.0
        quartheight = self.height / 4.0
        halfwidth = self.width / 2.0
        halfheight = self.height / 2.0
        x1 = self.center[0] - quartwidth
        x2 = self.center[0] + quartwidth
        y1 = self.center[1] - quartheight
        y2 = self.center[1] + quartheight
        new_depth = self._depth + 1
        self.children = [_QuadTree(x1, y1, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x1, y2, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x2, y1, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x2, y2, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth)]
        nodes = self.nodes
        self.nodes = []
        for node in nodes:
            self._insert_into_children(node.item, node.rect)

    def __len__(self):
        """
        Returns:

        - A count of the total number of members/items/nodes inserted
        into this quadtree and all of its child trees.
        """
        size = 0
        for child in self.children:
            size += len(child)
        size += len(self.nodes)
        return size

MAX_ITEMS = 10
MAX_DEPTH = 20

class Index(_QuadTree):
    """
    The top spatial index to be created by the user. Once created it can be
    populated with geographically placed members that can later be tested for
    intersection with a user inputted geographic bounding box. Note that the
    index can be iterated through in a for-statement, which loops through all
    all the quad instances and lets you access their properties.

    Example usage:

    >>> spindex = Index(bbox=(0, 0, 100, 100))
    >>> spindex.insert('duck', (50, 30, 53, 60))
    >>> spindex.insert('cookie', (10, 20, 15, 25))
    >>> spindex.insert('python', (40, 50, 95, 90))
    >>> results = spindex.intersect((51, 51, 86, 86))
    >>> sorted(results)
    ['duck', 'python']
    """

    def __init__(self, bbox=None, x=None, y=None, width=None, height=None, max_items=MAX_ITEMS, max_depth=MAX_DEPTH):
        """
        Initiate by specifying either 1) a bbox to keep track of, or 2) with an xy centerpoint and a width and height.

        Parameters:
        - **bbox**: The coordinate system bounding box of the area that the quadtree should
            keep track of, as a 4-length sequence (xmin,ymin,xmax,ymax)
        - **x**:
            The x center coordinate of the area that the quadtree should keep track of.
        - **y**
            The y center coordinate of the area that the quadtree should keep track of.
        - **width**:
            How far from the xcenter that the quadtree should look when keeping track.
        - **height**:
            How far from the ycenter that the quadtree should look when keeping track
        - **max_items** (optional): The maximum number of items allowed per quad before splitting
            up into four new subquads. Default is 10.
        - **max_depth** (optional): The maximum levels of nested subquads, after which no more splitting
            occurs and the bottommost quad nodes may grow indefinately. Default is 20.
        """
        if bbox is not None:
            x1, y1, x2, y2 = bbox
            width, height = abs(x2-x1), abs(y2-y1)
            midx, midy = x1+width/2.0, y1+height/2.0
            super(Index, self).__init__(midx, midy, width, height, max_items, max_depth)

        elif None not in (x, y, width, height):
            super(Index, self).__init__(x, y, width, height, max_items, max_depth)

        else:
            raise Exception("Either the bbox argument must be set, or the x, y, width, and height arguments must be set")

    def insert(self, item, bbox):
        """
        Inserts an item into the quadtree along with its bounding box.

        Parameters:
        - **item**: The item to insert into the index, which will be returned by the intersection method
        - **bbox**: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)
        """
        self._insert(item, bbox)

    def remove(self, item, bbox):
        """
        Removes an item from the quadtree.

        Parameters:
        - **item**: The item to remove from the index
        - **bbox**: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)

        Both parameters need to exactly match the parameters provided to the insert method.
        """
        self._remove(item, bbox)

    def intersect(self, bbox):
        """
        Intersects an input boundingbox rectangle with all of the items
        contained in the quadtree.

        Parameters:
        - **bbox**: A spatial bounding box tuple with four members (xmin,ymin,xmax,ymax)

        Returns:
        - A list of inserted items whose bounding boxes intersect with the input bbox.
        """
        return self._intersect(bbox)
<commit_msg>Fix duplicate intersect items<commit_after>"""
# Pyqtree

Pyqtree is a pure Python spatial index for GIS or rendering usage.
It stores and quickly retrieves items from a 2x2 rectangular grid area,
and grows in depth and detail as more items are added.
The actual quad tree implementation is adapted from
[Matt Rasmussen's compbio library](https://github.com/mdrasmus/compbio/blob/master/rasmus/quadtree.py)
and extended for geospatial use.

## Platforms

Python 2 and 3.

## Dependencies

Pyqtree is written in pure Python and has no dependencies.

## Installing It

Installing Pyqtree can be done by opening your terminal or commandline and typing:

    pip install pyqtree

Alternatively, you can simply download the "pyqtree.py" file and place
it anywhere Python can import it, such as the Python site-packages folder.

## Example Usage

Start your script by importing the quad tree.

    from pyqtree import Index

Setup the spatial index, giving it a bounding box area to keep track of.
The bounding box being in a four-tuple: (xmin, ymin, xmax, ymax).

    spindex = Index(bbox=(0, 0, 100, 100))

Populate the index with items that you want to be retrieved at a later point,
along with each item's geographic bbox.

    # this example assumes you have a list of items with bbox attribute
    for item in items:
        spindex.insert(item, item.bbox)

Then when you have a region of interest and you wish to retrieve items from that region,
just use the index's intersect method. This quickly gives you a list of the stored items
whose bboxes intersects your region of interests.

    overlapbbox = (51, 51, 86, 86)
    matches = spindex.intersect(overlapbbox)

There are other things that can be done as well, but that's it for the main usage!

## More Information:

- [Home Page](http://github.com/karimbahgat/Pyqtree)
- [API Documentation](http://pythonhosted.org/Pyqtree)

## License:

This code is free to share, use, reuse, and modify according to the MIT license, see LICENSE.txt.

## Credits:

- Karim Bahgat (2015)
- Joschua Gandert (2016)

"""

__version__ = "0.25.0"

#PYTHON VERSION CHECK
import sys
PYTHON3 = int(sys.version[0]) == 3
if PYTHON3:
    xrange = range

def _normalize_rect(rect):
    if len(rect) == 2:
        x1, y1 = rect
        x2, y2 = rect
    else:
        x1, y1, x2, y2 = rect
    if x1 > x2:
        x1, x2 = x2, x1
    if y1 > y2:
        y1, y2 = y2, y1
    return (x1, y1, x2, y2)

def _loopallchildren(parent):
    for child in parent.children:
        if child.children:
            for subchild in _loopallchildren(child):
                yield subchild
        yield child

class _QuadNode(object):
    def __init__(self, item, rect):
        self.item = item
        self.rect = rect

    def __eq__(self, other):
        return self.item == other.item and self.rect == other.rect

    def __hash__(self):
        return hash(self.item)

class _QuadTree(object):
    """
    Internal backend version of the index.
    The index being used behind the scenes. Has all the same methods as the user
    index, but requires more technical arguments when initiating it than the
    user-friendly version.
    """

    def __init__(self, x, y, width, height, max_items, max_depth, _depth=0):
        self.nodes = []
        self.children = []
        self.center = (x, y)
        self.width, self.height = width, height
        self.max_items = max_items
        self.max_depth = max_depth
        self._depth = _depth

    def __iter__(self):
        for child in _loopallchildren(self):
            yield child

    def _insert(self, item, bbox):
        rect = _normalize_rect(bbox)
        if len(self.children) == 0:
            node = _QuadNode(item, rect)
            self.nodes.append(node)

            if len(self.nodes) > self.max_items and self._depth < self.max_depth:
                self._split()
        else:
            self._insert_into_children(item, rect)

    def _remove(self, item, bbox):
        rect = _normalize_rect(bbox)
        if len(self.children) == 0:
            node = _QuadNode(item, rect)
            self.nodes.remove(node)
        else:
            self._remove_from_children(item, rect)

    def _intersect(self, rect, results=None):
        if results is None:
            rect = _normalize_rect(rect)
            results = []
        # search children
        if self.children:
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._intersect(rect, results)
                if rect[3] >= self.center[1]:
                    self.children[1]._intersect(rect, results)
            if rect[2] >= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._intersect(rect, results)
                if rect[3] >= self.center[1]:
                    self.children[3]._intersect(rect, results)
        # search node at this level
        for node in self.nodes:
            if (node.rect[2] >= rect[0] and node.rect[0] <= rect[2] and
                node.rect[3] >= rect[1] and node.rect[1] <= rect[3]):
                results.append(node.item)
        return results

    def _insert_into_children(self, item, rect):
        # if rect spans center then insert here
        if (rect[0] <= self.center[0] and rect[2] >= self.center[0] and
            rect[1] <= self.center[1] and rect[3] >= self.center[1]):
            node = _QuadNode(item, rect)
            self.nodes.append(node)
        else:
            # try to insert into children
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._insert(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[1]._insert(item, rect)
            if rect[2] > self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._insert(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[3]._insert(item, rect)

    def _remove_from_children(self, item, rect):
        # if rect spans center then insert here
        if (rect[0] <= self.center[0] and rect[2] >= self.center[0] and
            rect[1] <= self.center[1] and rect[3] >= self.center[1]):
            node = _QuadNode(item, rect)
            self.nodes.remove(node)
        else:
            # try to remove from children
            if rect[0] <= self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[0]._remove(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[1]._remove(item, rect)
            if rect[2] > self.center[0]:
                if rect[1] <= self.center[1]:
                    self.children[2]._remove(item, rect)
                if rect[3] >= self.center[1]:
                    self.children[3]._remove(item, rect)

    def _split(self):
        quartwidth = self.width / 4.0
        quartheight = self.height / 4.0
        halfwidth = self.width / 2.0
        halfheight = self.height / 2.0
        x1 = self.center[0] - quartwidth
        x2 = self.center[0] + quartwidth
        y1 = self.center[1] - quartheight
        y2 = self.center[1] + quartheight
        new_depth = self._depth + 1
        self.children = [_QuadTree(x1, y1, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x1, y2, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x2, y1, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth),
                         _QuadTree(x2, y2, halfwidth, halfheight,
                                   self.max_items, self.max_depth, new_depth)]
        nodes = self.nodes
        self.nodes = []
        for node in nodes:
            self._insert_into_children(node.item, node.rect)

    def __len__(self):
        """
        Returns:

        - A count of the total number of members/items/nodes inserted
        into this quadtree and all of its child trees.
        """
        size = 0
        for child in self.children:
            size += len(child)
        size += len(self.nodes)
        return size

MAX_ITEMS = 10
MAX_DEPTH = 20

class Index(_QuadTree):
    """
    The top spatial index to be created by the user. Once created it can be
    populated with geographically placed members that can later be tested for
    intersection with a user inputted geographic bounding box. Note that the
    index can be iterated through in a for-statement, which loops through all
    all the quad instances and lets you access their properties.

    Example usage:

    >>> spindex = Index(bbox=(0, 0, 100, 100))
    >>> spindex.insert('duck', (50, 30, 53, 60))
    >>> spindex.insert('cookie', (10, 20, 15, 25))
    >>> spindex.insert('python', (40, 50, 95, 90))
    >>> results = spindex.intersect((51, 51, 86, 86))
    >>> sorted(results)
    ['duck', 'python']
    """

    def __init__(self, bbox=None, x=None, y=None, width=None, height=None, max_items=MAX_ITEMS, max_depth=MAX_DEPTH):
        """
        Initiate by specifying either 1) a bbox to keep track of, or 2) with an xy centerpoint and a width and height.

        Parameters:
        - **bbox**: The coordinate system bounding box of the area that the quadtree should
            keep track of, as a 4-length sequence (xmin,ymin,xmax,ymax)
        - **x**:
            The x center coordinate of the area that the quadtree should keep track of.
        - **y**
            The y center coordinate of the area that the quadtree should keep track of.
        - **width**:
            How far from the xcenter that the quadtree should look when keeping track.
        - **height**:
            How far from the ycenter that the quadtree should look when keeping track
        - **max_items** (optional): The maximum number of items allowed per quad before splitting
            up into four new subquads. Default is 10.
        - **max_depth** (optional): The maximum levels of nested subquads, after which no more splitting
            occurs and the bottommost quad nodes may grow indefinately. Default is 20.
        """
        if bbox is not None:
            x1, y1, x2, y2 = bbox
            width, height = abs(x2-x1), abs(y2-y1)
            midx, midy = x1+width/2.0, y1+height/2.0
            super(Index, self).__init__(midx, midy, width, height, max_items, max_depth)

        elif None not in (x, y, width, height):
            super(Index, self).__init__(x, y, width, height, max_items, max_depth)

        else:
            raise Exception("Either the bbox argument must be set, or the x, y, width, and height arguments must be set")

    def insert(self, item, bbox):
        """
        Inserts an item into the quadtree along with its bounding box.

        Parameters:
        - **item**: The item to insert into the index, which will be returned by the intersection method
        - **bbox**: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)
        """
        self._insert(item, bbox)

    def remove(self, item, bbox):
        """
        Removes an item from the quadtree.

        Parameters:
        - **item**: The item to remove from the index
        - **bbox**: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)

        Both parameters need to exactly match the parameters provided to the insert method.
        """
        self._remove(item, bbox)

    def intersect(self, bbox):
        """
        Intersects an input boundingbox rectangle with all of the items
        contained in the quadtree.

        Parameters:
        - **bbox**: A spatial bounding box tuple with four members (xmin,ymin,xmax,ymax)

        Returns:
        - A list of inserted items whose bounding boxes intersect with the input bbox.
        """
        return list(set(self._intersect(bbox)))
<|endoftext|>


<commit_before>from pyschedule import Scenario, solvers, plotters

def run(S) :
    if solvers.mip.solve(S):
#         %matplotlib inline
        plotters.matplotlib.plot(S,task_colors=task_colors,fig_size=(10,5))
    else:
        print('no solution exists')

S = Scenario('bike_paint_shop', horizon=10)

Alice = S.Resource('Alice')
Bob = S.Resource('Bob')

green_paint, red_paint = S.Task('green_paint', length=2), S.Task('red_paint', length=2)
green_post, red_post = S.Task('green_post'), S.Task('red_post')

S += green_paint < green_post, red_paint + 1 <= red_post

# green_paint += Alice|Bob
# green_post += Alice|Bob
#
# red_paint += Alice|Bob
# red_post += Alice|Bob

S.clear_solution()
S.use_makespan_objective()
print(S)

task_colors = { green_paint   : '#A1D372',
                green_post    : '#A1D372',
                red_paint     : '#EB4845',
                red_post      : '#EB4845',
                S['MakeSpan'] : '#7EA7D8'}

# First remove the old resource to task assignments
# green_paint -= Alice|Bob
# green_post -= Alice|Bob
# red_paint -= Alice|Bob
# red_post -= Alice|Bob

# Add new shared ones
green_resource = Alice|Bob
green_paint += green_resource
green_post += green_resource

red_resource = Alice|Bob
red_paint += red_resource
red_post += red_resource

Paint_Shop = S.Resource('Paint_Shop')
red_paint += Paint_Shop
green_paint += Paint_Shop

Lunch = S.Task('Lunch')
Lunch += {Alice, Bob}
S += Lunch > 3, Lunch < 5
task_colors[Lunch] = '#7EA7D8'
S += red_paint > 2

#Alice is a morning bird
S += Alice['length'][:3] >= 3

run(S)
<commit_msg>print whole<commit_after>from pyschedule import Scenario, solvers, plotters

def run(S) :
    if solvers.mip.solve(S):
#         %matplotlib inline
        plotters.matplotlib.plot(S,task_colors=task_colors,fig_size=(10,5))
    else:
        print('no solution exists')

S = Scenario('bike_paint_shop', horizon=10)

Alice = S.Resource('Alice')
Bob = S.Resource('Bob')

green_paint, red_paint = S.Task('green_paint', length=2), S.Task('red_paint', length=2)
green_post, red_post = S.Task('green_post'), S.Task('red_post')

S += green_paint < green_post, red_paint + 1 <= red_post

# green_paint += Alice|Bob
# green_post += Alice|Bob
#
# red_paint += Alice|Bob
# red_post += Alice|Bob

S.clear_solution()
S.use_makespan_objective()

task_colors = { green_paint   : '#A1D372',
                green_post    : '#A1D372',
                red_paint     : '#EB4845',
                red_post      : '#EB4845',
                S['MakeSpan'] : '#7EA7D8'}

# First remove the old resource to task assignments
# green_paint -= Alice|Bob
# green_post -= Alice|Bob
# red_paint -= Alice|Bob
# red_post -= Alice|Bob

# Add new shared ones
green_resource = Alice|Bob
green_paint += green_resource
green_post += green_resource

red_resource = Alice|Bob
red_paint += red_resource
red_post += red_resource

Paint_Shop = S.Resource('Paint_Shop')
red_paint += Paint_Shop
green_paint += Paint_Shop

Lunch = S.Task('Lunch')
Lunch += {Alice, Bob}
S += Lunch > 3, Lunch < 5
task_colors[Lunch] = '#7EA7D8'
S += red_paint > 2

#Alice is a morning bird
S += Alice['length'][:3] >= 3

print(S)
run(S)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import (absolute_import, print_function, unicode_literals)
from acli.output.account import output_account_info
from acli.connections import get_boto3_session
from botocore.exceptions import NoCredentialsError
from contextlib import contextmanager

@contextmanager
def cred_checked(iam_client):
    try:
        assert iam_client.get_user()
        yield iam_client
    except NoCredentialsError:
        exit('No credentials found.')
    except Exception as e:
        exit('Unhanded exception: {0}'.format(e))

def account_info(aws_config):
    """
    @type aws_config: Config
    """
    session = get_boto3_session(aws_config)
    iam_client = session.client('iam')
    with cred_checked(iam_client) as checked_iam_client:
        account_id = checked_iam_client.get_user().get('User').get('Arn').split(':')[4]
        aliases = checked_iam_client.list_account_aliases().get('AccountAliases')
        output_account_info(output_media='console',
                            account_id=account_id,
                            account_aliases=aliases)
<commit_msg>handle situation without user id.<commit_after># -*- coding: utf-8 -*-
from __future__ import (absolute_import, print_function, unicode_literals)
from acli.output.account import output_account_info
from acli.connections import get_boto3_session
from botocore.exceptions import NoCredentialsError
from contextlib import contextmanager

@contextmanager
def cred_checked(iam_client):
    try:
        assert iam_client.list_users()
        yield iam_client
    except NoCredentialsError:
        exit('No credentials found.')
    except Exception as e:
        exit('Unhanded exception: {0}'.format(e))

def account_info(aws_config):
    """
    @type aws_config: Config
    """
    session = get_boto3_session(aws_config)
    iam_client = session.client('iam')
    with cred_checked(iam_client) as checked_iam_client:
        users = checked_iam_client.list_users(MaxItems=1)
        if users.get('Users', None):
            account_id = users.get('Users')[0]['Arn'].split(':')[4]
            aliases = checked_iam_client.list_account_aliases().get('AccountAliases')
            output_account_info(output_media='console',
                                account_id=account_id,
                                account_aliases=aliases)
<|endoftext|>