Merge pull request #516 from open5e/397-exportconversion-to-other-for…

…mats-and-systems

Adding a functional csv export.

api_v2/management/commands/export.py

@@ -1,6 +1,7 @@
 
 
 import os
+import csv
 import json
 import time
 
@@ -26,6 +27,14 @@ def add_arguments(self, parser):
                             type=str,
                             help="Directory to write files to.")
 
+        parser.add_argument("-f",
+                            "--format",
+                            type=str,
+                            choices = ['csv','json'],
+                            default = 'json',
+                            help="File format of the output files.")
+
+
     def handle(self, *args, **options) -> None:
 
         self.stdout.write('Checking if directory exists.')
@@ -46,8 +55,9 @@ def handle(self, *args, **options) -> None:
                 "Document",
                 "api",
                 doc_key=v1doc.slug,
-                base_path=options['dir'])
-            write_queryset_data(v1doc_path, v1docq)
+                base_path=options['dir'],
+                format=options['format'])
+            write_queryset_data(v1doc_path, v1docq, format=options['format'])
 
             for model in app_models:
                 if model._meta.app_label == 'api':
@@ -63,8 +73,9 @@ def handle(self, *args, **options) -> None:
                     model.__name__,
                     model._meta.app_label,
                     doc_key=v1doc.slug,
-                    base_path=options['dir'])
-                write_queryset_data(model_path, modelq)
+                    base_path=options['dir'],
+                    format=options['format'])
+                write_queryset_data(model_path, modelq, format=options['format'])
 
             self.stdout.write(self.style.SUCCESS(
                 'Wrote {} to {}'.format(v1doc.slug, v1doc_path)))
@@ -76,15 +87,17 @@ def handle(self, *args, **options) -> None:
         ruleset_path = get_filepath_by_model(
             'Ruleset',
             'api_v2',
-            base_path=options['dir'])
-        write_queryset_data(ruleset_path, rulesets)
+            base_path=options['dir'],
+            format=options['format'])
+        write_queryset_data(ruleset_path, rulesets, format=options['format'])
 
         license_path = get_filepath_by_model(
             'License',
             'api_v2',
-            base_path=options['dir'])
+            base_path=options['dir'],
+            format=options['format'])
         licenses = License.objects.all()
-        write_queryset_data(license_path, licenses)
+        write_queryset_data(license_path, licenses, format=options['format'])
 
         # Create a folder and Publisher fixture for each pubishing org.
         for pub in Publisher.objects.order_by('key'):
@@ -93,8 +106,9 @@ def handle(self, *args, **options) -> None:
                 "Publisher",
                 "api_v2",
                 pub_key=pub.key,
-                base_path=options['dir'])
-            write_queryset_data(pub_path, pubq)
+                base_path=options['dir'],
+                format=options['format'])
+            write_queryset_data(pub_path, pubq, format=options['format'])
 
             # Create a Document fixture for each document.
             for doc in Document.objects.filter(publisher=pub):
@@ -104,8 +118,9 @@ def handle(self, *args, **options) -> None:
                     "api_v2",
                     pub_key=pub.key,
                     doc_key=doc.key,
-                    base_path=options['dir'])
-                write_queryset_data(doc_path, docq)
+                    base_path=options['dir'],
+                    format=options['format'])
+                write_queryset_data(doc_path, docq, format=options['format'])
 
                 for model in app_models:
                     SKIPPED_MODEL_NAMES = ['Document', 'Ruleset', 'License', 'Publisher','SearchResult']
@@ -125,49 +140,67 @@ def handle(self, *args, **options) -> None:
                             model._meta.app_label,
                             pub_key=pub.key,
                             doc_key=doc.key,
-                            base_path=options['dir'])
-                        write_queryset_data(model_path, modelq)
+                            base_path=options['dir'],
+                            format=options['format'])
+                        write_queryset_data(model_path, modelq, format=options['format'])
 
                 self.stdout.write(self.style.SUCCESS(
                     'Wrote {} to {}'.format(doc.key, doc_path)))
 
         self.stdout.write(self.style.SUCCESS('Data for v2 data complete.'))
 
 
-def get_filepath_by_model(model_name, app_label, pub_key=None, doc_key=None, base_path=None):
+def get_filepath_by_model(model_name, app_label, pub_key=None, doc_key=None, base_path=None, format='json'):
+    if not format.startswith('.'):
+        file_ext = "."+format
+    else:
+        file_ext = format
 
     if app_label == "api_v2":
         root_folder_name = 'v2'
         root_models = ['License', 'Ruleset']
         pub_models = ['Publisher']
-
+        
         if model_name in root_models:
-            return "/".join((base_path, root_folder_name, model_name+".json"))
+            return "/".join((base_path, root_folder_name, model_name+file_ext))
 
         if model_name in pub_models:
-            return "/".join((base_path, root_folder_name, pub_key, model_name+".json"))
+            return "/".join((base_path, root_folder_name, pub_key, model_name+file_ext))
 
         else:
-            return "/".join((base_path, root_folder_name, pub_key, doc_key, model_name+".json"))
+            return "/".join((base_path, root_folder_name, pub_key, doc_key, model_name+file_ext))
 
     if app_label == "api":
         root_folder_name = 'v1'
         root_models = ['Manifest']
         doc_folder_name = doc_key
 
         if model_name in root_models:
-            return "/".join((base_path, root_folder_name, model_name+".json"))
+            return "/".join((base_path, root_folder_name, model_name+file_ext))
 
         else:
-            return "/".join((base_path, root_folder_name, doc_key, model_name+".json"))
+            return "/".join((base_path, root_folder_name, doc_key, model_name+file_ext))
 
 
-def write_queryset_data(filepath, queryset):
+def write_queryset_data(filepath, queryset, format='json'):
     if queryset.count() > 0:
         dir = os.path.dirname(filepath)
         if not os.path.exists(dir):
             os.makedirs(dir)
 
         output_filepath = filepath
+
         with open(output_filepath, 'w', encoding='utf-8') as f:
-            serializers.serialize("json", queryset, indent=2, stream=f)
+            if format=='json':
+                serializers.serialize("json", queryset, indent=2, stream=f)
+            if format=='csv':
+                # Create headers:
+                fieldnames = []
+                for field in queryset.first().__dict__.keys():
+                    if not field.startswith("_"):
+                        fieldnames.append(field)
+
+                writer = csv.DictWriter(f, fieldnames=fieldnames)
+                writer.writeheader()
+                writer.writerows(queryset.values())
+