|
1 | 1 | """ |
2 | 2 | Testing that we work in the downstream packages |
3 | 3 | """ |
| 4 | +import array |
4 | 5 | import importlib |
5 | 6 | import subprocess |
6 | 7 | import sys |
|
14 | 15 | import pandas as pd |
15 | 16 | from pandas import ( |
16 | 17 | DataFrame, |
| 18 | + DatetimeIndex, |
17 | 19 | Series, |
| 20 | + TimedeltaIndex, |
18 | 21 | ) |
19 | 22 | import pandas._testing as tm |
| 23 | +from pandas.core.arrays import ( |
| 24 | + DatetimeArray, |
| 25 | + TimedeltaArray, |
| 26 | +) |
| 27 | +from pandas.core.arrays.datetimes import _sequence_to_dt64ns |
| 28 | +from pandas.core.arrays.timedeltas import sequence_to_td64ns |
20 | 29 |
|
21 | 30 |
|
22 | 31 | def import_module(name): |
@@ -277,3 +286,70 @@ def __radd__(self, other): |
277 | 286 |
|
278 | 287 | assert right.__add__(left) is NotImplemented |
279 | 288 | assert right + left is left |
| 289 | + |
| 290 | + |
| 291 | +@pytest.fixture( |
| 292 | + params=[ |
| 293 | + "memoryview", |
| 294 | + "array", |
| 295 | + pytest.param("dask", marks=td.skip_if_no("dask.array")), |
| 296 | + pytest.param("xarray", marks=td.skip_if_no("xarray")), |
| 297 | + ] |
| 298 | +) |
| 299 | +def array_likes(request): |
| 300 | + """ |
| 301 | + Fixture giving a numpy array and a parametrized 'data' object, which can |
| 302 | + be a memoryview, array, dask or xarray object created from the numpy array. |
| 303 | + """ |
| 304 | + # GH#24539 recognize e.g xarray, dask, ... |
| 305 | + arr = np.array([1, 2, 3], dtype=np.int64) |
| 306 | + |
| 307 | + name = request.param |
| 308 | + if name == "memoryview": |
| 309 | + data = memoryview(arr) |
| 310 | + elif name == "array": |
| 311 | + data = array.array("i", arr) |
| 312 | + elif name == "dask": |
| 313 | + import dask.array |
| 314 | + |
| 315 | + data = dask.array.array(arr) |
| 316 | + elif name == "xarray": |
| 317 | + import xarray as xr |
| 318 | + |
| 319 | + data = xr.DataArray(arr) |
| 320 | + |
| 321 | + return arr, data |
| 322 | + |
| 323 | + |
| 324 | +@pytest.mark.parametrize("dtype", ["M8[ns]", "m8[ns]"]) |
| 325 | +def test_from_obscure_array(dtype, array_likes): |
| 326 | + # GH#24539 recognize e.g xarray, dask, ... |
| 327 | + # Note: we dont do this for PeriodArray bc _from_sequence won't accept |
| 328 | + # an array of integers |
| 329 | + # TODO: could check with arraylike of Period objects |
| 330 | + arr, data = array_likes |
| 331 | + |
| 332 | + cls = {"M8[ns]": DatetimeArray, "m8[ns]": TimedeltaArray}[dtype] |
| 333 | + |
| 334 | + expected = cls(arr) |
| 335 | + result = cls._from_sequence(data) |
| 336 | + tm.assert_extension_array_equal(result, expected) |
| 337 | + |
| 338 | + func = {"M8[ns]": _sequence_to_dt64ns, "m8[ns]": sequence_to_td64ns}[dtype] |
| 339 | + result = func(arr)[0] |
| 340 | + expected = func(data)[0] |
| 341 | + tm.assert_equal(result, expected) |
| 342 | + |
| 343 | + if not isinstance(data, memoryview): |
| 344 | + # FIXME(GH#44431) these raise on memoryview and attempted fix |
| 345 | + # fails on py3.10 |
| 346 | + func = {"M8[ns]": pd.to_datetime, "m8[ns]": pd.to_timedelta}[dtype] |
| 347 | + result = func(arr).array |
| 348 | + expected = func(data).array |
| 349 | + tm.assert_equal(result, expected) |
| 350 | + |
| 351 | + # Let's check the Indexes while we're here |
| 352 | + idx_cls = {"M8[ns]": DatetimeIndex, "m8[ns]": TimedeltaIndex}[dtype] |
| 353 | + result = idx_cls(arr) |
| 354 | + expected = idx_cls(data) |
| 355 | + tm.assert_index_equal(result, expected) |
0 commit comments