remove duplicate code in files.models.FileManager

core/files/models.py

@@ -1,6 +1,5 @@
 from django.core.files.base import ContentFile
 from django.db import models, IntegrityError
-from django_softdelete.models import SoftDeleteModel
 
 from inventory.models import Item
 
@@ -10,7 +9,8 @@ def hash_upload(instance, filename):
 
 
 class FileManager(models.Manager):
-    def get_or_create(self, **kwargs):
+
+    def __file_data_helper(self, **kwargs):
         if 'data' in kwargs and type(kwargs['data']) == str:
             import base64
             from hashlib import sha256
@@ -31,6 +31,10 @@ class FileManager(models.Manager):
             pass
         else:
             raise ValueError('data must be a base64 encoded string or file and hash must be provided')
+        return kwargs
+
+    def get_or_create(self, **kwargs):
+        kwargs = self.__file_data_helper(**kwargs)
         try:
             return self.get(hash=kwargs['hash']), False
         except self.model.DoesNotExist:
@@ -39,26 +43,7 @@ class FileManager(models.Manager):
             return obj, True
 
     def create(self, **kwargs):
-        if 'data' in kwargs and type(kwargs['data']) == str:
-            import base64
-            from hashlib import sha256
-            raw = kwargs['data']
-            if not raw.startswith('data:'):
-                raise ValueError('data must be a base64 encoded string or file and hash must be provided')
-            raw = raw.split(';base64,')
-            if len(raw) != 2:
-                raise ValueError('data must be a base64 encoded string or file and hash must be provided')
-            mime_type = raw[0].split(':')[1]
-            content = base64.b64decode(raw[1], validate=True)
-            kwargs.pop('data')
-            content_hash = sha256(content).hexdigest()
-            kwargs['file'] = ContentFile(content, content_hash)
-            kwargs['hash'] = content_hash
-            kwargs['mime_type'] = mime_type
-        elif 'file' in kwargs and 'hash' in kwargs and type(kwargs['file']) == ContentFile and 'mime_type' in kwargs:
-            pass
-        else:
-            raise ValueError('data must be a base64 encoded string or file and hash must be provided')
+        kwargs = self.__file_data_helper(**kwargs)
         if not self.filter(hash=kwargs['hash']).exists():
             obj = super().create(**kwargs)
             obj.file.save(content=kwargs['file'], name=kwargs['hash'])