from django.contrib import messages from django.contrib.admin import helpers from django.contrib.admin.utils import model_ngettext from django.core.exceptions import PermissionDenied from django.template.response import TemplateResponse from documents.classifier import DocumentClassifier from documents.models import Tag, Correspondent, DocumentType def select_action(modeladmin, request, queryset, title, action, modelclass, success_message="", document_action=None, queryset_action=None): opts = modeladmin.model._meta app_label = opts.app_label if not modeladmin.has_change_permission(request): raise PermissionDenied if request.POST.get('post'): n = queryset.count() selected_object = modelclass.objects.get(id=request.POST.get('obj_id')) if n: for document in queryset: if document_action: document_action(document, selected_object) document_display = str(document) modeladmin.log_change(request, document, document_display) if queryset_action: queryset_action(queryset, selected_object) modeladmin.message_user(request, success_message % { "selected_object": selected_object.name, "count": n, "items": model_ngettext(modeladmin.opts, n) }, messages.SUCCESS) # Return None to display the change list page again. return None context = dict( modeladmin.admin_site.each_context(request), title=title, queryset=queryset, opts=opts, action_checkbox_name=helpers.ACTION_CHECKBOX_NAME, media=modeladmin.media, action=action, objects=modelclass.objects.all(), itemname=model_ngettext(modelclass, 1) ) request.current_app = modeladmin.admin_site.name return TemplateResponse(request, "admin/%s/%s/select_object.html" % (app_label, opts.model_name), context) def simple_action(modeladmin, request, queryset, success_message="", document_action=None, queryset_action=None): if not modeladmin.has_change_permission(request): raise PermissionDenied n = queryset.count() if n: for document in queryset: if document_action: document_action(document) document_display = str(document) modeladmin.log_change(request, document, document_display) if queryset_action: queryset_action(queryset) modeladmin.message_user(request, success_message % { "count": n, "items": model_ngettext(modeladmin.opts, n) }, messages.SUCCESS) # Return None to display the change list page again. return None def add_tag_to_selected(modeladmin, request, queryset): return select_action(modeladmin=modeladmin, request=request, queryset=queryset, title="Add tag to multiple documents", action="add_tag_to_selected", modelclass=Tag, success_message="Successfully added tag %(selected_object)s to %(count)d %(items)s.", document_action=lambda doc, tag: doc.tags.add(tag)) add_tag_to_selected.short_description = "Add tag to selected documents" def remove_tag_from_selected(modeladmin, request, queryset): return select_action(modeladmin=modeladmin, request=request, queryset=queryset, title="Remove tag from multiple documents", action="remove_tag_from_selected", modelclass=Tag, success_message="Successfully removed tag %(selected_object)s from %(count)d %(items)s.", document_action=lambda doc, tag: doc.tags.remove(tag)) remove_tag_from_selected.short_description = "Remove tag from selected documents" def set_correspondent_on_selected(modeladmin, request, queryset): return select_action(modeladmin=modeladmin, request=request, queryset=queryset, title="Set correspondent on multiple documents", action="set_correspondent_on_selected", modelclass=Correspondent, success_message="Successfully set correspondent %(selected_object)s on %(count)d %(items)s.", queryset_action=lambda qs, correspondent: qs.update(correspondent=correspondent)) set_correspondent_on_selected.short_description = "Set correspondent on selected documents" def remove_correspondent_from_selected(modeladmin, request, queryset): return simple_action(modeladmin=modeladmin, request=request, queryset=queryset, success_message="Successfully removed correspondent from %(count)d %(items)s.", queryset_action=lambda qs: qs.update(correspondent=None)) remove_correspondent_from_selected.short_description = "Remove correspondent from selected documents" def set_document_type_on_selected(modeladmin, request, queryset): return select_action(modeladmin=modeladmin, request=request, queryset=queryset, title="Set document type on multiple documents", action="set_document_type_on_selected", modelclass=DocumentType, success_message="Successfully set document type %(selected_object)s on %(count)d %(items)s.", queryset_action=lambda qs, document_type: qs.update(document_type=document_type)) set_document_type_on_selected.short_description = "Set document type on selected documents" def remove_document_type_from_selected(modeladmin, request, queryset): return simple_action(modeladmin=modeladmin, request=request, queryset=queryset, success_message="Successfully removed document type from %(count)d %(items)s.", queryset_action=lambda qs: qs.update(document_type=None)) remove_document_type_from_selected.short_description = "Remove document type from selected documents" def run_document_classifier_on_selected(modeladmin, request, queryset): try: clf = DocumentClassifier.load_classifier() return simple_action(modeladmin=modeladmin, request=request, queryset=queryset, success_message="Successfully applied document classifier to %(count)d %(items)s.", document_action=lambda doc: clf.classify_document(doc, classify_correspondent=True, classify_tags=True, classify_document_type=True)) except FileNotFoundError: modeladmin.message_user(request, "Classifier model file not found.", messages.ERROR) return None run_document_classifier_on_selected.short_description = "Run document classifier on selected"