⬇️🔗 url file download helper method

shroominic · Dec 3, 2024 · f8f4ada · f8f4ada
1 parent 26a65ed
commit f8f4ada
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 1 deletion.
diff --git a/src/codeboxapi/codebox.py b/src/codeboxapi/codebox.py
@@ -52,7 +52,8 @@ def __new__(cls, *args, **kwargs) -> "CodeBox":
         """
         Creates a CodeBox session
         """
-        api_key = kwargs.get("api_key") or os.getenv("CODEBOX_API_KEY", "local")
+        api_key = kwargs.get("api_key") or os.getenv("CODEBOX_API_KEY")
+        # todo make sure "local" is not hardcoded default
         if api_key == "local":
             return import_module("codeboxapi.local").LocalBox(*args, **kwargs)
 
@@ -175,6 +176,24 @@ async def ainstall(self, *packages: str) -> str:
         )
         return " ".join(packages) + " installed successfully"
 
+    async def afile_from_url(self, url: str, file_path: str) -> "RemoteFile":
+        """
+        Download a file from a URL to the specified destination in the CodeBox.
+        Example:
+        >>> codebox.afile_from_url("https://github.com/org/repo/file.txt", "file.txt")
+        """
+        code = (
+            "import httpx\n"
+            "async with httpx.AsyncClient() as client:\n"
+            f"    async with client.stream('GET', '{url}') as response:\n"
+            "        response.raise_for_status()\n"
+            f"        with open('{file_path}', 'wb') as f:\n"
+            "            async for chunk in response.aiter_bytes():\n"
+            "                f.write(chunk)\n"
+        )
+        await self.aexec(code)
+        return await self.adownload(file_path)
+
     async def alist_files(self) -> list["RemoteFile"]:
         from .types import RemoteFile
 
@@ -246,6 +265,9 @@ def healthcheck(self) -> str:
     def install(self, *packages: str) -> str:
         return syncify(self.ainstall)(*packages)
 
+    def file_from_url(self, url: str, file_path: str) -> "RemoteFile":
+        return syncify(self.afile_from_url)(url, file_path)
+
     def list_files(self) -> list["RemoteFile"]:
         return syncify(self.alist_files)()
 

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,11 +1,14 @@
 import os
 
 import pytest
+from dotenv import load_dotenv
 
 from codeboxapi import CodeBox
 
 LOCALBOX = CodeBox(api_key="local")
 
+load_dotenv()
+
 
 @pytest.fixture(
     scope="session",

diff --git a/tests/test_v02.py b/tests/test_v02.py
@@ -312,6 +312,27 @@ async def test_async_bash_commands(codebox: CodeBox):
     assert result.text.strip() == "Hello!", "Execution result should be 'Hello!'"
 
 
+def test_file_from_url(codebox: CodeBox):
+    url = "https://raw.githubusercontent.com/shroominic/codebox-api/main/README.md"
+    file_path = "README.md"
+    remote_file = codebox.file_from_url(url, file_path)
+    assert isinstance(remote_file, RemoteFile), "Should return a RemoteFile"
+    assert remote_file.path == file_path, "File path should match"
+    assert len(remote_file.get_content()) > 0, "File should have content"
+    assert file_path in [file.path for file in codebox.list_files()]
+
+
+@pytest.mark.asyncio
+async def test_file_from_url_async(codebox: CodeBox):
+    url = "https://raw.githubusercontent.com/shroominic/codebox-api/main/README.md"
+    file_path = "README.md"
+    remote_file = await codebox.afile_from_url(url, file_path)
+    assert isinstance(remote_file, RemoteFile), "Should return a RemoteFile"
+    assert remote_file.path == file_path, "File path should match"
+    assert len(remote_file.get_content()) > 0, "File should have content"
+    assert file_path in [file.path for file in await codebox.alist_files()]
+
+
 def test_local_box_singleton():
     from codeboxapi.local import LocalBox