Upath type narrow on protocol (#431)

* upath: add type overloads to type narrow on protocol

* typesafety: test UPath constructor protocol type narrowing

* typesafety: update tests

* upath: update typing of metaclass

* upath.core: adjust overloads

* upath.implementations.local: force mypy to believe PosixUPath/WindowsUPath are actual UPath subclasses

* upath.registry: adjust overload formatting

* upath.core: restore metaclass.__call__ behavior

Author: ap--

typesafety/test_upath_types.yml

@@ -90,3 +90,58 @@
 
     path_cls = get_upath_class("some-unknown-protocol")
     reveal_type(path_cls)  # N: Revealed type is "type[upath.core.UPath] | None"
+
+- case: upath__new__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
+    - cls_fqn: upath.core.UPath
+      protocol: unknown-protocol
+  main: |
+    import upath
+
+    p = upath.UPath(".", protocol="{{ protocol }}")
+    reveal_type(p)  # N: Revealed type is "{{ cls_fqn }}"
+
+- case: upath__new__empty_protocol
+  disable_cache: true
+  skip: sys.platform == "win32"
+  main: |
+    from upath.core import UPath
+
+    p = UPath("asd", protocol="")
+    reveal_type(p)  # N: Revealed type is "upath.implementations.local.PosixUPath"
+
+- case: get_upath_class_pathlib_win
+  disable_cache: true
+  skip: sys.platform != "win32"
+  main: |
+    from upath.core import UPath
+
+    p = UPath("", protocol="")
+    reveal_type(p)  # N: Revealed type is "upath.implementations.local.WindowsUPath"

upath/core.py

@@ -48,6 +48,8 @@
 from upath.types import WritablePathLike
 
 if TYPE_CHECKING:
+    import upath.implementations as _uimpl
+
     if sys.version_info >= (3, 11):
         from typing import Self
     else:
@@ -56,6 +58,7 @@
     from pydantic import GetCoreSchemaHandler
     from pydantic_core.core_schema import CoreSchema
 
+    _MT = TypeVar("_MT")
     _WT = TypeVar("_WT", bound="WritablePath")
 
 __all__ = ["UPath"]
@@ -109,17 +112,19 @@ class _UPathMeta(ABCMeta):
         def __getitem__(cls, key):
             return cls
 
-    def __call__(cls, *args, **kwargs):
+    def __call__(cls: type[_MT], *args: Any, **kwargs: Any) -> _MT:
         # create a copy if UPath class
         try:
             (arg0,) = args
         except ValueError:
             pass
         else:
             if isinstance(arg0, UPath) and not kwargs:
-                return copy(arg0)
+                return copy(arg0)  # type: ignore[return-value]
+        # We do this call manually, because cls could be a registered
+        # subclass of UPath that is not directly inheriting from UPath.
         inst = cls.__new__(cls, *args, **kwargs)
-        inst.__init__(*args, **kwargs)
+        inst.__init__(*args, **kwargs)  # type: ignore[misc]
         return inst
 
 
@@ -297,9 +302,8 @@ def __new__(
         **storage_options: Any,
     ) -> UPath:
         # narrow type
-        assert issubclass(
-            cls, UPath
-        ), "UPath.__new__ can't instantiate non-UPath classes"
+        if not issubclass(cls, UPath):
+            raise TypeError("UPath.__new__ can't instantiate non-UPath classes")
 
         # deprecate 'scheme'
         if "scheme" in storage_options:
@@ -317,6 +321,7 @@ def __new__(
             storage_options=storage_options,
         )
         # determine which UPath subclass to dispatch to
+        upath_cls: type[UPath] | None
         if cls._protocol_dispatch or cls._protocol_dispatch is None:
             upath_cls = get_upath_class(protocol=pth_protocol)
             if upath_cls is None:
@@ -326,9 +331,12 @@ def __new__(
             # by setting MyUPathSubclass._protocol_dispatch to `False`.
             # This will effectively ignore the registered UPath
             # implementations and return an instance of MyUPathSubclass.
-            # This can be useful if a subclass wants to extend the UPath
+            # This be useful if a subclass wants to extend the UPath
             # api, and it is fine to rely on the default implementation
             # for all supported user protocols.
+            #
+            # THIS IS DEPRECATED!
+            # Use upath.extensions.ProxyUPath to extend the UPath API
             upath_cls = cls
 
         if issubclass(upath_cls, cls):
@@ -438,13 +446,161 @@ class UPath(_UPathMixin, WritablePath, ReadablePath):
         "_relative_base",
     )
 
-    if TYPE_CHECKING:
+    if TYPE_CHECKING:  # noqa: C901
         _chain: Chain
         _chain_parser: FSSpecChainParser
         _fs_cached: bool
         _raw_urlpaths: Sequence[JoinablePathLike]
         _relative_base: str | None
 
+        @overload
+        def __new__(
+            cls,
+        ) -> Self: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["simplecache"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.cached.SimpleCachePath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["gcs", "gs"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.cloud.GCSPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["s3", "s3a"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.cloud.S3Path: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["az", "abfs", "abfss", "adl"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.cloud.AzurePath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["data"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.data.DataPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["github"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.github.GitHubPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["hdfs"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.hdfs.HDFSPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["http", "https"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.http.HTTPPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["file", "local"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.local.FilePath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["memory"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.memory.MemoryPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["sftp", "ssh"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.sftp.SFTPPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["smb"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.smb.SMBPath: ...
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: Literal["webdav"],
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> _uimpl.webdav.WebdavPath: ...
+
+        if sys.platform == "win32":
+
+            @overload  # noqa: E301
+            def __new__(
+                cls,
+                *args: JoinablePathLike,
+                protocol: Literal[""],
+                chain_parser: FSSpecChainParser = ...,
+                **storage_options: Any,
+            ) -> _uimpl.local.WindowsUPath: ...
+
+        else:
+
+            @overload  # noqa: E301
+            def __new__(
+                cls,
+                *args: JoinablePathLike,
+                protocol: Literal[""],
+                chain_parser: FSSpecChainParser = ...,
+                **storage_options: Any,
+            ) -> _uimpl.local.PosixUPath: ...
+
+        @overload  # noqa: E301
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: str | None = ...,
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> Self: ...
+
+        def __new__(
+            cls,
+            *args: JoinablePathLike,
+            protocol: str | None = ...,
+            chain_parser: FSSpecChainParser = ...,
+            **storage_options: Any,
+        ) -> Self: ...
+
     # === JoinablePath attributes =====================================
 
     parser: UPathParser = LazyFlavourDescriptor()  # type: ignore[assignment]

upath/implementations/local.py

@@ -432,38 +432,53 @@ def _copy_from(
 UPath.register(LocalPath)
 
 
-class WindowsUPath(LocalPath, pathlib.WindowsPath):
-    __slots__ = ()
-
-    if os.name != "nt":
-
-        def __new__(
-            cls,
-            *args,
-            protocol: str | None = None,
-            chain_parser: FSSpecChainParser = DEFAULT_CHAIN_PARSER,
-            **storage_options: Any,
-        ) -> WindowsUPath:
-            raise NotImplementedError(
-                f"cannot instantiate {cls.__name__} on your system"
-            )
-
+# Mypy will ignore the ABC.register call above, so we need to force it to
+# think PosixUPath and WindowsUPath are subclasses of UPath.
+# This is really not a good pattern, but it's the best we can do without
+# either introducing a duck-type protocol for UPath or come up with a
+# better solution for the UPath versions of the pathlib.Path subclasses.
 
-class PosixUPath(LocalPath, pathlib.PosixPath):
-    __slots__ = ()
-
-    if os.name == "nt":
+if TYPE_CHECKING:
 
-        def __new__(
-            cls,
-            *args,
-            protocol: str | None = None,
-            chain_parser: FSSpecChainParser = DEFAULT_CHAIN_PARSER,
-            **storage_options: Any,
-        ) -> PosixUPath:
-            raise NotImplementedError(
-                f"cannot instantiate {cls.__name__} on your system"
-            )
+    class WindowsUPath(LocalPath, pathlib.WindowsPath, UPath):  # type: ignore[misc]
+        __slots__ = ()
+
+    class PosixUPath(LocalPath, pathlib.PosixPath, UPath):  # type: ignore[misc]
+        __slots__ = ()
+
+else:
+
+    class WindowsUPath(LocalPath, pathlib.WindowsPath):
+        __slots__ = ()
+
+        if os.name != "nt":
+
+            def __new__(
+                cls,
+                *args,
+                protocol: str | None = None,
+                chain_parser: FSSpecChainParser = DEFAULT_CHAIN_PARSER,
+                **storage_options: Any,
+            ) -> WindowsUPath:
+                raise NotImplementedError(
+                    f"cannot instantiate {cls.__name__} on your system"
+                )
+
+    class PosixUPath(LocalPath, pathlib.PosixPath):
+        __slots__ = ()
+
+        if os.name == "nt":
+
+            def __new__(
+                cls,
+                *args,
+                protocol: str | None = None,
+                chain_parser: FSSpecChainParser = DEFAULT_CHAIN_PARSER,
+                **storage_options: Any,
+            ) -> PosixUPath:
+                raise NotImplementedError(
+                    f"cannot instantiate {cls.__name__} on your system"
+                )
 
 
 class FilePath(UPath):

upath/registry.py

@@ -213,8 +213,7 @@ def get_upath_class(protocol: Literal["simplecache"]) -> type[_SimpleCachePath]:
     def get_upath_class(protocol: Literal["s3", "s3a"]) -> type[_S3Path]: ...
     @overload
     def get_upath_class(protocol: Literal["gcs", "gs"]) -> type[_GCSPath]: ...
-
-    @overload
+    @overload  # noqa: E301
     def get_upath_class(
         protocol: Literal["abfs", "abfss", "adl", "az"],
     ) -> type[_AzurePath]: ...