Skip to content

Commit b4b6211

Browse files
authored
add support for np.recarray in pd.factorize (#665)
One use case from whithin pandas is to pass `df.to_records` which returns an `np.recarray` directly to `pd.factorize`
1 parent b667761 commit b4b6211

File tree

2 files changed

+5
-1
lines changed

2 files changed

+5
-1
lines changed

pandas-stubs/core/algorithms.pyi

+1-1
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def unique(values: np.ndarray | list) -> np.ndarray: ...
4242
def unique(values: ExtensionArray) -> ExtensionArray: ...
4343
@overload
4444
def factorize(
45-
values: Sequence,
45+
values: Sequence | np.recarray,
4646
sort: bool = ...,
4747
use_na_sentinel: bool = ...,
4848
size_hint: int | None = ...,

tests/test_pandas.py

+4
Original file line numberDiff line numberDiff line change
@@ -756,6 +756,10 @@ def test_factorize() -> None:
756756
check(assert_type(codes, np.ndarray), np.ndarray)
757757
check(assert_type(uniques, np.ndarray), np.ndarray)
758758

759+
codes, uniques = pd.factorize(np.recarray((1,), dtype=[("x", int)]))
760+
check(assert_type(codes, np.ndarray), np.ndarray)
761+
check(assert_type(uniques, np.ndarray), np.ndarray)
762+
759763
codes, cat_uniques = pd.factorize(pd.Categorical(["b", "b", "a", "c", "b"]))
760764
check(assert_type(codes, np.ndarray), np.ndarray)
761765
check(assert_type(cat_uniques, pd.Categorical), pd.Categorical)

0 commit comments

Comments
 (0)