Add URL opening utilities (#142)

Author: cthoyt

src/pystow/utils/__init__.py

@@ -88,7 +88,7 @@
     stream_write_pydantic_jsonl,
     write_pydantic_jsonl,
 )
-from .safe_open import open_inner_zipfile, safe_open, safe_open_dict_reader
+from .safe_open import is_url, open_inner_zipfile, open_url, safe_open, safe_open_dict_reader
 from ..constants import README_TEXT, TimeoutHint
 
 if TYPE_CHECKING:
@@ -134,6 +134,7 @@
     "getenv_path",
     "gunzip",
     "gzip_compress",
+    "is_url",
     "iter_pydantic_jsonl",
     "iter_pydantic_tsv",
     "iter_tarred_csvs",
@@ -148,6 +149,7 @@
     "name_from_url",
     "open_inner_zipfile",
     "open_tarfile",
+    "open_url",
     "open_zip_reader",
     "open_zip_writer",
     "open_zipfile",

src/pystow/utils/safe_open.py

@@ -7,10 +7,11 @@
 import gzip
 import io
 import typing
+import urllib.request
 import zipfile
 from collections.abc import Generator, Mapping
 from pathlib import Path
-from typing import Any, Literal, TextIO, cast
+from typing import Any, BinaryIO, Literal, TextIO, cast, overload
 
 from .io_typing import (
     _MODE_TO_SIMPLE,
@@ -26,7 +27,9 @@
 )
 
 __all__ = [
+    "is_url",
     "open_inner_zipfile",
+    "open_url",
     "safe_open",
     "safe_open_dict_reader",
 ]
@@ -195,3 +198,39 @@ def safe_open_dict_reader(
     """
     with safe_open(f, operation="read", representation="text") as file:
         yield csv.DictReader(file, delimiter=delimiter, **kwargs)
+
+
+def is_url(s: str | Path | TextIO | Any) -> bool:
+    """Check if the object is a URL."""
+    if isinstance(s, str) and (s.startswith("http://") or s.startswith("https://")):
+        return True
+    return False
+
+
+# docstr-coverage:excused `overload`
+@overload
+@contextlib.contextmanager
+def open_url(
+    url: str, *, representation: Literal["text"] = ...
+) -> Generator[TextIO, None, None]: ...
+
+
+# docstr-coverage:excused `overload`
+@overload
+@contextlib.contextmanager
+def open_url(
+    url: str, *, representation: Literal["binary"] = ...
+) -> Generator[BinaryIO, None, None]: ...
+
+
+@contextlib.contextmanager
+def open_url(
+    url: str, *, representation: Representation = "text"
+) -> Generator[TextIO, None, None] | Generator[BinaryIO, None, None]:
+    """Get a file-like object from a URL."""
+    with urllib.request.urlopen(url) as response:  # noqa:S310
+        match representation:
+            case "text":
+                yield io.TextIOWrapper(response, encoding="utf-8")
+            case "binary":
+                yield io.BufferedReader(response)

tests/test_utils/test_utils.py

@@ -28,13 +28,15 @@
     getenv_path,
     gunzip,
     gzip_compress,
+    is_url,
     iter_tarred_csvs,
     iter_zipped_csvs,
     mkdir,
     mock_envvar,
     n,
     name_from_url,
     open_tarfile,
+    open_url,
     open_zip_reader,
     open_zip_writer,
     open_zipfile,
@@ -475,6 +477,27 @@ def test_gzip(self) -> None:
             self.assertTrue(path_new.is_file())
             self.assertEqual(TEST_TXT_CONTENT, path_new.read_text(encoding="utf-8"))
 
+    def test_is_url(self) -> None:
+        """Test checking URL."""
+        self.assertTrue(is_url("https://zenodo.org/records/15504009/files/startup.sh?download=1"))
+        self.assertTrue(is_url("https://zenodo.org/records/15504009/files/startup.sh"))
+        self.assertTrue(is_url("http://zenodo.org/records/15504009/files/startup.sh"))
+        self.assertFalse(is_url("ftp://zenodo.org/records/15504009/files/startup.sh"))
+        self.assertFalse(is_url("nope"))
+        self.assertFalse(is_url(Path().cwd()))
+
+    def test_open_url(self) -> None:
+        """Test opening a URL."""
+        with open_url(
+            "https://zenodo.org/records/15504009/files/startup.sh", representation="text"
+        ) as file:
+            self.assertIn("sleep 5", file.read())
+
+        with open_url(
+            "https://zenodo.org/records/15504009/files/startup.sh", representation="binary"
+        ) as file:
+            self.assertIn("sleep 5", file.read().decode("utf-8"))
+
 
 class TestDownload(unittest.TestCase):
     """Tests for downloading."""