fix: add arguments and fix export commands made for SEC-SEQ Symposium.

Author: mihow

ami/exports/base.py

@@ -1,5 +1,6 @@
 import csv
 import logging
+import typing
 
 from django.core.files.storage import default_storage
 from django.db import models
@@ -28,7 +29,11 @@ class BaseExportView(APIView):
     pass
 
 
-def get_data_in_batches(QuerySet: models.QuerySet, Serializer: type[serializers.Serializer], batch_size=1000):
+def get_data_in_batches(
+    QuerySet: models.QuerySet,
+    Serializer: type[serializers.Serializer],
+    batch_size: int = 1000,
+) -> typing.Iterator[list[dict]]:
     items = QuerySet.iterator(chunk_size=batch_size)
     batch = []
     for i, item in enumerate(items):
@@ -55,7 +60,11 @@ def get_data_in_batches(QuerySet: models.QuerySet, Serializer: type[serializers.
         yield batch
 
 
-def write_export(report_name, Serializer: type[serializers.Serializer], QuerySet: models.QuerySet):
+def write_export(
+    report_name: str,
+    Serializer: type[serializers.Serializer],
+    QuerySet: models.QuerySet,
+) -> str:
     timestamp = timezone.now().strftime("%Y%m%d-%H%M%S")
     file_name = f"{slugify(report_name)}-{timestamp}.csv"
     file_path = file_name

ami/exports/by_capture.py

@@ -40,11 +40,15 @@ class DetectionsByDeterminationAndCaptureTabularSerializer(serializers.Serialize
 
     def to_representation(self, instance: typing.Any) -> dict[str, typing.Any]:
         data = super().to_representation(instance)
-        taxon: Taxon = Taxon.objects.get(id=data["taxon_id"])
-
-        for taxon_rank in taxon.parents_json:
-            field_name = f"taxon_{taxon_rank.rank.name.lower()}"
-            data[field_name] = taxon_rank.name
+        try:
+            taxon: Taxon = Taxon.objects.get(id=data["taxon_id"])
+        except Taxon.DoesNotExist:
+            logger.warning(f"Taxon with ID '{data['taxon_id']}' not found")
+            pass
+        else:
+            for taxon_rank in taxon.parents_json:
+                field_name = f"taxon_{taxon_rank.rank.name.lower()}"
+                data[field_name] = taxon_rank.name
 
         return data
 

ami/exports/management/commands/export_by_capture.py

@@ -4,8 +4,10 @@
 """
 
 import logging
+import typing
 
 from django.core.management.base import BaseCommand
+from django.db import models
 
 from ami.exports import by_capture
 from ami.exports.base import write_export
@@ -16,18 +18,34 @@
 class Command(BaseCommand):
     help = "Export data by capture"
 
-    def handle(self, *args, **options):
-        # for i, batch in enumerate(by_capture.get_data_in_batches())
-        #     # print(f"Processing batch {batch}")
-        #     print(f"Processing batch {i}")
+    def add_arguments(self, parser) -> None:
+        parser.add_argument(
+            "--project-id",
+            type=int,
+            required=True,
+            help="Project ID to export data from",
+        )
+        parser.add_argument(
+            "--collection-ids",
+            type=int,
+            nargs="+",
+            required=False,
+            default=[],
+            help="Collection IDs to export data from (space-separated list)",
+        )
+
+    def handle(self, *args, **options) -> None:
+        project_id: int = options["project_id"]
+        collection_ids: list[int] = options["collection_ids"]
+
+        qs = by_capture.get_queryset().filter(occurrence__project=project_id)
+        if collection_ids:
+            qs = qs.filter(source_image__collections__in=collection_ids)
 
         fname = write_export(
             "detections_by_determination_and_capture",
             Serializer=by_capture.DetectionsByDeterminationAndCaptureTabularSerializer,
-            QuerySet=by_capture.get_queryset()
-            .filter(occurrence__project=85)
-            .filter(source_image__collections__in=[82, 79]),
-            # .filter(source_image__collections__in=[82]),
+            QuerySet=typing.cast(models.QuerySet, qs),
         )
         # get full path to the file
         print(f"Exported to {fname}")

ami/exports/management/commands/export_by_detection.py

@@ -14,23 +14,42 @@
 
 
 class Command(BaseCommand):
-    help = "Export data by capture"
+    help = "Export data by detection and determination"
+
+    def add_arguments(self, parser) -> None:
+        parser.add_argument(
+            "--project-id",
+            type=int,
+            required=True,
+            help="Project ID to export data from",
+        )
+        parser.add_argument(
+            "--collection-ids",
+            type=int,
+            nargs="+",
+            required=False,
+            default=[],
+            help="Collection IDs to export data from (space-separated list)",
+        )
 
     def handle(self, *args, **options):
         # for i, batch in enumerate(by_capture.get_data_in_batches())
         #     # print(f"Processing batch {batch}")
         #     print(f"Processing batch {i}")
+        project_id: int = options["project_id"]
+        collection_ids: list[int] = options["collection_ids"]
+
+        qs = by_detection.get_queryset().filter(occurrence__project=project_id)
+        if collection_ids:
+            qs = qs.filter(source_image__collections__in=collection_ids)
 
         fname = write_export(
             "detections",
             Serializer=by_detection.DetectionsTabularSerializer,
-            QuerySet=by_detection.get_queryset()
-            .filter(occurrence__project=85)
-            .filter(source_image__collections__in=[82, 79]),
-            # .filter(source_image__collections__in=[82]),
+            QuerySet=qs,
         )
         # get full path to the file
         print(f"Exported to {fname}")
 
-        logger.info("Export by capture completed")
-        self.stdout.write(self.style.SUCCESS("Export by capture completed"))
+        logger.info("Export by detection completed")
+        self.stdout.write(self.style.SUCCESS("Export by detection completed"))

ami/exports/management/commands/export_captures.py

@@ -16,16 +16,37 @@
 class Command(BaseCommand):
     help = "Export data by capture"
 
+    def add_arguments(self, parser) -> None:
+        parser.add_argument(
+            "--project-id",
+            type=int,
+            required=True,
+            help="Project ID to export data from",
+        )
+        parser.add_argument(
+            "--collection-ids",
+            type=int,
+            nargs="+",
+            required=False,
+            default=[],
+            help="Collection IDs to export data from (space-separated list)",
+        )
+
     def handle(self, *args, **options):
         # for i, batch in enumerate(by_capture.get_data_in_batches())
         #     # print(f"Processing batch {batch}")
         #     print(f"Processing batch {i}")
+        project_id: int = options["project_id"]
+        collection_ids: list[int] = options["collection_ids"]
+
+        qs = all_captures.get_queryset().filter(project=project_id)
+        if collection_ids:
+            qs = qs.filter(collections__in=collection_ids)
 
         fname = write_export(
             "captures",
             Serializer=all_captures.CapturesTabularSerializer,
-            QuerySet=all_captures.get_queryset().filter(project=85).filter(collections__in=[82, 79]),
-            # .filter(collections__in=[82]),
+            QuerySet=qs,
         )
         # get full path to the file
         print(f"Exported to {fname}")

ami/exports/management/commands/export_sessions.py

@@ -5,7 +5,7 @@
 
 import logging
 
-from django.core.management.base import BaseCommand
+from django.core.management.base import BaseCommand, CommandParser
 
 from ami.exports import all_sessions
 from ami.exports.base import write_export
@@ -16,15 +16,24 @@
 class Command(BaseCommand):
     help = "Export data by capture"
 
+    def add_arguments(self, parser: CommandParser) -> None:
+        parser.add_argument(
+            "--project-id",
+            type=int,
+            required=True,
+            help="Project ID to export data from",
+        )
+
     def handle(self, *args, **options):
         # for i, batch in enumerate(by_capture.get_data_in_batches())
         #     # print(f"Processing batch {batch}")
         #     print(f"Processing batch {i}")
+        project_id: int = options["project_id"]
 
         fname = write_export(
             "sessions",
             Serializer=all_sessions.SessionsTabularSerializer,
-            QuerySet=all_sessions.get_queryset().filter(project=85),
+            QuerySet=all_sessions.get_queryset().filter(project=project_id),
         )
         # get full path to the file
         print(f"Exported to {fname}")