Add missing tests for data: URIs

Author: mattt

Diff for: python/cog/types.py

@@ -319,6 +319,14 @@ def get_filename_from_urlopen(resp: urllib.response.addinfourl) -> str:  # type:
 def get_filename_from_url(url: str) -> str:
     parsed_url = urllib.parse.urlparse(url)
 
+    if parsed_url.scheme == "data":
+        resp = urllib.request.urlopen(url)  # noqa: S310
+        mime_type = resp.headers.get_content_type()
+        extension = mimetypes.guess_extension(mime_type)
+        if extension is None:
+            return "file"
+        return "file" + extension
+
     filename = os.path.basename(parsed_url.path)
     filename = urllib.parse.unquote_plus(filename)
 

Diff for: python/tests/test_types.py

@@ -90,6 +90,19 @@ def test_urlfile_can_be_pickled_even_once_loaded():
             "https://example.com/%E1%9E%A0_%E1%9E%8F_%E1%9E%A2_%E1%9E%9C_%E1%9E%94_%E1%9E%93%E1%9E%87_%E1%9E%80_%E1%9E%9A%E1%9E%9F_%E1%9E%82%E1%9E%8F%E1%9E%9A%E1%9E%94%E1%9E%9F_%E1%9E%96_%E1%9E%9A_%E1%9E%99_%E1%9E%9F_%E1%9E%98_%E1%9E%93%E1%9E%A2_%E1%9E%8E_%E1%9E%85%E1%9E%98_%E1%9E%9B_Why_Was_The_Death_Of_Jesus_So_Powerful_.m4a",
             "ហ_ត_អ_វ_ប_នជ_ក_រស_គតរបស_ព_រ_យ_ស_ម_នអ_ណ_ចម_ល_Why_Was_The_Death_Of_Jesus_So_Powerful_.m4a",
         ),
+        # Data URIs
+        (
+            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUAAAAFCAYAAACNbyblAAAAHElEQVQI12P4//8/w38GIAXDIBKE0DHxgljNBAAO9TXL0Y4OHwAAAABJRU5ErkJggg==",
+            "file.png",
+        ),
+        (
+            "data:text/plain,hello world",
+            "file.txt",
+        ),
+        (
+            "data:application/data;base64,aGVsbG8gd29ybGQ=",
+            "file",
+        ),
         # Illegal characters
         ("https://example.com/nulbytes\u0000.wav", "nulbytes_.wav"),
         ("https://example.com/nulbytes%00.wav", "nulbytes_.wav"),