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