Fix type annotations for upath.core, upath.extensions and upath.implementations

upath/core.py

@@ -17,6 +17,7 @@
 from typing import Literal
 from typing import NoReturn
 from typing import TextIO
+from typing import TypeVar
 from typing import overload
 from urllib.parse import SplitResult
 from urllib.parse import urlsplit
@@ -41,6 +42,7 @@
 from upath.types import PathInfo
 from upath.types import ReadablePath
 from upath.types import ReadablePathLike
+from upath.types import SupportsPathLike
 from upath.types import UPathParser
 from upath.types import WritablePath
 from upath.types import WritablePathLike
@@ -54,10 +56,10 @@
     from pydantic import GetCoreSchemaHandler
     from pydantic_core.core_schema import CoreSchema
 
+    _WT = TypeVar("_WT", bound="WritablePath")
 
 __all__ = ["UPath"]
 
-
 _FSSPEC_HAS_WORKING_GLOB = None
 
 
@@ -526,7 +528,7 @@ def parts(self) -> Sequence[str]:
             parts = [*names, drive + sep]
         return tuple(reversed(parts))
 
-    def with_name(self, name) -> Self:
+    def with_name(self, name: str) -> Self:
         """Return a new path with the file name changed."""
         split = self.parser.split
         if self.parser.sep in name:  # `split(name)[0]`
@@ -572,6 +574,21 @@ def parents(self) -> Sequence[Self]:
             return parents
         return super().parents
 
+    def joinpath(self, *pathsegments: JoinablePathLike) -> Self:
+        return self.with_segments(self.__vfspath__(), *pathsegments)
+
+    def __truediv__(self, key: JoinablePathLike) -> Self:
+        try:
+            return self.with_segments(self.__vfspath__(), key)
+        except TypeError:
+            return NotImplemented
+
+    def __rtruediv__(self, key: JoinablePathLike) -> Self:
+        try:
+            return self.with_segments(key, self.__vfspath__())
+        except TypeError:
+            return NotImplemented
+
     # === ReadablePath attributes =====================================
 
     @property
@@ -600,6 +617,32 @@ def __open_reader__(self) -> BinaryIO:
     def readlink(self) -> Self:
         _raise_unsupported(type(self).__name__, "readlink")
 
+    @overload
+    def copy(self, target: _WT, **kwargs: Any) -> _WT: ...
+
+    @overload
+    def copy(self, target: SupportsPathLike | str, **kwargs: Any) -> Self: ...
+
+    def copy(self, target: _WT | SupportsPathLike | str, **kwargs: Any) -> _WT | UPath:
+        if not isinstance(target, UPath):
+            return super().copy(self.with_segments(target), **kwargs)
+        else:
+            return super().copy(target, **kwargs)
+
+    @overload
+    def copy_into(self, target_dir: _WT, **kwargs: Any) -> _WT: ...
+
+    @overload
+    def copy_into(self, target_dir: SupportsPathLike | str, **kwargs: Any) -> Self: ...
+
+    def copy_into(
+        self, target_dir: _WT | SupportsPathLike | str, **kwargs: Any
+    ) -> _WT | UPath:
+        if not isinstance(target_dir, UPath):
+            return super().copy_into(self.with_segments(target_dir), **kwargs)
+        else:
+            return super().copy_into(target_dir, **kwargs)
+
     # --- WritablePath attributes -------------------------------------
 
     def symlink_to(
@@ -715,7 +758,7 @@ def open(
     def stat(
         self,
         *,
-        follow_symlinks=True,
+        follow_symlinks: bool = True,
     ) -> UPathStatResult:
         if not follow_symlinks:
             warnings.warn(
@@ -732,7 +775,7 @@ def lstat(self) -> UPathStatResult:
     def chmod(self, mode: int, *, follow_symlinks: bool = True) -> None:
         _raise_unsupported(type(self).__name__, "chmod")
 
-    def exists(self, *, follow_symlinks=True) -> bool:
+    def exists(self, *, follow_symlinks: bool = True) -> bool:
         return self.fs.exists(self.path)
 
     def is_dir(self) -> bool:
@@ -780,7 +823,7 @@ def glob(
         *,
         case_sensitive: bool = UNSET_DEFAULT,
         recurse_symlinks: bool = UNSET_DEFAULT,
-    ) -> Iterator[UPath]:
+    ) -> Iterator[Self]:
         if case_sensitive is not UNSET_DEFAULT:
             warnings.warn(
                 "UPath.glob(): case_sensitive is currently ignored.",
@@ -808,7 +851,7 @@ def rglob(
         *,
         case_sensitive: bool = UNSET_DEFAULT,
         recurse_symlinks: bool = UNSET_DEFAULT,
-    ) -> Iterator[UPath]:
+    ) -> Iterator[Self]:
         if case_sensitive is not UNSET_DEFAULT:
             warnings.warn(
                 "UPath.glob(): case_sensitive is currently ignored.",
@@ -941,7 +984,7 @@ def resolve(self, strict: bool = False) -> Self:
 
         return self.with_segments(*_parts[:1], *resolved)
 
-    def touch(self, mode=0o666, exist_ok=True) -> None:
+    def touch(self, mode: int = 0o666, exist_ok: bool = True) -> None:
         exists = self.fs.exists(self.path)
         if exists and not exist_ok:
             raise FileExistsError(str(self))

upath/extensions.py

@@ -96,12 +96,15 @@ def __init__(
     def parser(self) -> UPathParser:
         return self.__wrapped__.parser
 
-    def with_segments(self) -> Self:
-        return self._from_upath(self.__wrapped__.with_segments())
+    def with_segments(self, *pathsegments: JoinablePathLike) -> Self:
+        return self._from_upath(self.__wrapped__.with_segments(*pathsegments))
 
     def __str__(self) -> str:
         return self.__wrapped__.__str__()
 
+    def __vfspath__(self) -> str:
+        return self.__wrapped__.__vfspath__()
+
     def __repr__(self) -> str:
         return (
             f"{type(self).__name__}"
@@ -425,10 +428,10 @@ def walk(
         ):
             yield self._from_upath(pth), dirnames, filenames
 
-    def copy(self, target: UPath, **kwargs: Any) -> Self:
+    def copy(self, target: WritablePathLike, **kwargs: Any) -> Self:  # type: ignore[override]  # noqa: E501
         return self._from_upath(self.__wrapped__.copy(target, **kwargs))
 
-    def copy_into(self, target_dir: UPath, **kwargs: Any) -> Self:
+    def copy_into(self, target_dir: WritablePathLike, **kwargs: Any) -> Self:  # type: ignore[override]  # noqa: E501
         return self._from_upath(self.__wrapped__.copy_into(target_dir, **kwargs))
 
     def write_bytes(self, data: bytes) -> int:
@@ -445,8 +448,10 @@ def write_text(
             data, encoding=encoding, errors=errors, newline=newline
         )
 
-    def _copy_from(self, source: ReadablePath, follow_symlinks: bool = True) -> None:
-        self.__wrapped__._copy_from(source, follow_symlinks=follow_symlinks)
+    def _copy_from(
+        self, source: ReadablePath | Self, follow_symlinks: bool = True
+    ) -> None:
+        self.__wrapped__._copy_from(source, follow_symlinks=follow_symlinks)  # type: ignore  # noqa: E501
 
     @property
     def anchor(self) -> str:
@@ -461,7 +466,7 @@ def suffix(self) -> str:
         return self.__wrapped__.suffix
 
     @property
-    def suffixes(self) -> Sequence[str]:
+    def suffixes(self) -> list[str]:
         return self.__wrapped__.suffixes
 
     @property
@@ -474,7 +479,7 @@ def with_stem(self, stem: str) -> Self:
     def with_suffix(self, suffix: str) -> Self:
         return self._from_upath(self.__wrapped__.with_suffix(suffix))
 
-    def joinpath(self, *pathsegments: str) -> Self:
+    def joinpath(self, *pathsegments: JoinablePathLike) -> Self:
         return self._from_upath(self.__wrapped__.joinpath(*pathsegments))
 
     def __truediv__(self, other: str | Self) -> Self:

upath/implementations/cloud.py

@@ -50,7 +50,7 @@ def root(self) -> str:
             return ""
         return self.parser.sep
 
-    def __vfspath__(self):
+    def __vfspath__(self) -> str:
         path = super().__vfspath__()
         if self._relative_base is None:
             drive = self.parser.splitdrive(path)[0]
@@ -93,7 +93,7 @@ def mkdir(
             if "unexpected keyword argument 'create_parents'" in str(err):
                 self.fs.mkdir(self.path)
 
-    def exists(self, *, follow_symlinks=True):
+    def exists(self, *, follow_symlinks: bool = True) -> bool:
         # required for gcsfs<2025.5.0, see: https://github.com/fsspec/gcsfs/pull/676
         path = self.path
         if len(path) > 1:

upath/implementations/data.py

@@ -1,28 +1,45 @@
 from __future__ import annotations
 
+import sys
+from collections.abc import Sequence
+
 from upath.core import UPath
+from upath.types import JoinablePathLike
+
+if sys.version_info > (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
 
 
 class DataPath(UPath):
 
     @property
-    def parts(self):
+    def parts(self) -> Sequence[str]:
         return (self.path,)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.parser.join(*self._raw_urlpaths)
 
-    def with_segments(self, *pathsegments):
+    def with_segments(self, *pathsegments: JoinablePathLike) -> Self:
         raise NotImplementedError("path operation not supported by DataPath")
 
-    def with_suffix(self, suffix: str):
+    def with_suffix(self, suffix: str) -> Self:
         raise NotImplementedError("path operation not supported by DataPath")
 
-    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
+    def mkdir(
+        self, mode: int = 0o777, parents: bool = False, exist_ok: bool = False
+    ) -> None:
         raise FileExistsError(str(self))
 
-    def write_bytes(self, data):
+    def write_bytes(self, data: bytes) -> int:
         raise NotImplementedError("DataPath does not support writing")
 
-    def write_text(self, data, **kwargs):
+    def write_text(
+        self,
+        data: str,
+        encoding: str | None = ...,
+        errors: str | None = ...,
+        newline: str | None = ...,
+    ) -> int:
         raise NotImplementedError("DataPath does not support writing")

upath/implementations/github.py

@@ -2,10 +2,17 @@
 GitHub file system implementation
 """
 
+import sys
+from collections.abc import Iterator
 from collections.abc import Sequence
 
 import upath.core
 
+if sys.version_info > (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 
 class GitHubPath(upath.core.UPath):
     """
@@ -19,7 +26,7 @@ def path(self) -> str:
             return ""
         return pth
 
-    def iterdir(self):
+    def iterdir(self) -> Iterator[Self]:
         if self.is_file():
             raise NotADirectoryError(str(self))
         yield from super().iterdir()

upath/implementations/hdfs.py

@@ -1,19 +1,29 @@
 from __future__ import annotations
 
+import sys
+from collections.abc import Iterator
+
 from upath.core import UPath
 
+if sys.version_info > (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 __all__ = ["HDFSPath"]
 
 
 class HDFSPath(UPath):
     __slots__ = ()
 
-    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
+    def mkdir(
+        self, mode: int = 0o777, parents: bool = False, exist_ok: bool = False
+    ) -> None:
         if not exist_ok and self.exists():
             raise FileExistsError(str(self))
         super().mkdir(mode=mode, parents=parents, exist_ok=exist_ok)
 
-    def iterdir(self):
+    def iterdir(self) -> Iterator[Self]:
         if self.is_file():
             raise NotADirectoryError(str(self))
         yield from super().iterdir()

upath/implementations/memory.py

@@ -1,22 +1,30 @@
 from __future__ import annotations
 
+import sys
+from collections.abc import Iterator
+
 from upath.core import UPath
 
+if sys.version_info > (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 __all__ = ["MemoryPath"]
 
 
 class MemoryPath(UPath):
-    def iterdir(self):
+    def iterdir(self) -> Iterator[Self]:
         if not self.is_dir():
             raise NotADirectoryError(str(self))
         yield from super().iterdir()
 
     @property
-    def path(self):
+    def path(self) -> str:
         path = super().path
         return "/" if path == "." else path
 
-    def __str__(self):
+    def __str__(self) -> str:
         s = super().__str__()
         if s.startswith("memory:///"):
             s = s.replace("memory:///", "memory://", 1)

upath/implementations/smb.py

@@ -2,6 +2,7 @@
 
 import sys
 import warnings
+from collections.abc import Iterator
 from typing import TYPE_CHECKING
 from typing import Any
 
@@ -21,7 +22,12 @@
 class SMBPath(UPath):
     __slots__ = ()
 
-    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
+    def mkdir(
+        self,
+        mode: int = 0o777,
+        parents: bool = False,
+        exist_ok: bool = False,
+    ) -> None:
         # smbclient does not support setting mode externally
         if parents and not exist_ok and self.exists():
             raise FileExistsError(str(self))
@@ -36,7 +42,7 @@ def mkdir(self, mode=0o777, parents=False, exist_ok=False):
             if not self.is_dir():
                 raise FileExistsError(str(self))
 
-    def iterdir(self):
+    def iterdir(self) -> Iterator[Self]:
         if not self.is_dir():
             raise NotADirectoryError(str(self))
         else: