Skip to content

Commit a9d0eea

Browse files
committed
separate dtype args
1 parent fc9cf76 commit a9d0eea

File tree

1 file changed

+115
-102
lines changed

1 file changed

+115
-102
lines changed

pandas-stubs/_typing.pyi

Lines changed: 115 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -224,33 +224,35 @@ Dtype: TypeAlias = ExtensionDtype | NpDtype
224224
# m
225225
# datetime64
226226

227+
# Builtin bool type and its string alias
228+
BuiltinBooleanDtypeArg: TypeAlias = type[bool] | Literal["bool"]
229+
# Pandas nullable boolean type and its string alias
230+
PandasBooleanDtypeArg: TypeAlias = pd.BooleanDtype | Literal["boolean"]
231+
# Numpy bool type
232+
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bool_
233+
NumpyBooleanDtypeArg: TypeAlias = type[np.bool_] | Literal["?", "b1", "bool_"]
234+
# PyArrow boolean type and its string alias
235+
PyArrowBooleanDtypeArg: TypeAlias = Literal["bool[pyarrow]", "boolean[pyarrow]"]
227236
BooleanDtypeArg: TypeAlias = (
228-
# Builtin bool type and its string alias
229-
type[bool] # noqa: PYI030
230-
| Literal["bool"]
231-
# Pandas nullable boolean type and its string alias
232-
| pd.BooleanDtype
233-
| Literal["boolean"]
234-
# Numpy bool type
235-
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bool_
236-
| type[np.bool_]
237-
| Literal["?", "b1", "bool_"]
238-
# PyArrow boolean type and its string alias
239-
| Literal["bool[pyarrow]", "boolean[pyarrow]"]
237+
BuiltinBooleanDtypeArg
238+
| PandasBooleanDtypeArg
239+
| NumpyBooleanDtypeArg
240+
| PyArrowBooleanDtypeArg
240241
)
241-
IntDtypeArg: TypeAlias = (
242-
# Builtin integer type and its string alias
243-
type[int] # noqa: PYI030
244-
| Literal["int"]
245-
# Pandas nullable integer types and their string aliases
246-
| pd.Int8Dtype
242+
# Builtin integer type and its string alias
243+
BuiltinIntDtypeArg: TypeAlias = type[int] | Literal["int"]
244+
# Pandas nullable integer types and their string aliases
245+
PandasIntDtypeArg: TypeAlias = (
246+
pd.Int8Dtype
247247
| pd.Int16Dtype
248248
| pd.Int32Dtype
249249
| pd.Int64Dtype
250250
| Literal["Int8", "Int16", "Int32", "Int64"]
251-
# Numpy signed integer types and their string aliases
251+
)
252+
# Numpy signed integer types and their string aliases
253+
NumpyIntDtypeArg: TypeAlias = (
252254
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.byte
253-
| type[np.byte]
255+
type[np.byte] # noqa: PYI030
254256
| Literal["b", "i1", "int8", "byte"]
255257
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.short
256258
| type[np.short]
@@ -267,19 +269,26 @@ IntDtypeArg: TypeAlias = (
267269
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.intp
268270
| type[np.intp] # signed pointer (=`intptr_t`, platform dependent)
269271
| Literal["p", "intp"]
270-
# PyArrow integer types and their string aliases
271-
| Literal["int8[pyarrow]", "int16[pyarrow]", "int32[pyarrow]", "int64[pyarrow]"]
272272
)
273-
UIntDtypeArg: TypeAlias = (
274-
# Pandas nullable unsigned integer types and their string aliases
275-
pd.UInt8Dtype # noqa: PYI030
273+
# PyArrow integer types and their string aliases
274+
PyArrowIntDtypeArg: TypeAlias = Literal[
275+
"int8[pyarrow]", "int16[pyarrow]", "int32[pyarrow]", "int64[pyarrow]"
276+
]
277+
IntDtypeArg: TypeAlias = (
278+
BuiltinIntDtypeArg | PandasIntDtypeArg | NumpyIntDtypeArg | PyArrowIntDtypeArg
279+
)
280+
# Pandas nullable unsigned integer types and their string aliases
281+
PandasUIntDtypeArg: TypeAlias = (
282+
pd.UInt8Dtype
276283
| pd.UInt16Dtype
277284
| pd.UInt32Dtype
278285
| pd.UInt64Dtype
279286
| Literal["UInt8", "UInt16", "UInt32", "UInt64"]
280-
# Numpy unsigned integer types and their string aliases
287+
)
288+
# Numpy unsigned integer types and their string aliases
289+
NumpyUIntDtypeArg: TypeAlias = (
281290
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.ubyte
282-
| type[np.ubyte]
291+
type[np.ubyte] # noqa: PYI030
283292
| Literal["B", "u1", "uint8", "ubyte"]
284293
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.ushort
285294
| type[np.ushort]
@@ -296,21 +305,23 @@ UIntDtypeArg: TypeAlias = (
296305
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.uintp
297306
| type[np.uintp] # unsigned pointer (=`uintptr_t`, platform dependent)
298307
| Literal["P", "uintp"]
299-
# PyArrow unsigned integer types and their string aliases
300-
| Literal["uint8[pyarrow]", "uint16[pyarrow]", "uint32[pyarrow]", "uint64[pyarrow]"]
301308
)
302-
FloatDtypeArg: TypeAlias = (
303-
# Builtin float type and its string alias
304-
type[float] # noqa: PYI030
305-
| Literal["float"]
306-
# Pandas nullable float types and their string aliases
307-
| pd.Float32Dtype
308-
| pd.Float64Dtype
309-
| Literal["Float32", "Float64"]
310-
# Numpy float types and their string aliases
309+
# PyArrow unsigned integer types and their string aliases
310+
PyArrowUIntDtypeArg: TypeAlias = Literal[
311+
"uint8[pyarrow]", "uint16[pyarrow]", "uint32[pyarrow]", "uint64[pyarrow]"
312+
]
313+
UIntDtypeArg: TypeAlias = PandasUIntDtypeArg | NumpyUIntDtypeArg | PyArrowUIntDtypeArg
314+
# Builtin float type and its string alias
315+
BuiltinFloatDtypeArg: TypeAlias = type[float] | Literal["float"]
316+
# Pandas nullable float types and their string aliases
317+
PandasFloatDtypeArg: TypeAlias = (
318+
pd.Float32Dtype | pd.Float64Dtype | Literal["Float32", "Float64"]
319+
)
320+
# Numpy float types and their string aliases
321+
NumpyFloatDtypeArg: TypeAlias = (
311322
# NOTE: Alias np.float16 only on Linux x86_64, use np.half instead
312323
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.half
313-
| type[np.half]
324+
type[np.half] # noqa: PYI030
314325
| Literal["e", "f2", "<f2", "float16", "half"]
315326
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.single
316327
| type[np.single]
@@ -321,40 +332,39 @@ FloatDtypeArg: TypeAlias = (
321332
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.longdouble
322333
| type[np.longdouble]
323334
| Literal["g", "f16", "float128", "longdouble", "longfloat"]
324-
# PyArrow floating point types and their string aliases
325-
| Literal[
326-
"float[pyarrow]",
327-
"double[pyarrow]",
328-
"float16[pyarrow]",
329-
"float32[pyarrow]",
330-
"float64[pyarrow]",
331-
]
332335
)
333-
ComplexDtypeArg: TypeAlias = (
334-
# Builtin complex type and its string alias
335-
type[complex] # noqa: PYI030
336-
| Literal["complex"]
337-
# Numpy complex types and their aliases
336+
# PyArrow floating point types and their string aliases
337+
PyArrowFloatDtypeArg: TypeAlias = Literal[
338+
"float[pyarrow]",
339+
"double[pyarrow]",
340+
"float16[pyarrow]",
341+
"float32[pyarrow]",
342+
"float64[pyarrow]",
343+
]
344+
FloatDtypeArg: TypeAlias = (
345+
BuiltinFloatDtypeArg
346+
| PandasFloatDtypeArg
347+
| NumpyFloatDtypeArg
348+
| PyArrowFloatDtypeArg
349+
)
350+
# Builtin complex type and its string alias
351+
BuiltinComplexDtypeArg: TypeAlias = type[complex] | Literal["complex"]
352+
# Numpy complex types and their aliases
353+
NumpyComplexDtypeArg: TypeAlias = (
338354
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.csingle
339-
| type[np.csingle]
355+
type[np.csingle] # noqa: PYI030
340356
| Literal["F", "c8", "complex64", "csingle", "singlecomplex"]
341357
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.cdouble
342358
| type[np.cdouble]
343359
| Literal["D", "c16", "complex128", "cdouble", "cfloat", "complex_"]
344360
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.clongdouble
345361
# NOTE: Alias np.complex256 only on Linux x86_64, use np.clongdouble instead
346362
| type[np.clongdouble]
347-
| Literal[
348-
"G",
349-
"c32",
350-
"complex256",
351-
"clongdouble",
352-
"clongfloat",
353-
"longcomplex",
354-
]
363+
| Literal["G", "c32", "complex256", "clongdouble", "clongfloat", "longcomplex"]
355364
)
365+
ComplexDtypeArg: TypeAlias = BuiltinComplexDtypeArg | NumpyComplexDtypeArg
356366
# Refer to https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
357-
TimedeltaDtypeArg: TypeAlias = Literal[
367+
NumpyTimedeltaDtypeArg: TypeAlias = Literal[
358368
"timedelta64[Y]",
359369
"timedelta64[M]",
360370
"timedelta64[W]",
@@ -399,13 +409,16 @@ TimedeltaDtypeArg: TypeAlias = Literal[
399409
"<m8[ps]",
400410
"<m8[fs]",
401411
"<m8[as]",
402-
# PyArrow duration type and its string alias
412+
]
413+
# PyArrow duration type and its string alias
414+
PyArrowTimedeltaDtypeArg: TypeAlias = Literal[
403415
"duration[s][pyarrow]",
404416
"duration[ms][pyarrow]",
405417
"duration[us][pyarrow]",
406418
"duration[ns][pyarrow]",
407419
]
408-
TimestampDtypeArg: TypeAlias = Literal[
420+
TimedeltaDtypeArg: TypeAlias = NumpyTimedeltaDtypeArg | PyArrowTimedeltaDtypeArg
421+
NumpyTimestampDtypeArg: TypeAlias = Literal[
409422
"datetime64[Y]",
410423
"datetime64[M]",
411424
"datetime64[W]",
@@ -450,58 +463,58 @@ TimestampDtypeArg: TypeAlias = Literal[
450463
"<M8[ps]",
451464
"<M8[fs]",
452465
"<M8[as]",
453-
# PyArrow timestamp type and its string alias
466+
]
467+
# PyArrow timestamp type and its string alias
468+
PyArrowTimestampDtypeArg: TypeAlias = Literal[
454469
"date32[pyarrow]",
455470
"date64[pyarrow]",
456471
"timestamp[s][pyarrow]",
457472
"timestamp[ms][pyarrow]",
458473
"timestamp[us][pyarrow]",
459474
"timestamp[ns][pyarrow]",
460475
]
461-
476+
TimestampDtypeArg: TypeAlias = NumpyTimestampDtypeArg | PyArrowTimestampDtypeArg
477+
# Builtin str type and its string alias
478+
BuiltinStrDtypeArg: TypeAlias = type[str] | Literal["str"]
479+
# Pandas nullable string type and its string alias
480+
PandasStrDtypeArg: TypeAlias = pd.StringDtype | Literal["string"]
481+
# Numpy string type and its string alias
482+
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.str_
483+
NumpyStrDtypeArg: TypeAlias = (
484+
type[np.str_] | Literal["U", "str_", "str0", "unicode", "unicode_"]
485+
)
486+
# PyArrow string type and its string alias
487+
PyArrowStrDtypeArg: TypeAlias = Literal["string[pyarrow]"]
462488
StrDtypeArg: TypeAlias = (
463-
# Builtin str type and its string alias
464-
type[str] # noqa: PYI030
465-
| Literal["str"]
466-
# Pandas nullable string type and its string alias
467-
| pd.StringDtype
468-
| Literal["string"]
469-
# Numpy string type and its string alias
470-
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.str_
471-
| type[np.str_]
472-
| Literal["U", "str_", "str0", "unicode", "unicode_"]
473-
# PyArrow string type and its string alias
474-
| Literal["string[pyarrow]"]
489+
BuiltinStrDtypeArg | PandasStrDtypeArg | NumpyStrDtypeArg | PyArrowStrDtypeArg
490+
)
491+
# Builtin bytes type and its string alias
492+
BuiltinBytesDtypeArg: TypeAlias = type[bytes] | Literal["bytes"]
493+
# Numpy bytes type and its string alias
494+
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_
495+
NumpyBytesDtypeArg: TypeAlias = (
496+
type[np.bytes_] | Literal["S", "bytes_", "bytes0", "string_"]
475497
)
498+
# PyArrow binary type and its string alias
499+
PyArrowBytesDtypeArg: TypeAlias = Literal["binary[pyarrow]"]
476500
BytesDtypeArg: TypeAlias = (
477-
# Builtin bytes type and its string alias
478-
type[bytes] # noqa: PYI030
479-
| Literal["bytes"]
480-
# Numpy bytes type and its string alias
481-
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_
482-
| type[np.bytes_]
483-
| Literal["S", "bytes_", "bytes0", "string_"]
484-
# PyArrow binary type and its string alias
485-
| Literal["binary[pyarrow]"]
501+
BuiltinBytesDtypeArg | NumpyBytesDtypeArg | PyArrowBytesDtypeArg
486502
)
487503
CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"]
488504

489-
ObjectDtypeArg: TypeAlias = (
490-
# Builtin object type and its string alias
491-
type[object] # noqa: PYI030
492-
| Literal["object"]
493-
# Numpy object type and its string alias
494-
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.object_
495-
| type[np.object_]
496-
| Literal["O"] # NOTE: "object_" not assigned
497-
)
505+
# Builtin object type and its string alias
506+
BuiltinObjectDtypeArg: TypeAlias = type[object] | Literal["object"]
507+
# Numpy object type and its string alias
508+
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.object_
509+
# NOTE: "object_" not assigned
510+
NumpyObjectDtypeArg: TypeAlias = type[np.object_] | Literal["O"]
498511

499-
VoidDtypeArg: TypeAlias = (
500-
# Numpy void type and its string alias
501-
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.void
502-
type[np.void]
503-
| Literal["V", "void", "void0"]
504-
)
512+
ObjectDtypeArg: TypeAlias = BuiltinObjectDtypeArg | NumpyObjectDtypeArg
513+
514+
# Numpy void type and its string alias
515+
# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.void
516+
NumpyVoidDtypeArg: TypeAlias = type[np.void] | Literal["V", "void", "void0"]
517+
VoidDtypeArg: TypeAlias = NumpyVoidDtypeArg
505518

506519
# DtypeArg specifies all allowable dtypes in a functions its dtype argument
507520
DtypeArg: TypeAlias = Dtype | Mapping[Hashable, Dtype]

0 commit comments

Comments
 (0)