@@ -44,6 +44,12 @@ def is_categorical_dtype(dt):
44
44
return dt == "category"
45
45
46
46
47
+ try :
48
+ from pandas .core .arrays .integer import IntegerDtype
49
+ except ImportError :
50
+ IntegerDtype = ()
51
+
52
+
47
53
def dtype_for_elements_strategy (s ):
48
54
return st .shared (
49
55
s .map (lambda x : pandas .Series ([x ]).dtype ),
@@ -79,6 +85,12 @@ def elements_and_dtype(elements, dtype, source=None):
79
85
f"{ prefix } dtype is categorical, which is currently unsupported"
80
86
)
81
87
88
+ if isinstance (dtype , type ) and issubclass (dtype , IntegerDtype ):
89
+ raise InvalidArgument (
90
+ f"Passed dtype={ dtype !r} is a dtype class, please pass in an instance of this class."
91
+ "Otherwise it would be treated as dtype=object"
92
+ )
93
+
82
94
if isinstance (dtype , type ) and np .dtype (dtype ).kind == "O" and dtype is not object :
83
95
note_deprecation (
84
96
f"Passed dtype={ dtype !r} is not a valid Pandas dtype. We'll treat it as "
@@ -92,13 +104,31 @@ def elements_and_dtype(elements, dtype, source=None):
92
104
f"Passed dtype={ dtype !r} is a strategy, but we require a concrete dtype "
93
105
"here. See https://stackoverflow.com/q/74355937 for workaround patterns."
94
106
)
95
- dtype = try_convert (np .dtype , dtype , "dtype" )
107
+
108
+ pd_dtype_map = {
109
+ t .name : t for t in getattr (IntegerDtype , "__subclasses__" , lambda : [])()
110
+ }
111
+
112
+ dtype = pd_dtype_map .get (dtype , dtype )
113
+
114
+ if isinstance (dtype , IntegerDtype ):
115
+ is_na_dtype = True
116
+ dtype = np .dtype (dtype .name .lower ())
117
+ elif dtype is not None :
118
+ is_na_dtype = False
119
+ dtype = try_convert (np .dtype , dtype , "dtype" )
120
+ else :
121
+ is_na_dtype = False
96
122
97
123
if elements is None :
98
124
elements = npst .from_dtype (dtype )
125
+ if is_na_dtype :
126
+ elements = st .none () | elements
99
127
elif dtype is not None :
100
128
101
129
def convert_element (value ):
130
+ if value is None :
131
+ return None
102
132
name = f"draw({ prefix } elements)"
103
133
try :
104
134
return np .array ([value ], dtype = dtype )[0 ]
@@ -282,9 +312,17 @@ def series(
282
312
else :
283
313
check_strategy (index , "index" )
284
314
285
- elements , dtype = elements_and_dtype (elements , dtype )
315
+ elements , np_dtype = elements_and_dtype (elements , dtype )
286
316
index_strategy = index
287
317
318
+ # if it is converted to an object, use object for series type
319
+ if (
320
+ np_dtype is not None
321
+ and np_dtype .kind == "O"
322
+ and not isinstance (dtype , IntegerDtype )
323
+ ):
324
+ dtype = np_dtype
325
+
288
326
@st .composite
289
327
def result (draw ):
290
328
index = draw (index_strategy )
@@ -293,13 +331,13 @@ def result(draw):
293
331
if dtype is not None :
294
332
result_data = draw (
295
333
npst .arrays (
296
- dtype = dtype ,
334
+ dtype = object ,
297
335
elements = elements ,
298
336
shape = len (index ),
299
337
fill = fill ,
300
338
unique = unique ,
301
339
)
302
- )
340
+ ). tolist ()
303
341
else :
304
342
result_data = list (
305
343
draw (
@@ -310,9 +348,8 @@ def result(draw):
310
348
fill = fill ,
311
349
unique = unique ,
312
350
)
313
- )
351
+ ). tolist ()
314
352
)
315
-
316
353
return pandas .Series (result_data , index = index , dtype = dtype , name = draw (name ))
317
354
else :
318
355
return pandas .Series (
@@ -549,7 +586,7 @@ def row():
549
586
550
587
column_names .add (c .name )
551
588
552
- c .elements , c . dtype = elements_and_dtype (c .elements , c .dtype , label )
589
+ c .elements , _ = elements_and_dtype (c .elements , c .dtype , label )
553
590
554
591
if c .dtype is None and rows is not None :
555
592
raise InvalidArgument (
@@ -589,7 +626,9 @@ def just_draw_columns(draw):
589
626
if columns_without_fill :
590
627
for c in columns_without_fill :
591
628
data [c .name ] = pandas .Series (
592
- np .zeros (shape = len (index ), dtype = c .dtype ), index = index
629
+ np .zeros (shape = len (index ), dtype = object ),
630
+ index = index ,
631
+ dtype = c .dtype ,
593
632
)
594
633
seen = {c .name : set () for c in columns_without_fill if c .unique }
595
634
0 commit comments