upath.registry: narrow get_upath_class types on protocol

typesafety/test_upath_types.yml

@@ -19,3 +19,74 @@
     b: JoinablePathLike = PurePath()
     c: JoinablePathLike = Path()
     d: JoinablePathLike = UPath()
+
+- case: get_upath_class_fsspec_protocols
+  disable_cache: false
+  parametrized:
+    - cls_fqn: upath.implementations.cached.SimpleCachePath
+      protocol: simplecache
+    - cls_fqn: upath.implementations.cloud.S3Path
+      protocol: s3
+    - cls_fqn: upath.implementations.cloud.GCSPath
+      protocol: gcs
+    - cls_fqn: upath.implementations.cloud.AzurePath
+      protocol: abfs
+    - cls_fqn: upath.implementations.data.DataPath
+      protocol: data
+    - cls_fqn: upath.implementations.github.GitHubPath
+      protocol: github
+    - cls_fqn: upath.implementations.hdfs.HDFSPath
+      protocol: hdfs
+    - cls_fqn: upath.implementations.http.HTTPPath
+      protocol: http
+    - cls_fqn: upath.implementations.local.FilePath
+      protocol: file
+    - cls_fqn: upath.implementations.memory.MemoryPath
+      protocol: memory
+    - cls_fqn: upath.implementations.sftp.SFTPPath
+      protocol: sftp
+    - cls_fqn: upath.implementations.smb.SMBPath
+      protocol: smb
+    - cls_fqn: upath.implementations.webdav.WebdavPath
+      protocol: webdav
+  main: |
+    from upath.registry import get_upath_class
+
+    path_cls = get_upath_class("{{ protocol }}")
+    reveal_type(path_cls)  # N: Revealed type is "type[{{ cls_fqn }}]"
+
+- case: get_upath_class_pathlib_unix
+  disable_cache: false
+  skip: sys.platform == "win32"
+  main: |
+    from upath.registry import get_upath_class
+
+    path_cls = get_upath_class("")
+    reveal_type(path_cls)  # N: Revealed type is "type[upath.implementations.local.PosixUPath]"
+
+- case: get_upath_class_pathlib_win
+  disable_cache: false
+  skip: sys.platform != "win32"
+  main: |
+    from upath.registry import get_upath_class
+
+    path_cls = get_upath_class("")
+    reveal_type(path_cls)  # N: Revealed type is "type[upath.implementations.local.WindowsUPath]"
+
+- case: get_upath_class_unknown_protocol_py39
+  disable_cache: false
+  skip: sys.version_info >= (3, 10)
+  main: |
+    from upath.registry import get_upath_class
+
+    path_cls = get_upath_class("some-unknown-protocol")
+    reveal_type(path_cls)  # N: Revealed type is "Union[type[upath.core.UPath], None]"
+
+- case: get_upath_class_unknown_protocol_py310plus
+  disable_cache: false
+  skip: sys.version_info < (3, 10)
+  main: |
+    from upath.registry import get_upath_class
+
+    path_cls = get_upath_class("some-unknown-protocol")
+    reveal_type(path_cls)  # N: Revealed type is "type[upath.core.UPath] | None"

upath/implementations/_experimental.py

@@ -10,8 +10,17 @@
 
 def __getattr__(name: str) -> type[UPath]:
     if name.startswith("_") and name.endswith("Path"):
+        from upath import UPath
+
         protocol = name[1:-4].lower()
-        cls = get_upath_class(protocol, fallback=False)
-        assert cls is not None
+        cls = get_upath_class(protocol)
+        if cls is None:
+            raise RuntimeError(
+                f"Could not find fsspec implementation for protocol {protocol!r}"
+            )
+        elif not issubclass(cls, UPath):
+            raise RuntimeError(
+                "UPath implementation not a subclass of upath.UPath, {cls!r}"
+            )
         return cls
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

upath/registry.py

@@ -48,6 +48,27 @@
 
 import upath
 
+if TYPE_CHECKING:
+    from typing import Literal
+    from typing import overload
+
+    from upath.implementations.cached import SimpleCachePath as _SimpleCachePath
+    from upath.implementations.cloud import AzurePath as _AzurePath
+    from upath.implementations.cloud import GCSPath as _GCSPath
+    from upath.implementations.cloud import S3Path as _S3Path
+    from upath.implementations.data import DataPath as _DataPath
+    from upath.implementations.github import GitHubPath as _GitHubPath
+    from upath.implementations.hdfs import HDFSPath as _HDFSPath
+    from upath.implementations.http import HTTPPath as _HTTPPath
+    from upath.implementations.local import FilePath as _FilePath
+    from upath.implementations.local import PosixUPath as _PosixUPath
+    from upath.implementations.local import WindowsUPath as _WindowsUPath
+    from upath.implementations.memory import MemoryPath as _MemoryPath
+    from upath.implementations.sftp import SFTPPath as _SFTPPath
+    from upath.implementations.smb import SMBPath as _SMBPath
+    from upath.implementations.webdav import WebdavPath as _WebdavPath
+
+
 __all__ = [
     "get_upath_class",
     "available_implementations",
@@ -125,7 +146,7 @@ def __setitem__(self, item: str, value: type[upath.UPath] | str) -> None:
                 f"expected UPath subclass or FQN-string, got: {type(value).__name__!r}"
             )
         if not item or item in self._m:
-            get_upath_class.cache_clear()
+            get_upath_class.cache_clear()  # type: ignore[attr-defined]
         self._m[item] = value
 
     def __delitem__(self, __v: str) -> None:
@@ -182,7 +203,57 @@ def register_implementation(
     _registry[protocol] = cls
 
 
-@lru_cache
+# --- get_upath_class type overloads ------------------------------------------
+
+if TYPE_CHECKING:  # noqa: C901
+
+    @overload
+    def get_upath_class(protocol: Literal["simplecache"]) -> type[_SimpleCachePath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["s3", "s3a"]) -> type[_S3Path]: ...
+    @overload
+    def get_upath_class(protocol: Literal["gcs", "gs"]) -> type[_GCSPath]: ...
+
+    @overload
+    def get_upath_class(
+        protocol: Literal["abfs", "abfss", "adl", "az"],
+    ) -> type[_AzurePath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["data"]) -> type[_DataPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["github"]) -> type[_GitHubPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["hdfs"]) -> type[_HDFSPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["http", "https"]) -> type[_HTTPPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["file", "local"]) -> type[_FilePath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["memory"]) -> type[_MemoryPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["sftp", "ssh"]) -> type[_SFTPPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["smb"]) -> type[_SMBPath]: ...
+    @overload
+    def get_upath_class(protocol: Literal["webdav"]) -> type[_WebdavPath]: ...
+
+    if sys.platform == "win32":
+
+        @overload
+        def get_upath_class(protocol: Literal[""]) -> type[_WindowsUPath]: ...
+
+    else:
+
+        @overload
+        def get_upath_class(protocol: Literal[""]) -> type[_PosixUPath]: ...  # type: ignore[overload-overlap]  # noqa: E501
+
+    @overload
+    def get_upath_class(
+        protocol: str, *, fallback: bool = ...
+    ) -> type[upath.UPath] | None: ...
+
+
+@lru_cache  # type: ignore[misc] # see: https://github.com/python/typeshed/issues/11280
 def get_upath_class(
     protocol: str,
     *,

upath/tests/implementations/test_github.py

@@ -30,6 +30,16 @@ def path(self):
         path = "github://ap--:universal_pathlib@test_data/data"
         self.path = UPath(path)
 
+    @pytest.fixture(autouse=True)
+    def _xfail_on_rate_limit_errors(self):
+        try:
+            yield
+        except Exception as e:
+            if "rate limit exceeded" in str(e):
+                pytest.xfail("GitHub API rate limit exceeded")
+            else:
+                raise
+
     def test_is_GitHubPath(self):
         """
         Test that the path is a GitHubPath instance.