@@ -1030,9 +1030,13 @@ def test_cut() -> None:
10301030 g = pd .cut (pd .Series ([1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ]), 4 , precision = 1 , duplicates = "drop" )
10311031 h = pd .cut (pd .Series ([1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ]), 4 , labels = False , duplicates = "raise" )
10321032 i = pd .cut (pd .Series ([1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ]), 4 , labels = ["1" , "2" , "3" , "4" ])
1033- check (assert_type (g , pd .Series ), pd .Series )
1034- check (assert_type (h , pd .Series ), pd .Series )
1035- check (assert_type (i , pd .Series ), pd .Series )
1033+ check (assert_type (g , pd .Series [Any , pd .Categorical ]), pd .Series , pd .Interval [float ])
1034+ check (assert_type (h , pd .Series [Any , pd .Categorical ]), pd .Series , pd .Interval [float ])
1035+ check (
1036+ assert_type (i , pd .Series [Any , pd .Categorical ]),
1037+ pd .Series ,
1038+ pd .Interval [float ],
1039+ )
10361040
10371041 j0 , j1 = pd .cut (
10381042 pd .Series ([1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ]),
@@ -1059,13 +1063,13 @@ def test_cut() -> None:
10591063 intval_idx ,
10601064 retbins = True ,
10611065 )
1062- check (assert_type (j0 , pd .Series ), pd .Series )
1066+ check (assert_type (j0 , pd .Series [ Any , pd . Categorical ] ), pd .Series )
10631067 check (assert_type (j1 , npt .NDArray ), np .ndarray )
1064- check (assert_type (k0 , pd .Series ), pd .Series )
1068+ check (assert_type (k0 , pd .Series [ Any , pd . Categorical ] ), pd .Series )
10651069 check (assert_type (k1 , npt .NDArray ), np .ndarray )
1066- check (assert_type (l0 , pd .Series ), pd .Series )
1070+ check (assert_type (l0 , pd .Series [ Any , pd . Categorical ] ), pd .Series )
10671071 check (assert_type (l1 , npt .NDArray ), np .ndarray )
1068- check (assert_type (m0 , pd .Series ), pd .Series )
1072+ check (assert_type (m0 , pd .Series [ Any , pd . Categorical ] ), pd .Series )
10691073 check (assert_type (m1 , pd .IntervalIndex ), pd .IntervalIndex )
10701074
10711075 n0 , n1 = pd .cut ([1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ], intval_idx , retbins = True )
@@ -1076,26 +1080,24 @@ def test_cut() -> None:
10761080 check (
10771081 assert_type (
10781082 pd .cut (s1 , bins = [np .datetime64 ("2020-01-03" ), np .datetime64 ("2020-09-01" )]),
1079- "pd.Series[pd.CategoricalDtype ]" ,
1083+ "pd.Series[Any, pd.Categorical ]" ,
10801084 ),
10811085 pd .Series ,
1086+ pd .Interval [pd .Timestamp ],
10821087 )
10831088 check (
1084- assert_type (
1085- pd .cut (s1 , bins = 10 ),
1086- "pd.Series[pd.CategoricalDtype]" ,
1087- ),
1089+ assert_type (pd .cut (s1 , bins = 10 ), pd .Series [Any , pd .Categorical ]),
10881090 pd .Series ,
10891091 pd .Interval ,
10901092 )
10911093 s0r , s1r = pd .cut (s1 , bins = 10 , retbins = True )
1092- check (assert_type (s0r , pd .Series ), pd .Series , pd .Interval )
1094+ check (assert_type (s0r , pd .Series [ Any , pd . Categorical ] ), pd .Series , pd .Interval )
10931095 check (assert_type (s1r , pd .DatetimeIndex ), pd .DatetimeIndex , pd .Timestamp )
10941096 s0rlf , s1rlf = pd .cut (s1 , bins = 10 , labels = False , retbins = True )
1095- check (assert_type (s0rlf , pd .Series ), pd .Series , np .integer )
1097+ check (assert_type (s0rlf , pd .Series [ Any , pd . Categorical ] ), pd .Series , np .integer )
10961098 check (assert_type (s1rlf , pd .DatetimeIndex ), pd .DatetimeIndex , pd .Timestamp )
10971099 s0rls , s1rls = pd .cut (s1 , bins = 4 , labels = ["1" , "2" , "3" , "4" ], retbins = True )
1098- check (assert_type (s0rls , pd .Series ), pd .Series , str )
1100+ check (assert_type (s0rls , pd .Series [ Any , pd . Categorical ] ), pd .Series , str )
10991101 check (assert_type (s1rls , pd .DatetimeIndex ), pd .DatetimeIndex , pd .Timestamp )
11001102
11011103
0 commit comments