From 7406183a0c7ca72f202a0e8b6f132e90bdbf21c7 Mon Sep 17 00:00:00 2001
From: Hironsan <hiroki.nakayama.py@gmail.com>
Date: Fri, 3 Aug 2018 15:54:56 +0900
Subject: [PATCH] Update data download api

---
 app/server/models.py | 14 ++++++++++++++
 app/server/urls.py   |  4 ++--
 app/server/views.py  | 15 ++++++++++-----
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/app/server/models.py b/app/server/models.py
index 76bb1e51..30bbc407 100644
--- a/app/server/models.py
+++ b/app/server/models.py
@@ -131,6 +131,20 @@ class Document(models.Model):
         elif self.project.is_type_of(Project.Seq2seq):
             return self.seq2seq_annotations.all()
 
+    def make_dataset(self):
+        if self.project.is_type_of(Project.DOCUMENT_CLASSIFICATION):
+            return self.make_dataset_for_classification()
+        elif self.project.is_type_of(Project.SEQUENCE_LABELING):
+            return self.seq_annotations.all()
+        elif self.project.is_type_of(Project.Seq2seq):
+            return self.seq2seq_annotations.all()
+
+    def make_dataset_for_classification(self):
+        annotations = self.get_annotations()
+        dataset = [[a.document.id, a.document.text, a.label.text, a.user.username]
+                   for a in annotations]
+        return dataset
+
     def __str__(self):
         return self.text[:50]
 
diff --git a/app/server/urls.py b/app/server/urls.py
index 5745043c..72c298f0 100644
--- a/app/server/urls.py
+++ b/app/server/urls.py
@@ -2,7 +2,7 @@ from django.urls import path
 
 from .views import IndexView
 from .views import ProjectView, DatasetView, DatasetUpload
-from .views import ProjectsView, ProjectAdminView, RawDataAPI, DataDownloadAPI
+from .views import ProjectsView, ProjectAdminView, RawDataAPI, DataDownload
 from rest_framework import routers
 from .views import ProjectViewSet
 from .views import ProjectLabelsAPI, ProjectLabelAPI, ProjectDocsAPI, AnnotationsAPI, AnnotationAPI
@@ -21,7 +21,7 @@ urlpatterns = [
     path('api/projects/<int:project_id>/docs/<int:doc_id>/annotations/<int:annotation_id>', AnnotationAPI.as_view(), name='ann'),
     path('projects/', ProjectsView.as_view(), name='projects'),
     path('projects/<int:pk>/admin', ProjectAdminView.as_view(), name='project-admin'),
-    path('projects/<int:project_id>/download', DataDownloadAPI.as_view(), name='download'),
+    path('projects/<int:project_id>/download', DataDownload.as_view(), name='download'),
     path('projects/<int:project_id>/', ProjectView.as_view(), name='annotation'),
     path('projects/<int:pk>/apis/raw_data', RawDataAPI.as_view(), name='data_api'),
     path('projects/<int:pk>/docs', DatasetView.as_view(), name='dataset'),
diff --git a/app/server/views.py b/app/server/views.py
index 93bae389..39b02b85 100644
--- a/app/server/views.py
+++ b/app/server/views.py
@@ -93,13 +93,18 @@ class RawDataAPI(View):
         return JsonResponse({'status': 'ok'})
 
 
-class DataDownloadAPI(View):
+class DataDownload(View):
 
     def get(self, request, *args, **kwargs):
-        annotated_docs = [a.as_dict() for a in Annotation.objects.filter(manual=True)]
-        json_str = json.dumps(annotated_docs)
-        response = HttpResponse(json_str, content_type='application/json')
-        response['Content-Disposition'] = 'attachment; filename=annotation_data.json'
+        project_id = self.kwargs['project_id']
+        project = get_object_or_404(Project, pk=project_id)
+        docs = project.get_documents(is_null=False).distinct()
+        response = HttpResponse(content_type='text/csv')
+        response['Content-Disposition'] = 'attachment; filename="{}.csv"'.format(project.name)
+
+        writer = csv.writer(response)
+        for d in docs:
+            writer.writerows(d.make_dataset())
 
         return response