From b48bcd114c2996454e3a3e68375b4d2370a202ba Mon Sep 17 00:00:00 2001
From: Hironsan <light.tree.1.13@gmail.com>
Date: Mon, 19 Apr 2021 18:54:00 +0900
Subject: [PATCH] Add writer factory

---
 app/api/views/download/factory.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/app/api/views/download/factory.py b/app/api/views/download/factory.py
index 5974cc99..eb750be9 100644
--- a/app/api/views/download/factory.py
+++ b/app/api/views/download/factory.py
@@ -1,5 +1,7 @@
+from typing import Type
+
 from ...models import DOCUMENT_CLASSIFICATION, SEQ2SEQ, SEQUENCE_LABELING
-from . import repositories
+from . import catalog, repositories, writer
 
 
 def create_repository(project) -> repositories.BaseRepository:
@@ -12,3 +14,14 @@ def create_repository(project) -> repositories.BaseRepository:
         ValueError(f'Invalid project type: {project.project_type}')
     repository = mapping.get(project.project_type)(project)
     return repository
+
+
+def create_writer(format: str) -> Type[writer.BaseWriter]:
+    mapping = {
+        catalog.CSV.name: writer.CsvWriter,
+        catalog.JSONL.name: writer.JSONLWriter,
+        catalog.FastText.name: writer.FastTextWriter,
+    }
+    if format not in mapping:
+        ValueError(f'Invalid format: {format}')
+    return mapping[format]