_npyio_impl.pyi 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  1. import types
  2. import zipfile
  3. from _typeshed import (
  4. StrOrBytesPath,
  5. StrPath,
  6. SupportsKeysAndGetItem,
  7. SupportsRead,
  8. SupportsWrite,
  9. )
  10. from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence
  11. from re import Pattern
  12. from typing import (
  13. IO,
  14. Any,
  15. ClassVar,
  16. Generic,
  17. Literal as L,
  18. Protocol,
  19. Self,
  20. TypeAlias,
  21. overload,
  22. type_check_only,
  23. )
  24. from typing_extensions import TypeVar, override
  25. import numpy as np
  26. from numpy._core.multiarray import packbits, unpackbits
  27. from numpy._typing import ArrayLike, DTypeLike, NDArray, _DTypeLike, _SupportsArrayFunc
  28. from numpy.ma.mrecords import MaskedRecords
  29. from ._datasource import DataSource as DataSource
  30. __all__ = [
  31. "fromregex",
  32. "genfromtxt",
  33. "load",
  34. "loadtxt",
  35. "packbits",
  36. "save",
  37. "savetxt",
  38. "savez",
  39. "savez_compressed",
  40. "unpackbits",
  41. ]
  42. _T = TypeVar("_T")
  43. _T_co = TypeVar("_T_co", covariant=True)
  44. _ScalarT = TypeVar("_ScalarT", bound=np.generic)
  45. _ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, default=Any, covariant=True)
  46. _FName: TypeAlias = StrPath | Iterable[str] | Iterable[bytes]
  47. _FNameRead: TypeAlias = StrPath | SupportsRead[str] | SupportsRead[bytes]
  48. _FNameWriteBytes: TypeAlias = StrPath | SupportsWrite[bytes]
  49. _FNameWrite: TypeAlias = _FNameWriteBytes | SupportsWrite[str]
  50. @type_check_only
  51. class _SupportsReadSeek(SupportsRead[_T_co], Protocol[_T_co]):
  52. def seek(self, offset: int, whence: int, /) -> object: ...
  53. class BagObj(Generic[_T_co]):
  54. def __init__(self, /, obj: SupportsKeysAndGetItem[str, _T_co]) -> None: ...
  55. def __getattribute__(self, key: str, /) -> _T_co: ...
  56. def __dir__(self) -> list[str]: ...
  57. class NpzFile(Mapping[str, NDArray[_ScalarT_co]]):
  58. _MAX_REPR_ARRAY_COUNT: ClassVar[int] = 5
  59. zip: zipfile.ZipFile | None = None
  60. fid: IO[str] | None = None
  61. files: list[str]
  62. allow_pickle: bool
  63. pickle_kwargs: Mapping[str, Any] | None
  64. f: BagObj[NpzFile[_ScalarT_co]]
  65. #
  66. def __init__(
  67. self,
  68. /,
  69. fid: IO[Any],
  70. own_fid: bool = False,
  71. allow_pickle: bool = False,
  72. pickle_kwargs: Mapping[str, object] | None = None,
  73. *,
  74. max_header_size: int = 10_000,
  75. ) -> None: ...
  76. def __del__(self) -> None: ...
  77. def __enter__(self) -> Self: ...
  78. def __exit__(self, cls: type[BaseException] | None, e: BaseException | None, tb: types.TracebackType | None, /) -> None: ...
  79. @override
  80. def __len__(self) -> int: ...
  81. @override
  82. def __iter__(self) -> Iterator[str]: ...
  83. @override
  84. def __getitem__(self, key: str, /) -> NDArray[_ScalarT_co]: ...
  85. #
  86. @override
  87. @overload
  88. def get(self, key: str, default: None = None, /) -> NDArray[_ScalarT_co] | None: ...
  89. @overload
  90. def get(self, key: str, default: NDArray[_ScalarT_co] | _T, /) -> NDArray[_ScalarT_co] | _T: ... # pyright: ignore[reportIncompatibleMethodOverride]
  91. #
  92. def close(self) -> None: ...
  93. # NOTE: Returns a `NpzFile` if file is a zip file;
  94. # returns an `ndarray`/`memmap` otherwise
  95. def load(
  96. file: StrOrBytesPath | _SupportsReadSeek[bytes],
  97. mmap_mode: L["r+", "r", "w+", "c"] | None = None,
  98. allow_pickle: bool = False,
  99. fix_imports: bool = True,
  100. encoding: L["ASCII", "latin1", "bytes"] = "ASCII",
  101. *,
  102. max_header_size: int = 10_000,
  103. ) -> Any: ...
  104. def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool = True) -> None: ...
  105. def savez(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ...
  106. def savez_compressed(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ...
  107. # File-like objects only have to implement `__iter__` and,
  108. # optionally, `encoding`
  109. @overload
  110. def loadtxt(
  111. fname: _FName,
  112. dtype: None = None,
  113. comments: str | Sequence[str] | None = "#",
  114. delimiter: str | None = None,
  115. converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None,
  116. skiprows: int = 0,
  117. usecols: int | Sequence[int] | None = None,
  118. unpack: bool = False,
  119. ndmin: L[0, 1, 2] = 0,
  120. encoding: str | None = None,
  121. max_rows: int | None = None,
  122. *,
  123. quotechar: str | None = None,
  124. like: _SupportsArrayFunc | None = None,
  125. ) -> NDArray[np.float64]: ...
  126. @overload
  127. def loadtxt(
  128. fname: _FName,
  129. dtype: _DTypeLike[_ScalarT],
  130. comments: str | Sequence[str] | None = "#",
  131. delimiter: str | None = None,
  132. converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None,
  133. skiprows: int = 0,
  134. usecols: int | Sequence[int] | None = None,
  135. unpack: bool = False,
  136. ndmin: L[0, 1, 2] = 0,
  137. encoding: str | None = None,
  138. max_rows: int | None = None,
  139. *,
  140. quotechar: str | None = None,
  141. like: _SupportsArrayFunc | None = None,
  142. ) -> NDArray[_ScalarT]: ...
  143. @overload
  144. def loadtxt(
  145. fname: _FName,
  146. dtype: DTypeLike | None,
  147. comments: str | Sequence[str] | None = "#",
  148. delimiter: str | None = None,
  149. converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None,
  150. skiprows: int = 0,
  151. usecols: int | Sequence[int] | None = None,
  152. unpack: bool = False,
  153. ndmin: L[0, 1, 2] = 0,
  154. encoding: str | None = None,
  155. max_rows: int | None = None,
  156. *,
  157. quotechar: str | None = None,
  158. like: _SupportsArrayFunc | None = None,
  159. ) -> NDArray[Any]: ...
  160. def savetxt(
  161. fname: _FNameWrite,
  162. X: ArrayLike,
  163. fmt: str | Sequence[str] = "%.18e",
  164. delimiter: str = " ",
  165. newline: str = "\n",
  166. header: str = "",
  167. footer: str = "",
  168. comments: str = "# ",
  169. encoding: str | None = None,
  170. ) -> None: ...
  171. @overload
  172. def fromregex(
  173. file: _FNameRead,
  174. regexp: str | bytes | Pattern[Any],
  175. dtype: _DTypeLike[_ScalarT],
  176. encoding: str | None = None,
  177. ) -> NDArray[_ScalarT]: ...
  178. @overload
  179. def fromregex(
  180. file: _FNameRead,
  181. regexp: str | bytes | Pattern[Any],
  182. dtype: DTypeLike | None,
  183. encoding: str | None = None,
  184. ) -> NDArray[Any]: ...
  185. @overload
  186. def genfromtxt(
  187. fname: _FName,
  188. dtype: None = None,
  189. comments: str = "#",
  190. delimiter: str | int | Iterable[int] | None = None,
  191. skip_header: int = 0,
  192. skip_footer: int = 0,
  193. converters: Mapping[int | str, Callable[[str], Any]] | None = None,
  194. missing_values: Any = None,
  195. filling_values: Any = None,
  196. usecols: Sequence[int] | None = None,
  197. names: L[True] | str | Collection[str] | None = None,
  198. excludelist: Sequence[str] | None = None,
  199. deletechars: str = " !#$%&'()*+,-./:;<=>?@[\\]^{|}~",
  200. replace_space: str = "_",
  201. autostrip: bool = False,
  202. case_sensitive: bool | L["upper", "lower"] = True,
  203. defaultfmt: str = "f%i",
  204. unpack: bool | None = None,
  205. usemask: bool = False,
  206. loose: bool = True,
  207. invalid_raise: bool = True,
  208. max_rows: int | None = None,
  209. encoding: str | None = None,
  210. *,
  211. ndmin: L[0, 1, 2] = 0,
  212. like: _SupportsArrayFunc | None = None,
  213. ) -> NDArray[Any]: ...
  214. @overload
  215. def genfromtxt(
  216. fname: _FName,
  217. dtype: _DTypeLike[_ScalarT],
  218. comments: str = "#",
  219. delimiter: str | int | Iterable[int] | None = None,
  220. skip_header: int = 0,
  221. skip_footer: int = 0,
  222. converters: Mapping[int | str, Callable[[str], Any]] | None = None,
  223. missing_values: Any = None,
  224. filling_values: Any = None,
  225. usecols: Sequence[int] | None = None,
  226. names: L[True] | str | Collection[str] | None = None,
  227. excludelist: Sequence[str] | None = None,
  228. deletechars: str = " !#$%&'()*+,-./:;<=>?@[\\]^{|}~",
  229. replace_space: str = "_",
  230. autostrip: bool = False,
  231. case_sensitive: bool | L["upper", "lower"] = True,
  232. defaultfmt: str = "f%i",
  233. unpack: bool | None = None,
  234. usemask: bool = False,
  235. loose: bool = True,
  236. invalid_raise: bool = True,
  237. max_rows: int | None = None,
  238. encoding: str | None = None,
  239. *,
  240. ndmin: L[0, 1, 2] = 0,
  241. like: _SupportsArrayFunc | None = None,
  242. ) -> NDArray[_ScalarT]: ...
  243. @overload
  244. def genfromtxt(
  245. fname: _FName,
  246. dtype: DTypeLike | None,
  247. comments: str = "#",
  248. delimiter: str | int | Iterable[int] | None = None,
  249. skip_header: int = 0,
  250. skip_footer: int = 0,
  251. converters: Mapping[int | str, Callable[[str], Any]] | None = None,
  252. missing_values: Any = None,
  253. filling_values: Any = None,
  254. usecols: Sequence[int] | None = None,
  255. names: L[True] | str | Collection[str] | None = None,
  256. excludelist: Sequence[str] | None = None,
  257. deletechars: str = " !#$%&'()*+,-./:;<=>?@[\\]^{|}~",
  258. replace_space: str = "_",
  259. autostrip: bool = False,
  260. case_sensitive: bool | L["upper", "lower"] = True,
  261. defaultfmt: str = "f%i",
  262. unpack: bool | None = None,
  263. usemask: bool = False,
  264. loose: bool = True,
  265. invalid_raise: bool = True,
  266. max_rows: int | None = None,
  267. encoding: str | None = None,
  268. *,
  269. ndmin: L[0, 1, 2] = 0,
  270. like: _SupportsArrayFunc | None = None,
  271. ) -> NDArray[Any]: ...
  272. @overload
  273. def recfromtxt(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ...
  274. @overload
  275. def recfromtxt(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ...
  276. @overload
  277. def recfromcsv(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ...
  278. @overload
  279. def recfromcsv(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ...