Carrinho e exportação com seleção de campos

sigi/apps/casas/admin.py

@@ -6,7 +6,6 @@ from django.http import HttpResponseRedirect
 from django.urls import reverse
 from django.utils.safestring import mark_safe
 from django.utils.translation import gettext as _
-from import_export import resources
 from import_export.fields import Field
 from sigi.apps.casas.forms import OrgaoForm
 from sigi.apps.casas.models import Orgao, Presidente, Funcionario, TipoOrgao
@@ -19,11 +18,11 @@ from sigi.apps.contatos.models import Telefone
 from sigi.apps.convenios.models import Convenio, Projeto
 # from sigi.apps.ocorrencias.models import Ocorrencia
 # from sigi.apps.servicos.models import Servico, TipoServico
-from sigi.apps.utils import queryset_ascii
-from sigi.apps.utils.mixins import CartExportMixin
+from sigi.apps.utils import field_label, queryset_ascii
+from sigi.apps.utils.mixins import CartExportMixin, LabeledResourse
 
 
-class OrgaoExportResourse(resources.ModelResource):
+class OrgaoExportResourse(LabeledResourse):
     presidente = Field(column_name='presidente')
     telefone = Field(column_name='telefone')
     num_parlamentares = Field(column_name='num_parlamentares')
@@ -34,7 +33,7 @@ class OrgaoExportResourse(resources.ModelResource):
         fields = ('municipio__codigo_ibge', 'cnpj', 'municipio__codigo_tse',
                   'nome', 'municipio__nome', 'municipio__uf__sigla',
                   'presidente', 'logradouro', 'bairro', 'cep', 'telefone',
-                  'pagina_web', 'email', 'telefone', 'num_parlamentares',
+                  'pagina_web', 'email', 'num_parlamentares',
                   'ult_alt_endereco', 'contato')
         export_order = fields
 
@@ -305,6 +304,10 @@ class OrgaoAdmin(CartExportMixin, admin.ModelAdmin):
                      'cep', 'municipio__nome', 'municipio__uf__nome',
                      'municipio__codigo_ibge', 'pagina_web', 'observacoes')
 
+    def get_queryset(self, request):
+        queryset = super(OrgaoAdmin, self).get_queryset(request)
+        return queryset.prefetch_related('gerentes_interlegis', 'convenio_set')
+
     def get_uf(self, obj):
         return obj.municipio.uf.nome
     get_uf.short_description = _('Unidade da Federação')

sigi/apps/utils/__init__.py

@@ -1,8 +1,8 @@
-# -*- coding: utf-8 -*-
 from unicodedata import normalize
-
 from django.contrib import admin
+from django.core.exceptions import FieldDoesNotExist
 from django.db import models
+from django.utils.encoding import force_str
 
 
 class SearchField(models.TextField):
@@ -10,7 +10,7 @@ class SearchField(models.TextField):
     def pre_save(self, model_instance, add):
         search_text = []
         for field_name in self.field_names:
-            val = unicode(to_ascii(getattr(model_instance, field_name)))
+            val = force_str(to_ascii(getattr(model_instance, field_name)))
             search_text.append(val)
         value = ' '.join(search_text)
         setattr(model_instance, self.name, value)
@@ -28,9 +28,9 @@ class SearchField(models.TextField):
 
 
 def to_ascii(txt, codif='utf-8'):
-    if not isinstance(txt, basestring):
-        txt = unicode(txt)
-    if isinstance(txt, unicode):
+    if not isinstance(txt, str):
+        txt = force_str(txt)
+    if isinstance(txt, str):
         txt = txt.encode('utf-8')
     return normalize('NFKD', txt.decode(codif)).encode('ASCII', 'ignore')
 
@@ -40,3 +40,21 @@ def queryset_ascii(self, request):
         request.GET._mutable = True
         request.GET['q'] = to_ascii(request.GET['q'])
     return admin.ModelAdmin.get_queryset(self, request)
+
+def field_label(name, model):
+    name = name.split('__')
+    try:
+        field = model._meta.get_field(name[0])
+    except FieldDoesNotExist:
+        return force_str(" ".join(name))
+
+    try:
+        label = force_str(field.verbose_name)
+    except AttributeError:
+        # field is likely a RelatedObject
+        label = force_str(field.name).capitalize()
+    if len(name) > 1:
+        to_model = field.get_path_info()[0].to_opts.model
+        label = label + "/" + field_label("__".join(name[1:]), to_model)
+
+    return label

sigi/apps/utils/mixins.py

@@ -1,13 +1,54 @@
 from collections import OrderedDict
 from functools import update_wrapper
+from django import forms
 from django.contrib import admin
 from django.contrib.admin import helpers
 from django.contrib.admin.options import csrf_protect_m
-from django.http.response import HttpResponseRedirect
+from django.core.exceptions import PermissionDenied
+from django.http.response import HttpResponse, HttpResponseRedirect
+from django.template.response import TemplateResponse
 from django.utils.translation import gettext as _, ngettext
+from import_export import resources
 from import_export.admin import ExportMixin
+from import_export.forms import ExportForm
+from import_export.signals import post_export
+from sigi.apps.utils import field_label
+
+class ExportFormFields(ExportForm):
+   def __init__(self, formats, field_list, *args, **kwargs):
+        super().__init__(formats, *args, **kwargs)
+        self.fields['selected_fields'] = forms.MultipleChoiceField(
+            label=_('Campos a exportar'),
+            required=True,
+            choices=field_list,
+            initial=[f[0] for f in field_list],
+            widget=forms.CheckboxSelectMultiple,
+        )
+
+class LabeledResourse(resources.ModelResource):
+    selected_fields = None
+    def get_export_headers(self):
+        headers = []
+        for field in self.get_export_fields():
+            if field.attribute == field.column_name:
+                label = field_label(field.attribute, self._meta.model)
+            else:
+                label = field.column_name
+            headers.append(label)
+        return headers
+
+    def get_export_fields(self):
+        fields = self.get_fields()
+        if self.selected_fields:
+            fields = [f for f in fields if f.column_name in self.selected_fields]
+        return fields
+
+    def export(self, queryset=None, selected_fields=None, *args, **kwargs):
+        self.selected_fields = selected_fields
+        return super().export(queryset, *args, **kwargs)
 
 class CartExportMixin(ExportMixin):
+    to_encoding = 'utf-8'
     actions = ['add_to_cart']
     change_list_template = 'admin/cart/change_list_cart_export.html'
     _cart_session_name = None
@@ -135,3 +176,44 @@ class CartExportMixin(ExportMixin):
         request.session.pop(self._cart_viewing_name, None)
         self.message_user(request, _(u"Carrinho vazio"))
         return HttpResponseRedirect('..')
+
+    def export_action(self, request, *args, **kwargs):
+        if not self.has_export_permission(request):
+            raise PermissionDenied
+
+        formats = self.get_export_formats()
+        resource = (self.get_export_resource_class())()
+        field_list = list(zip(resource.get_export_order(),
+                              resource.get_export_headers()))
+        form = ExportFormFields(formats, field_list, request.POST or None)
+        if form.is_valid():
+            file_format = formats[
+                int(form.cleaned_data['file_format'])
+            ]()
+
+            queryset = self.get_export_queryset(request)
+            export_data = self.get_export_data(
+                file_format,
+                queryset,
+                request=request,
+                encoding=self.to_encoding,
+                selected_fields=form.cleaned_data['selected_fields'])
+            content_type = file_format.get_content_type()
+            response = HttpResponse(export_data, content_type=content_type)
+            response['Content-Disposition'] = 'attachment; filename="%s"' % (
+                self.get_export_filename(request, queryset, file_format),
+            )
+
+            post_export.send(sender=None, model=self.model)
+            return response
+
+        context = self.get_export_context_data()
+
+        context.update(self.admin_site.each_context(request))
+
+        context['title'] = _("Export")
+        context['form'] = form
+        context['opts'] = self.model._meta
+        request.current_app = self.admin_site.name
+        return TemplateResponse(request, [self.export_template_name],
+                                context)

sigi/settings/development.py

@@ -28,7 +28,7 @@ DATABASES = {
         'NAME': 'sigi',
         'USER': 'sigi',
         'PASSWORD': '123456',
-        'HOST': 'database',
+        'HOST': 'localhost',
     }
 }