| 289 | |
| 290 | @st.composite |
| 291 | def arrays(draw, type, size=None, nullable=True): |
| 292 | if isinstance(type, st.SearchStrategy): |
| 293 | ty = draw(type) |
| 294 | elif isinstance(type, pa.DataType): |
| 295 | ty = type |
| 296 | else: |
| 297 | raise TypeError('Type must be a pyarrow DataType') |
| 298 | |
| 299 | if isinstance(size, st.SearchStrategy): |
| 300 | size = draw(size) |
| 301 | elif size is None: |
| 302 | size = draw(_default_array_sizes) |
| 303 | elif not isinstance(size, int): |
| 304 | raise TypeError('Size must be an integer') |
| 305 | |
| 306 | if pa.types.is_null(ty): |
| 307 | h.assume(nullable) |
| 308 | value = st.none() |
| 309 | elif pa.types.is_boolean(ty): |
| 310 | value = st.booleans() |
| 311 | elif pa.types.is_integer(ty): |
| 312 | values = draw(npst.arrays(ty.to_pandas_dtype(), shape=(size,))) |
| 313 | return pa.array(values, type=ty) |
| 314 | elif pa.types.is_floating(ty): |
| 315 | values = draw(npst.arrays(ty.to_pandas_dtype(), shape=(size,))) |
| 316 | # Workaround ARROW-4952: no easy way to assert array equality |
| 317 | # in a NaN-tolerant way. |
| 318 | values[np.isnan(values)] = -42.0 |
| 319 | return pa.array(values, type=ty) |
| 320 | elif pa.types.is_decimal(ty): |
| 321 | # TODO(kszucs): properly limit the precision |
| 322 | # value = st.decimals(places=type.scale, allow_infinity=False) |
| 323 | h.reject() |
| 324 | elif pa.types.is_time(ty): |
| 325 | value = st.times() |
| 326 | elif pa.types.is_date(ty): |
| 327 | value = st.dates() |
| 328 | elif pa.types.is_timestamp(ty): |
| 329 | if zoneinfo is None: |
| 330 | pytest.skip('no module named zoneinfo (or tzdata on Windows)') |
| 331 | if ty.tz is None: |
| 332 | pytest.skip('requires timezone not None') |
| 333 | min_int64 = -(2**63) |
| 334 | max_int64 = 2**63 - 1 |
| 335 | min_datetime = datetime.datetime.fromtimestamp( |
| 336 | min_int64 // 10**9) + datetime.timedelta(hours=12) |
| 337 | max_datetime = datetime.datetime.fromtimestamp( |
| 338 | max_int64 // 10**9) - datetime.timedelta(hours=12) |
| 339 | try: |
| 340 | offset = ty.tz.split(":") |
| 341 | offset_hours = int(offset[0]) |
| 342 | offset_min = int(offset[1]) |
| 343 | tz = datetime.timedelta(hours=offset_hours, minutes=offset_min) |
| 344 | except ValueError: |
| 345 | tz = zoneinfo.ZoneInfo(ty.tz) |
| 346 | value = st.datetimes(timezones=st.just(tz), min_value=min_datetime, |
| 347 | max_value=max_datetime) |
| 348 | elif pa.types.is_duration(ty): |