GEAR-348 Export selected dataset

backend/data_export/celery_tasks.py

@@ -24,12 +24,16 @@
 logger = get_task_logger(__name__)
 
 
-def create_collaborative_dataset(project: Project, dirpath: str, confirmed_only: bool, formatters, writer):
+def create_collaborative_dataset(
+    project: Project, dirpath: str, confirmed_only: bool, formatters, writer, example_ids=None
+):
     is_text_project = project.is_text_project
     if confirmed_only:
         examples = ExportedExample.objects.confirmed(project)
     else:
         examples = ExportedExample.objects.filter(project=project)
+    if example_ids:
+        examples = examples.filter(id__in=example_ids)
     labels = create_labels(project, examples)
     comments = create_comment(examples)
     dataset = Dataset(examples, labels, comments, is_text_project)
@@ -40,14 +44,18 @@ def create_collaborative_dataset(project: Project, dirpath: str, confirmed_only:
     service.export(filepath)
 
 
-def create_individual_dataset(project: Project, dirpath: str, confirmed_only: bool, formatters, writer):
+def create_individual_dataset(
+    project: Project, dirpath: str, confirmed_only: bool, formatters, writer, example_ids=None
+):
     is_text_project = project.is_text_project
     members = Member.objects.filter(project=project)
     for member in members:
         if confirmed_only:
             examples = ExportedExample.objects.confirmed(project, user=member.user)
         else:
             examples = ExportedExample.objects.filter(project=project)
+        if example_ids:
+            examples = examples.filter(id__in=example_ids)
         labels = create_labels(project, examples, member.user)
         comments = create_comment(examples, member.user)
         dataset = Dataset(examples, labels, comments, is_text_project)
@@ -69,17 +77,17 @@ def create_individual_dataset(project: Project, dirpath: str, confirmed_only: bo
     retry_backoff_max=30,
     max_retries=3,
 )
-def export_dataset(project_id, file_format: str, confirmed_only=False):
+def export_dataset(project_id, file_format: str, confirmed_only=False, example_ids=None):
     try:
         project = get_object_or_404(Project, pk=project_id)
         dirpath = os.path.join(settings.MEDIA_ROOT, str(uuid.uuid4()))
         os.makedirs(dirpath, exist_ok=True)
         formatters = create_formatter(project, file_format)
         writer = create_writer(file_format)
         if project.collaborative_annotation:
-            create_collaborative_dataset(project, dirpath, confirmed_only, formatters, writer)
+            create_collaborative_dataset(project, dirpath, confirmed_only, formatters, writer, example_ids)
         else:
-            create_individual_dataset(project, dirpath, confirmed_only, formatters, writer)
+            create_individual_dataset(project, dirpath, confirmed_only, formatters, writer, example_ids)
         zip_file = shutil.make_archive(dirpath, "zip", dirpath)
         shutil.rmtree(dirpath)
         return zip_file

backend/data_export/views.py

@@ -39,7 +39,12 @@ def post(self, request, *args, **kwargs):
         project_id = self.kwargs["project_id"]
         file_format = request.data.pop("format")
         export_approved = request.data.pop("exportApproved", False)
+        example_ids = request.data.pop("ids", None)
         task = export_dataset.delay(
-            project_id=project_id, file_format=file_format, confirmed_only=export_approved, **request.data
+            project_id=project_id,
+            file_format=file_format,
+            confirmed_only=export_approved,
+            example_ids=example_ids,
+            **request.data,
         )
         return Response({"task_id": task.task_id})

frontend/pages/projects/_id/dataset/export.vue

@@ -72,6 +72,18 @@ export default Vue.extend({
     example(): string {
       const item = this.formats.find((item: Format) => item.name === this.selectedFormat)
       return item!.example.trim()
+    },
+
+    // Ids of the examples selected on the dataset list page, if any were passed along
+    selectedIds(): number[] {
+      const ids = this.$route.query.ids
+      if (!ids) {
+        return []
+      }
+      return String(ids)
+        .split(',')
+        .map((x) => Number(x.trim()))
+        .filter((n) => Number.isFinite(n))
     }
   },
 
@@ -98,7 +110,8 @@ export default Vue.extend({
       this.taskId = await this.$repositories.download.prepare(
         this.projectId,
         this.selectedFormat,
-        this.exportApproved
+        this.exportApproved,
+        this.selectedIds
       )
       this.pollData()
     },

frontend/pages/projects/_id/dataset/index.vue

@@ -5,7 +5,7 @@
       <!-- Top action menu for import/export and assignment ops -->
       <action-menu
         @upload="$router.push('dataset/import')"
-        @download="$router.push('dataset/export')"
+        @download="goToExport"
         @assign="dialogAssignment = true"
         @reset="dialogReset = true"
       />
@@ -238,6 +238,15 @@ export default Vue.extend({
       this.$router.push(query)
     },
 
+    // Navigate to the export page, carrying over the selected example ids (if any)
+    goToExport() {
+      const query: Record<string, string> = {}
+      if (this.selected.length > 0) {
+        query.ids = this.selected.map((item) => item.id).join(',')
+      }
+      this.$router.push({ path: 'dataset/export', query })
+    },
+
     /**
      * Navigate to the labeling/annotation page.
      * getLinkToAnnotationPage() builds the correct route for the project's type.

frontend/repositories/download/apiDownloadRepository.ts

@@ -3,12 +3,20 @@ import ApiService from '@/services/api.service'
 export class APIDownloadRepository {
   constructor(private readonly request = ApiService) {}
 
-  async prepare(projectId: string, format: string, exportApproved: boolean): Promise<string> {
+  async prepare(
+    projectId: string,
+    format: string,
+    exportApproved: boolean,
+    ids?: number[]
+  ): Promise<string> {
     const url = `/projects/${projectId}/download`
-    const data = {
+    const data: { format: string; exportApproved: boolean; ids?: number[] } = {
       format,
       exportApproved
     }
+    if (ids && ids.length > 0) {
+      data.ids = ids
+    }
     const response = await this.request.post(url, data)
     return response.data.task_id
   }