[async] URLFile code-review improvements

python/cog/types.py

@@ -22,13 +22,13 @@
 
 def Input(
     default: Any = ...,
-    description: str = None,
-    ge: float = None,
-    le: float = None,
-    min_length: int = None,
-    max_length: int = None,
-    regex: str = None,
-    choices: List[Union[str, int]] = None,
+    description: Optional[str] = None,
+    ge: Optional[float] = None,
+    le: Optional[float] = None,
+    min_length: Optional[int] = None,
+    max_length: Optional[int] = None,
+    regex: Optional[str] = None,
+    choices: Optional[List[Union[str, int]]] = None,
 ) -> Any:
     """Input is similar to pydantic.Field, but doesn't require a default value to be the first argument."""
     return Field(
@@ -205,6 +205,13 @@ class URLFile(io.IOBase):
 
     def __init__(self, url: str) -> None:
         parsed = urllib.parse.urlparse(url)
+        if parsed.scheme not in {
+            "http",
+            "https",
+        }:
+            raise ValueError(
+                "URLFile requires URL to conform to HTTP or HTTPS protocol"
+            )
         object.__setattr__(self, "name", os.path.basename(parsed.path))
         object.__setattr__(self, "__url__", url)
 

python/tests/server/test_clients.py

@@ -10,8 +10,6 @@
 import pytest
 from cog.server.clients import ClientManager
 
-pytest.mark.asyncio
-
 
 @pytest.mark.asyncio
 async def test_upload_files_without_url():

python/tests/test_types.py

@@ -19,6 +19,14 @@ def file_fixture(body: str):
     )
 
 
+def test_urlfile_protocol_validation():
+    with pytest.raises(ValueError):
+        URLFile("file:///etc/shadow")
+
+    with pytest.raises(ValueError):
+        URLFile("data:text/plain,hello")
+
+
 @mock.patch("urllib.request.urlopen", return_value=file_fixture("hello world"))
 def test_urlfile_acts_like_response(mock_urlopen):
     u = URLFile("https://example.com/some/url")