diff --git a/api_v2/management/commands/export.py b/api_v2/management/commands/export.py index 5513a6c7..cadf64c5 100644 --- a/api_v2/management/commands/export.py +++ b/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,8 +140,9 @@ 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))) @@ -134,21 +150,25 @@ def handle(self, *args, **options) -> None: 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' @@ -156,18 +176,31 @@ def get_filepath_by_model(model_name, app_label, pub_key=None, doc_key=None, bas 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()) +