diff --git a/archivist/sites/furaffinity/management/commands/import_data.py b/archivist/sites/furaffinity/management/commands/import_data.py index 63907bd..5dc4334 100644 --- a/archivist/sites/furaffinity/management/commands/import_data.py +++ b/archivist/sites/furaffinity/management/commands/import_data.py @@ -21,33 +21,35 @@ class Command(BaseCommand): def add_arguments(self, parser): parser.add_argument('path', type=str, help='Path to the folder containing JSON files or a single JSON file') + parser.add_argument('--delete', action='store_true', help='Delete imported files') def handle(self, *args, **kwargs): path = kwargs['path'] + delete = kwargs['delete'] if os.path.isfile(path): - self.process_json_file(path) + self.process_json_file(path, delete) elif os.path.isdir(path): - self.process_json_folder(path) + self.process_json_folder(path, delete) else: self.stdout.write(self.style.ERROR(f"The path '{path}' is not a valid file or folder.")) return - def process_json_file(self, file_path): + def process_json_file(self, file_path, delete): #self.stdout.write(self.style.NOTICE(f"Importing data from: {file_path}")) tqdm.write(self.style.NOTICE(f"Importing data from: {file_path}")) with open(file_path) as f: data = json.load(f) - self.import_data(data, file_path) + self.import_data(data, file_path, delete) #self.stdout.write(self.style.SUCCESS('Data imported successfully.')) tqdm.write(self.style.SUCCESS('Data imported successfully.')) - def process_json_folder(self, folder_path): + def process_json_folder(self, folder_path, delete): if not os.path.exists(folder_path): #self.stdout.write(self.style.ERROR(f"The folder '{folder_path}' does not exist.")) tqdm.write(self.style.ERROR(f"The folder '{folder_path}' does not exist.")) @@ -57,10 +59,11 @@ class Command(BaseCommand): for file_name in files: if file_name.endswith('.json'): file_path = os.path.join(root, file_name) - self.process_json_file(file_path) + self.process_json_file(file_path, delete) + - def import_data(self, data, json_file_path): + def import_data(self, data, json_file_path, delete): #self.stdout.write(self.style.NOTICE(data)) submission, created = FA_Submission.objects.get_or_create(submission_id=data["id"]) @@ -130,9 +133,13 @@ class Command(BaseCommand): file_instance.file_name = file_name file_instance.save() + + tqdm.write(self.style.NOTICE(f"Import media file: {file_path}")) # Now link the image_instance to your_model_instance submission.file = file_instance + + else: #self.stdout.write(self.style.WARNING(f"File not found: {file_path}")) tqdm.write(self.style.WARNING(f"File not found: {file_path}")) @@ -150,8 +157,11 @@ class Command(BaseCommand): # submission.image.save(os.path.basename(image_file_path), File(img_file), save=True) submission.save() - - + + self.delete_imported_file(json_file_path, delete) + self.delete_imported_file(file_path, delete) + + def compute_file_hash(self, file_path): try: # Compute BLAKE3 hash of the file @@ -165,3 +175,12 @@ class Command(BaseCommand): tqdm.write(self.style.WARNING(f"Error computing file hash: {e}")) return None + + # Delete the file if the --delete flag is used + def delete_imported_file(self, file_path, delete): + if delete: + if os.path.exists(file_path): + os.remove(file_path) + tqdm.write(self.style.SUCCESS(f"Deleted: {file_path}")) + else: + tqdm.write(self.style.WARNING(f"File not found: {file_path}"))