Open URLs with safe open (#143)

Author: cthoyt

src/pystow/utils/safe_open.py

@@ -11,7 +11,7 @@
 import zipfile
 from collections.abc import Generator, Mapping
 from pathlib import Path
-from typing import Any, BinaryIO, Literal, TextIO, cast, overload
+from typing import Any, BinaryIO, Literal, TextIO, TypeGuard, cast, overload
 
 from .io_typing import (
     _MODE_TO_SIMPLE,
@@ -87,7 +87,7 @@ def safe_open(
 
 
 @contextlib.contextmanager
-def safe_open(
+def safe_open(  # noqa:C901
     path: str | Path | typing.TextIO | typing.BinaryIO,
     *,
     operation: Operation = "read",
@@ -102,16 +102,26 @@ def safe_open(
         raise InvalidRepresentationError(representation)
 
     if isinstance(path, (str, Path)):
-        mode = MODE_MAP[operation, representation]
         encoding = ensure_sensible_default_encoding(encoding, representation=representation)
         newline = ensure_sensible_newline(newline, representation=representation)
-        path = Path(path).expanduser().resolve()
-        if path.suffix.endswith(".gz"):
-            with gzip.open(path, mode=mode, encoding=encoding, newline=newline) as file:
-                yield file  # type:ignore
+
+        if is_url(path):
+            if operation != "read":
+                raise ValueError('can only use operation="read" with URLs')
+            with open_url(
+                path, representation=representation, encoding=encoding, newline=newline
+            ) as file:
+                yield file
         else:
-            with open(path, mode=mode, encoding=encoding, newline=newline) as file:
-                yield file  # type:ignore
+            mode = MODE_MAP[operation, representation]
+            path = Path(path).expanduser().resolve()
+            if path.suffix.endswith(".gz"):
+                with gzip.open(path, mode=mode, encoding=encoding, newline=newline) as file:
+                    yield file  # type:ignore
+            else:
+                with open(path, mode=mode, encoding=encoding, newline=newline) as file:
+                    yield file  # type:ignore
+
     elif isinstance(path, typing.TextIO | io.TextIOWrapper | io.TextIOBase):
         if representation != "text":
             raise ValueError(
@@ -200,7 +210,7 @@ def safe_open_dict_reader(
         yield csv.DictReader(file, delimiter=delimiter, **kwargs)
 
 
-def is_url(s: str | Path | TextIO | Any) -> bool:
+def is_url(s: str | Path | TextIO | Any) -> TypeGuard[str]:
     """Check if the object is a URL."""
     if isinstance(s, str) and (s.startswith("http://") or s.startswith("https://")):
         return True
@@ -211,26 +221,38 @@ def is_url(s: str | Path | TextIO | Any) -> bool:
 @overload
 @contextlib.contextmanager
 def open_url(
-    url: str, *, representation: Literal["text"] = ...
+    url: str,
+    *,
+    representation: Literal["text"] = ...,
+    encoding: str | None = ...,
+    newline: str | None = ...,
 ) -> Generator[TextIO, None, None]: ...
 
 
 # docstr-coverage:excused `overload`
 @overload
 @contextlib.contextmanager
 def open_url(
-    url: str, *, representation: Literal["binary"] = ...
+    url: str,
+    *,
+    representation: Literal["binary"] = ...,
+    encoding: str | None = ...,
+    newline: str | None = ...,
 ) -> Generator[BinaryIO, None, None]: ...
 
 
 @contextlib.contextmanager
 def open_url(
-    url: str, *, representation: Representation = "text"
+    url: str,
+    *,
+    representation: Representation = "text",
+    encoding: str | None = None,
+    newline: str | None = None,
 ) -> 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")
+                yield io.TextIOWrapper(response, encoding=encoding, newline=newline)
             case "binary":
                 yield io.BufferedReader(response)

tests/test_utils/test_utils.py

@@ -391,6 +391,11 @@ def test_safe_open_exceptions(self) -> None:
                 with safe_open(passthrough, representation="binary") as _file:
                     pass
 
+        url = "https://zenodo.org/records/15504009/files/startup.sh"
+        with self.assertRaises(ValueError):
+            with safe_open(url, operation="write") as _file:
+                pass
+
     def test_safe_open_binary(self) -> None:
         """Test safe open in binary mode."""
         for path in [TEST_TXT, TEST_TXT_GZ]:
@@ -410,6 +415,13 @@ def test_safe_open_binary(self) -> None:
                             msg=f"failed to read bytes from {path} in a passthrough scenario",
                         )
 
+    def test_safe_open_url_binary(self) -> None:
+        """Test safe open in URL mode."""
+        with safe_open(
+            "https://zenodo.org/records/15504009/files/startup.sh", representation="binary"
+        ) as file:
+            self.assertIn("sleep 5", file.read().decode("utf-8"))
+
     def test_safe_open_text(self) -> None:
         """Test safe open in text mode."""
         for path, encoding, newline in itt.product(
@@ -433,6 +445,13 @@ def test_safe_open_text(self) -> None:
                             msg=f"failed to read text from {path} in a passthrough scenario",
                         )
 
+    def test_safe_open_url_text(self) -> None:
+        """Test safe open in URL mode."""
+        with safe_open(
+            "https://zenodo.org/records/15504009/files/startup.sh", representation="text"
+        ) as file:
+            self.assertIn("sleep 5", file.read())
+
     def test_encodings(self) -> None:
         """Test I/O in different encodings."""
         for encoding in ["ascii", "utf-16-be", "CP1252"]: