Skip to content

Commit fc1e498

Browse files
cmp0xffcmp0xff
committed
cat
1 parent a15df83 commit fc1e498

2 files changed

Lines changed: 29 additions & 33 deletions

File tree

‎pandas-stubs/core/arrays/categorical.pyi‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,9 @@ from pandas._typing import (
3737
SequenceNotStr,
3838
np_1darray,
3939
np_1darray_bool,
40-
np_1darray_str,
4140
np_ndarray_anyint,
4241
np_ndarray_float,
42+
np_ndarray_str,
4343
)
4444

4545
from pandas.core.dtypes.dtypes import (
@@ -51,7 +51,7 @@ from pandas.core.dtypes.dtypes import (
5151
class Categorical(NDArrayBackedExtensionArray, Generic[CategoricalValueT]):
5252
__array_priority__: int = ...
5353
@overload
54-
def __new__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
54+
def __new__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
5555
cls,
5656
values: Sequence[Never],
5757
categories: SequenceNotStr[Hashable] | AnyArrayLike | None = None,
@@ -63,10 +63,10 @@ class Categorical(NDArrayBackedExtensionArray, Generic[CategoricalValueT]):
6363
def __new__( # pyright: ignore[reportOverlappingOverload]
6464
cls,
6565
values: (
66-
list[str] | np_1darray_str | SequenceNotStr[str] | Series[str] | Index[str]
66+
list[str] | np_ndarray_str | SequenceNotStr[str] | Series[str] | Index[str]
6767
),
6868
categories: (
69-
SequenceNotStr[str] | Series[str] | Index[str] | np_1darray_str | None
69+
SequenceNotStr[str] | Series[str] | Index[str] | np_ndarray_str | None
7070
) = None,
7171
ordered: bool | None = None,
7272
dtype: CategoricalDtype | None = None,
@@ -102,7 +102,7 @@ class Categorical(NDArrayBackedExtensionArray, Generic[CategoricalValueT]):
102102
SequenceNotStr[CategoricalValueT1]
103103
| Series[CategoricalValueT1]
104104
| Index[CategoricalValueT1]
105-
| np_1darray_str
105+
| np_ndarray_str
106106
| None
107107
) = None,
108108
ordered: bool | None = None,

‎tests/arrays/test_categorical.py‎

Lines changed: 24 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77
import numpy as np
88
import pandas as pd
99
from pandas import Categorical
10-
from pandas.api.typing.aliases import (
11-
Ordered,
12-
)
10+
from pandas.api.typing.aliases import Ordered
1311
from pandas.core.arrays.categorical import CategoricalDtype
1412
from pandas.core.indexes.base import Index
1513

@@ -19,6 +17,7 @@
1917
from tests._typing import (
2018
np_1darray,
2119
np_1darray_bool,
20+
np_ndarray_str,
2221
)
2322

2423

@@ -53,48 +52,45 @@ def test_constructor() -> None:
5352
cat = Categorical(dd)
5453
check(assert_type(cat, "Categorical[str]"), Categorical, str)
5554

56-
values = np.array(["a", "b", "c", "a"])
57-
cat_np = Categorical(values)
58-
# np.array() is typed as ndarray[Any, Any] by numpy stubs, so mypy cannot infer
59-
# the element type; the actual type is Categorical[str]
60-
# TODO: facebook/pyrefly#0
61-
check(assert_type(cat_np, "Categorical[str]"), Categorical) # type: ignore[assert-type] # pyrefly: ignore[assert-type]
55+
values = check(
56+
assert_type(np.array(["a", "b", "c", "a"], np.str_), np_ndarray_str),
57+
np_1darray,
58+
np.str_,
59+
)
60+
check(assert_type(Categorical(values), "Categorical[str]"), Categorical, str)
6261

6362
cat = Categorical(["a", "b", "c"], categories=["a", "b", "c", "d"])
64-
check(assert_type(cat, "Categorical[str]"), Categorical)
63+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
6564

66-
cat = Categorical(["a", "b", "c"], categories=np.array(["a", "b", "c", "d"]))
67-
# TODO: facebook/pyrefly#0
68-
# pyrefly: ignore[assert-type]
69-
check(assert_type(cat, "Categorical[str]"), Categorical)
65+
cat = Categorical(
66+
["a", "b", "c"], categories=np.array(["a", "b", "c", "d"], np.str_)
67+
)
68+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
7069

7170
cat = Categorical(["a", "b", "c"], categories=["a", "b", "c"], ordered=True)
72-
check(assert_type(cat, "Categorical[str]"), Categorical)
71+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
7372

7473
cat = Categorical(["a", "b", "c"], categories=["a", "b", "c"], ordered=False)
75-
check(assert_type(cat, "Categorical[str]"), Categorical)
74+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
7675

7776
cat = Categorical(["a", "b", "c"], categories=["a", "b", "c"], ordered=None)
78-
check(assert_type(cat, "Categorical[str]"), Categorical)
77+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
7978

8079
cat = Categorical(
81-
values=["x", "y", "z", "x"],
82-
categories=["x", "y", "z"],
83-
ordered=True,
84-
copy=True,
80+
values=["x", "y", "z", "x"], categories=["x", "y", "z"], ordered=True, copy=True
8581
)
86-
check(assert_type(cat, "Categorical[str]"), Categorical)
82+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
8783

8884
dtype = pd.CategoricalDtype(categories=["x", "y", "z"], ordered=True)
8985
cat = Categorical(
9086
values=["x", "y", "z", "x"],
9187
dtype=dtype,
9288
copy=True,
9389
)
94-
check(assert_type(cat, "Categorical[str]"), Categorical)
90+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
9591

9692
cat_int = Categorical([1, 2, 3, 1, 2])
97-
check(assert_type(cat_int, "Categorical[int]"), Categorical)
93+
check(assert_type(cat_int, "Categorical[int]"), Categorical, int)
9894

9995
cat_mixed = Categorical(["a", 1, "b", 2])
10096
check(assert_type(cat_mixed, Categorical), Categorical)
@@ -103,19 +99,19 @@ def test_constructor() -> None:
10399
check(assert_type(cat_empty, Categorical), Categorical)
104100

105101
cat = Categorical(["a", "b", "c"], categories=None)
106-
check(assert_type(cat, "Categorical[str]"), Categorical)
102+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
107103

108104
cat1 = Categorical(["a", "b", "c"])
109105
cat = Categorical(cat1)
110-
check(assert_type(cat, "Categorical[str]"), Categorical)
106+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
111107

112108
values_series = pd.Series(["a", "b", "c", "a"])
113109
cat = Categorical(values_series)
114-
check(assert_type(cat, "Categorical[str]"), Categorical)
110+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
115111

116112
values_index = pd.Index(["a", "b", "c", "a"])
117113
cat = Categorical(values_index)
118-
check(assert_type(cat, "Categorical[str]"), Categorical)
114+
check(assert_type(cat, "Categorical[str]"), Categorical, str)
119115

120116

121117
def test_categorical_dtype() -> None:

0 commit comments

Comments
 (0)