@@ -3239,3 +3239,47 @@ def test_operator_constistency() -> None:
3239
3239
pd .Series ,
3240
3240
pd .Timedelta ,
3241
3241
)
3242
+
3243
+
3244
+ def test_map () -> None :
3245
+ s = pd .Series ([1 , 2 , 3 ])
3246
+
3247
+ mapping = {1 : "a" , 2 : "b" , 3 : "c" }
3248
+ check (
3249
+ assert_type (s .map (mapping , na_action = "ignore" ), "pd.Series[str]" ),
3250
+ pd .Series ,
3251
+ str ,
3252
+ )
3253
+
3254
+ def callable (x : int ) -> str :
3255
+ return str (x )
3256
+
3257
+ check (
3258
+ assert_type (s .map (callable , na_action = "ignore" ), "pd.Series[str]" ),
3259
+ pd .Series ,
3260
+ str ,
3261
+ )
3262
+
3263
+ series = pd .Series (["a" , "b" , "c" ])
3264
+ check (
3265
+ assert_type (s .map (series , na_action = "ignore" ), "pd.Series[str]" ), pd .Series , str
3266
+ )
3267
+
3268
+
3269
+ def test_map_na () -> None :
3270
+ s : pd .Series [int ] = pd .Series ([1 , pd .NA , 3 ])
3271
+
3272
+ mapping = {1 : "a" , 2 : "b" , 3 : "c" }
3273
+ check (assert_type (s .map (mapping , na_action = None ), "pd.Series[str]" ), pd .Series , str )
3274
+
3275
+ def callable (x : int | NAType ) -> str | NAType :
3276
+ if isinstance (x , int ):
3277
+ return str (x )
3278
+ return x
3279
+
3280
+ check (
3281
+ assert_type (s .map (callable , na_action = None ), "pd.Series[str]" ), pd .Series , str
3282
+ )
3283
+
3284
+ series = pd .Series (["a" , "b" , "c" ])
3285
+ check (assert_type (s .map (series , na_action = None ), "pd.Series[str]" ), pd .Series , str )
0 commit comments