Add dataset type field + add filters to admin interface (#9)

* Add dataset type field + add filters to admin interface

* lint code

Author: willemarcel

vbos/datasets/admin.py

@@ -32,17 +32,20 @@ class RasterFileAdmin(admin.ModelAdmin):
 
 @admin.register(RasterDataset)
 class RasterDatasetAdmin(admin.ModelAdmin):
-    list_display = ["id", "cluster", "name", "created", "updated", "file"]
+    list_display = ["id", "name", "cluster", "type", "updated", "file"]
+    list_filter = ["cluster", "type"]
 
 
 @admin.register(VectorDataset)
 class VectorDatasetAdmin(admin.ModelAdmin):
-    list_display = ["id", "cluster", "name", "created", "updated"]
+    list_display = ["id", "name", "cluster", "type", "updated"]
+    list_filter = ["cluster", "type"]
 
 
 @admin.register(VectorItem)
 class VectorItemAdmin(admin.GISModelAdmin):
     list_display = ["id", "dataset", "metadata"]
+    list_filter = ["dataset"]
 
     def get_urls(self):
         urls = super().get_urls()
@@ -117,12 +120,14 @@ def import_file(self, request):
 
 @admin.register(TabularDataset)
 class TabularDatasetAdmin(admin.ModelAdmin):
-    list_display = ["id", "cluster", "name", "created", "updated"]
+    list_display = ["id", "name", "cluster", "type", "updated"]
+    list_filter = ["cluster", "type"]
 
 
 @admin.register(TabularItem)
 class TabularItemAdmin(admin.GISModelAdmin):
     list_display = ["id", "dataset", "data"]
+    list_filter = ["dataset"]
 
     def get_urls(self):
         urls = super().get_urls()

vbos/datasets/filters.py

@@ -13,6 +13,7 @@
 class DatasetFilter(FilterSet):
     name = CharFilter(field_name="name", lookup_expr="icontains")
     source = CharFilter(field_name="source", lookup_expr="icontains")
+    type = CharFilter(field_name="type", lookup_expr="iexact")
     cluster = ModelChoiceFilter(
         field_name="cluster__name",
         to_field_name="name__iexact",

vbos/datasets/migrations/0008_rasterdataset_type_tabulardataset_type_and_more.py

@@ -0,0 +1,64 @@
+# Generated by Django 5.2.5 on 2025-09-12 21:28
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("datasets", "0007_alter_rasterdataset_cluster_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="rasterdataset",
+            name="type",
+            field=models.CharField(
+                choices=[
+                    ("baseline", "Baseline"),
+                    ("estimated_damage", "Estimated Damage"),
+                    (
+                        "aid_resources_needed",
+                        "Resources Needed to be Sent to Those Affected",
+                    ),
+                    ("estimate_financial_damage", "Estimate Financial Damage"),
+                ],
+                default="baseline",
+                max_length=55,
+            ),
+        ),
+        migrations.AddField(
+            model_name="tabulardataset",
+            name="type",
+            field=models.CharField(
+                choices=[
+                    ("baseline", "Baseline"),
+                    ("estimated_damage", "Estimated Damage"),
+                    (
+                        "aid_resources_needed",
+                        "Resources Needed to be Sent to Those Affected",
+                    ),
+                    ("estimate_financial_damage", "Estimate Financial Damage"),
+                ],
+                default="baseline",
+                max_length=55,
+            ),
+        ),
+        migrations.AddField(
+            model_name="vectordataset",
+            name="type",
+            field=models.CharField(
+                choices=[
+                    ("baseline", "Baseline"),
+                    ("estimated_damage", "Estimated Damage"),
+                    (
+                        "aid_resources_needed",
+                        "Resources Needed to be Sent to Those Affected",
+                    ),
+                    ("estimate_financial_damage", "Estimate Financial Damage"),
+                ],
+                default="baseline",
+                max_length=55,
+            ),
+        ),
+    ]

vbos/datasets/models.py

@@ -1,12 +1,20 @@
 from django.contrib.gis.db import models
 from django.conf import settings
+from django.utils.translation import gettext_lazy as _
 from django.core.validators import FileExtensionValidator
 from django.db.models.fields.files import default_storage
 from django.db.models.signals import pre_delete
 from django.dispatch import receiver
 
 UPLOAD_TO = "staging/raster/" if settings.DEBUG else "production/raster/"
 
+TYPE_CHOICES = {
+    "baseline": _("Baseline"),
+    "estimated_damage": _("Estimated Damage"),
+    "aid_resources_needed": _("Resources Needed to be Sent to Those Affected"),
+    "estimate_financial_damage": _("Estimate Financial Damage"),
+}
+
 
 class Cluster(models.Model):
     name = models.CharField(max_length=100, unique=True)
@@ -57,6 +65,7 @@ class RasterDataset(models.Model):
         Cluster,
         on_delete=models.PROTECT,
     )
+    type = models.CharField(max_length=55, choices=TYPE_CHOICES, default="baseline")
     source = models.CharField(max_length=155, blank=True, null=True)
     file = models.ForeignKey(RasterFile, on_delete=models.PROTECT)
 
@@ -75,6 +84,7 @@ class VectorDataset(models.Model):
         Cluster,
         on_delete=models.PROTECT,
     )
+    type = models.CharField(max_length=55, choices=TYPE_CHOICES, default="baseline")
     source = models.CharField(max_length=155, blank=True, null=True)
 
     def __str__(self):
@@ -104,6 +114,7 @@ class TabularDataset(models.Model):
         Cluster,
         on_delete=models.PROTECT,
     )
+    type = models.CharField(max_length=55, choices=TYPE_CHOICES, default="baseline")
     source = models.CharField(max_length=155, blank=True, null=True)
 
     def __str__(self):

vbos/datasets/serializers.py

@@ -23,15 +23,24 @@ class RasterDatasetSerializer(serializers.ModelSerializer):
 
     class Meta:
         model = RasterDataset
-        fields = ["id", "name", "created", "updated", "cluster", "source", "file"]
+        fields = [
+            "id",
+            "name",
+            "created",
+            "updated",
+            "cluster",
+            "type",
+            "source",
+            "file",
+        ]
 
 
 class VectorDatasetSerializer(serializers.ModelSerializer):
     cluster = serializers.ReadOnlyField(source="cluster.name")
 
     class Meta:
         model = VectorDataset
-        fields = ["id", "name", "created", "updated", "cluster", "source"]
+        fields = ["id", "name", "created", "updated", "cluster", "type", "source"]
 
 
 class VectorItemSerializer(GeoFeatureModelSerializer):
@@ -61,7 +70,7 @@ class TabularDatasetSerializer(serializers.ModelSerializer):
 
     class Meta:
         model = TabularDataset
-        fields = ["id", "name", "created", "updated", "cluster", "source"]
+        fields = ["id", "name", "created", "updated", "cluster", "type", "source"]
 
 
 class TabularItemSerializer(serializers.ModelSerializer):

vbos/datasets/test/test_raster_views.py

@@ -24,6 +24,7 @@ def setUp(self):
             file=self.r_2,
             source="OSM",
             cluster=Cluster.objects.create(name="Administrative"),
+            type="estimated_damage",
         )
         self.url = reverse("datasets:raster-list")
 
@@ -37,6 +38,8 @@ def test_raster_datasets_list(self):
         assert req.data.get("results")[1]["cluster"] == "Administrative"
         assert req.data.get("results")[0]["source"] == "WMO"
         assert req.data.get("results")[1]["source"] == "OSM"
+        assert req.data.get("results")[0]["type"] == "baseline"
+        assert req.data.get("results")[1]["type"] == "estimated_damage"
 
     def test_raster_datasets_list_filter(self):
         req = self.client.get(self.url, {"cluster": "transportation"})
@@ -50,6 +53,10 @@ def test_raster_datasets_list_filter(self):
         assert req.status_code == status.HTTP_200_OK
         assert req.data.get("count") == 1
 
+        req = self.client.get(self.url, {"type": "baseline"})
+        assert req.status_code == status.HTTP_200_OK
+        assert req.data.get("count") == 1
+
     def test_raster_datasets_detail(self):
         url = reverse("datasets:raster-detail", args=[self.dataset_1.id])
         req = self.client.get(url)

vbos/datasets/test/test_tabular_views.py

@@ -18,6 +18,7 @@ def setUp(self):
             name="Prices",
             cluster=Cluster.objects.create(name="Statistics"),
             source="Government",
+            type="estimated_damage",
         )
         self.url = reverse("datasets:tabular-list")
 
@@ -31,6 +32,8 @@ def test_tabular_datasets_list(self):
         assert req.data.get("results")[1]["source"] == "Government"
         assert req.data.get("results")[0]["cluster"] == "Administrative"
         assert req.data.get("results")[1]["cluster"] == "Statistics"
+        assert req.data.get("results")[0]["type"] == "baseline"
+        assert req.data.get("results")[1]["type"] == "estimated_damage"
 
     def test_raster_datasets_list_filter(self):
         req = self.client.get(self.url, {"cluster": "transportation"})

vbos/datasets/test/test_vector_views.py

@@ -14,7 +14,9 @@ def setUp(self):
             source="OSM",
         )
         self.dataset_2 = VectorDataset.objects.create(
-            name="Roads", cluster=Cluster.objects.create(name="Transportation")
+            name="Roads",
+            cluster=Cluster.objects.create(name="Transportation"),
+            type="estimated_damage",
         )
         self.url = reverse("datasets:vector-list")
 
@@ -24,6 +26,8 @@ def test_vector_datasets_list(self):
         assert req.data.get("count") == 2
         assert req.data.get("results")[0]["name"] == "Boundaries"
         assert req.data.get("results")[1]["name"] == "Roads"
+        assert req.data.get("results")[0]["type"] == "baseline"
+        assert req.data.get("results")[1]["type"] == "estimated_damage"
 
     def test_vector_datasets_list_filter(self):
         req = self.client.get(self.url, {"cluster": "transportation"})