fix: FileMetadata export_links type definition

GooeyAI · Dec 23, 2024 · 376e5d1 · 376e5d1
1 parent b91b325
commit 376e5d1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/daras_ai_v2/gdrive_downloader.py b/daras_ai_v2/gdrive_downloader.py
@@ -7,7 +7,7 @@
 from daras_ai_v2.functional import flatmap_parallel
 from daras_ai_v2.exceptions import raise_for_status
 
-docs_export_mimetype = {
+DOCS_EXPORT_MIMETYPES = {
     "application/vnd.google-apps.document": "text/plain",
     "application/vnd.google-apps.spreadsheet": "text/csv",
     "application/vnd.google-apps.presentation": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
@@ -85,7 +85,7 @@ def gdrive_download(
 
     if f.host != "drive.google.com":
         # export google docs to appropriate type
-        export_mime_type = docs_export_mimetype.get(mime_type, mime_type)
+        export_mime_type = DOCS_EXPORT_MIMETYPES.get(mime_type, mime_type)
         if f_url_export := export_links.get(export_mime_type, None):
             r = requests.get(f_url_export)
             file_bytes = r.content

diff --git a/files/models.py b/files/models.py
@@ -7,10 +7,7 @@ class FileMetadata(models.Model):
     etag = models.CharField(max_length=255, null=True)
     mime_type = models.CharField(max_length=255, default="", blank=True)
     total_bytes = models.PositiveIntegerField(default=0, blank=True)
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.export_links = {}
+    export_links: dict[str, str] | None = None
 
     def __str__(self):
         ret = f"{self.name or 'Unnamed'} - {self.mime_type}"