Refactor URLFile to use urllib.response instead of httpx

This commit refactors the existing change to use httpx instead to use
the urllib module from the standard library instead. This already
conforms to a file like object and fits in better with the design.

python/cog/server/clients.py

@@ -130,7 +130,11 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
         if self.fh.seekable():
             self.fh.seek(0)
 
-        for chunk in iter(self.fh):
+        while True:
+            chunk = self.fh.read(1024 * 1024)
+            if not chunk:
+                log.info("finished reading file")
+                break
             yield chunk
 
 

python/cog/types.py

@@ -237,15 +237,7 @@ def __delattr__(self, name: str) -> None:
     # Luckily the only dunder method on HTTPResponse is __iter__
     def __iter__(self) -> Iterator[bytes]:
         response = self.__wrapped__
-
-        # URLFile doesn't support `seek()` so the upload_file function
-        # cannot reset the request should it need to retry. As a hack
-        # we re-create the request each time the iterator is requested.
-        if response.num_bytes_downloaded > 0:
-            response.close()
-            object.__delattr__(self, "__target__")
-
-        yield from self.__wrapped__.iter_raw()
+        return iter(response)
 
     @property
     def __wrapped__(self) -> Any:
@@ -259,9 +251,7 @@ def __wrapped__(self) -> Any:
             # is that the book keeping for closing the response needs to be
             # handled elsewhere. There's probably a better design for this
             # in the long term.
-            client = httpx.Client()
-            req = client.build_request("GET", url)
-            res = client.send(request=req, stream=True)
+            res = urllib.request.urlopen(url)
             object.__setattr__(self, "__target__", res)
 
             return res

python/tests/server/test_clients.py

@@ -1,13 +1,18 @@
+from email.message import Message
+import io
 import os
 import tempfile
+from urllib.response import addinfourl
+from unittest import mock
 
 import cog
 import httpx
 import pytest
 from cog.server.clients import ClientManager
 
+pytest.mark.asyncio
+
 
-@pytest.mark.asyncio
 async def test_upload_files_without_url():
     client_manager = ClientManager()
     temp_dir = tempfile.mkdtemp()
@@ -111,9 +116,11 @@ async def test_upload_files_with_retry(respx_mock):
 
 @pytest.mark.asyncio
 @pytest.mark.respx(base_url="https://example.com")
-async def test_upload_files_with_url_file(respx_mock):
-    source = respx_mock.get("/cdn/my_file.txt").mock(
-        return_value=httpx.Response(200, text="hello world")
+@mock.patch("urllib.request.urlopen")
+async def test_upload_files_with_url_file(urlopen_mock, respx_mock):
+    fp = io.BytesIO(b"hello world")
+    urlopen_mock.return_value = addinfourl(
+        fp=fp, headers=Message(), url="https://example.com/cdn/my_file.txt"
     )
 
     uploader = respx_mock.put("/bucket/my_file.txt").mock(
@@ -131,14 +138,17 @@ async def test_upload_files_with_url_file(respx_mock):
     assert result == {"path": "https://cdn.example.com/bucket/my_file.txt"}
 
     assert uploader.call_count == 1
-    assert source.call_count == 1
+    assert urlopen_mock.call_count == 1
+    assert urlopen_mock.call_args[0][0] == "https://example.com/cdn/my_file.txt"
 
 
 @pytest.mark.asyncio
 @pytest.mark.respx(base_url="https://example.com")
-async def test_upload_files_with_url_file_with_retry(respx_mock):
-    source = respx_mock.get("/cdn/my_file.txt").mock(
-        return_value=httpx.Response(200, text="hello world")
+@mock.patch("urllib.request.urlopen")
+async def test_upload_files_with_url_file_with_retry(urlopen_mock, respx_mock):
+    fp = io.BytesIO(b"hello world")
+    urlopen_mock.return_value = addinfourl(
+        fp=fp, headers=Message(), url="https://example.com/cdn/my_file.txt"
     )
 
     uploader = respx_mock.put("/bucket/my_file.txt").mock(
@@ -154,4 +164,5 @@ async def test_upload_files_with_url_file_with_retry(respx_mock):
         )
 
     assert uploader.call_count == 3
-    assert source.call_count == 3
+    assert urlopen_mock.call_count == 1
+    assert urlopen_mock.call_args[0][0] == "https://example.com/cdn/my_file.txt"