right-tree/backend/right_tree/api/admin.py

from django.contrib import admin, messages
from django.http import HttpResponseRedirect, FileResponse
from django.utils.text import slugify
from django.utils import timezone

from right_tree.api import models
from right_tree.api.resource_generation_utils import storage


class ZoneAdmin(admin.ModelAdmin):
    ordering = ['name', 'variant', 'refined_variant', 'id']
    search_fields = ['name', 'habitat__name', 'variant', 'refined_variant', 'id']


class QuestionnaireAdmin(admin.ModelAdmin):
    list_display = ['address_display', 'location_display', 'soil_variant', 'ecological_district', 'habitat', 'zone']
    actions = ['export']

    @admin.display(description="Address")
    def address_display(self, obj):
        return obj.address.full_address

    @admin.display(description="Location")
    def location_display(self, obj):
        return f"({obj.location.x}, {obj.location.y})"

    @admin.action(description="Export planting guides for selected questionnaires")
    def export(self, request, queryset):
        export = models.Export.objects.create(creation_date=timezone.now())
        export.questionnaires.set(queryset)
        export.export()

        return HttpResponseRedirect('/admin/api/export')


class ExportAdmin(admin.ModelAdmin):
    list_display = ['creation_date', 'completion_date', 'completion_display']
    actions = ['download']

    @admin.display(description="Completion")
    def completion_display(self, obj):
        return f"{obj.completion:.1%}"

    @admin.action(description="Download completed exports")
    def download(self, request, queryset):
        if queryset.count() > 1:
            self.message_user(
                request,
                'Cannot download more than one export at a time.',
                messages.ERROR,
            )
            return

        export = queryset.first()

        if not export.complete:
            self.message_user(
                request,
                'Cannot download. Export is incomplete.',
                messages.ERROR,
            )
            return

        filename = f"export_{export.pk}/export.zip"
        filepath = storage.path(filename)

        if storage.exists(filename):
            return FileResponse(
                open(filepath, 'rb'),
                filename=f"export_{slugify(export.creation_date)}.zip",
            )

        self.message_user(
            request,
            'Cannot download. Export is corrupt.',
            messages.ERROR,
        )


admin.site.register(models.Plant)
admin.site.register(models.SoilOrder)
admin.site.register(models.SoilLayer)
admin.site.register(models.SoilVariant)
admin.site.register(models.EcologicalRegion)
admin.site.register(models.EcologicalDistrictLayer)
admin.site.register(models.ToleranceLevel)
admin.site.register(models.HabitatImage)
admin.site.register(models.Habitat)
admin.site.register(models.Zone, ZoneAdmin)
admin.site.register(models.ChristchurchRegion)
admin.site.register(models.Questionnaire, QuestionnaireAdmin)
admin.site.register(models.Export, ExportAdmin)