Skip to content

Commit 55d1aca

Browse files
authored
REF: implement get_unit_from_dtype (#46410)
* REF: implement get_unit_from_dtype * avoid cpplint complaint
1 parent fa7e31b commit 55d1aca

File tree

7 files changed

+89
-1
lines changed

7 files changed

+89
-1
lines changed

pandas/_libs/tslibs/conversion.pyx

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ from pandas._libs.tslibs.np_datetime cimport (
3636
dtstruct_to_dt64,
3737
get_datetime64_unit,
3838
get_datetime64_value,
39+
get_unit_from_dtype,
3940
npy_datetime,
4041
npy_datetimestruct,
4142
pandas_datetime_to_datetimestruct,
@@ -234,7 +235,9 @@ def ensure_datetime64ns(arr: ndarray, copy: bool = True):
234235
result = result.copy()
235236
return result
236237

237-
unit = get_datetime64_unit(arr.flat[0])
238+
if arr.dtype.kind != "M":
239+
raise TypeError("ensure_datetime64ns arr must have datetime64 dtype")
240+
unit = get_unit_from_dtype(arr.dtype)
238241
if unit == NPY_DATETIMEUNIT.NPY_FR_GENERIC:
239242
# without raising explicitly here, we end up with a SystemError
240243
# built-in function ensure_datetime64ns returned a result with an error

pandas/_libs/tslibs/np_datetime.pxd

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
cimport numpy as cnp
12
from cpython.datetime cimport (
23
date,
34
datetime,
@@ -79,3 +80,5 @@ cdef NPY_DATETIMEUNIT get_datetime64_unit(object obj) nogil
7980
cdef int _string_to_dts(str val, npy_datetimestruct* dts,
8081
int* out_local, int* out_tzoffset,
8182
bint want_exc) except? -1
83+
84+
cdef NPY_DATETIMEUNIT get_unit_from_dtype(cnp.dtype dtype)

pandas/_libs/tslibs/np_datetime.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,6 @@
1+
import numpy as np
2+
13
class OutOfBoundsDatetime(ValueError): ...
4+
5+
# only exposed for testing
6+
def py_get_unit_from_dtype(dtype: np.dtype): ...

pandas/_libs/tslibs/np_datetime.pyx

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ from cpython.object cimport (
1919

2020
PyDateTime_IMPORT
2121

22+
cimport numpy as cnp
23+
24+
cnp.import_array()
2225
from numpy cimport int64_t
2326

2427
from pandas._libs.tslibs.util cimport get_c_string_buf_and_size
@@ -42,6 +45,8 @@ cdef extern from "src/datetime/np_datetime.h":
4245

4346
npy_datetimestruct _NS_MIN_DTS, _NS_MAX_DTS
4447

48+
PyArray_DatetimeMetaData get_datetime_metadata_from_dtype(cnp.PyArray_Descr *dtype);
49+
4550
cdef extern from "src/datetime/np_datetime_strings.h":
4651
int parse_iso_8601_datetime(const char *str, int len, int want_exc,
4752
npy_datetimestruct *out,
@@ -74,6 +79,22 @@ cdef inline NPY_DATETIMEUNIT get_datetime64_unit(object obj) nogil:
7479
"""
7580
return <NPY_DATETIMEUNIT>(<PyDatetimeScalarObject*>obj).obmeta.base
7681

82+
83+
cdef NPY_DATETIMEUNIT get_unit_from_dtype(cnp.dtype dtype):
84+
# NB: caller is responsible for ensuring this is *some* datetime64 or
85+
# timedelta64 dtype, otherwise we can segfault
86+
cdef:
87+
cnp.PyArray_Descr* descr = <cnp.PyArray_Descr*>dtype
88+
PyArray_DatetimeMetaData meta
89+
meta = get_datetime_metadata_from_dtype(descr)
90+
return meta.base
91+
92+
93+
def py_get_unit_from_dtype(dtype):
94+
# for testing get_unit_from_dtype; adds 896 bytes to the .so file.
95+
return get_unit_from_dtype(dtype)
96+
97+
7798
# ----------------------------------------------------------------------
7899
# Comparison
79100

pandas/_libs/tslibs/src/datetime/np_datetime.c

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -768,3 +768,15 @@ void pandas_timedelta_to_timedeltastruct(npy_timedelta td,
768768
"invalid base unit");
769769
}
770770
}
771+
772+
773+
/*
774+
* This function returns a pointer to the DateTimeMetaData
775+
* contained within the provided datetime dtype.
776+
*
777+
* Copied near-verbatim from numpy/core/src/multiarray/datetime.c
778+
*/
779+
PyArray_DatetimeMetaData
780+
get_datetime_metadata_from_dtype(PyArray_Descr *dtype) {
781+
return (((PyArray_DatetimeDTypeMetaData *)dtype->c_metadata)->meta);
782+
}

pandas/_libs/tslibs/src/datetime/np_datetime.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,5 +75,12 @@ int cmp_npy_datetimestruct(const npy_datetimestruct *a,
7575
void
7676
add_minutes_to_datetimestruct(npy_datetimestruct *dts, int minutes);
7777

78+
/*
79+
* This function returns the DateTimeMetaData
80+
* contained within the provided datetime dtype.
81+
*/
82+
PyArray_DatetimeMetaData get_datetime_metadata_from_dtype(
83+
PyArray_Descr *dtype);
84+
7885

7986
#endif // PANDAS__LIBS_TSLIBS_SRC_DATETIME_NP_DATETIME_H_
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import numpy as np
2+
3+
from pandas._libs.tslibs.np_datetime import py_get_unit_from_dtype
4+
5+
6+
def test_get_unit_from_dtype():
7+
# datetime64
8+
assert py_get_unit_from_dtype(np.dtype("M8[Y]")) == 0
9+
assert py_get_unit_from_dtype(np.dtype("M8[M]")) == 1
10+
assert py_get_unit_from_dtype(np.dtype("M8[W]")) == 2
11+
# B has been deprecated and removed -> no 3
12+
assert py_get_unit_from_dtype(np.dtype("M8[D]")) == 4
13+
assert py_get_unit_from_dtype(np.dtype("M8[h]")) == 5
14+
assert py_get_unit_from_dtype(np.dtype("M8[m]")) == 6
15+
assert py_get_unit_from_dtype(np.dtype("M8[s]")) == 7
16+
assert py_get_unit_from_dtype(np.dtype("M8[ms]")) == 8
17+
assert py_get_unit_from_dtype(np.dtype("M8[us]")) == 9
18+
assert py_get_unit_from_dtype(np.dtype("M8[ns]")) == 10
19+
assert py_get_unit_from_dtype(np.dtype("M8[ps]")) == 11
20+
assert py_get_unit_from_dtype(np.dtype("M8[fs]")) == 12
21+
assert py_get_unit_from_dtype(np.dtype("M8[as]")) == 13
22+
23+
# timedelta64
24+
assert py_get_unit_from_dtype(np.dtype("m8[Y]")) == 0
25+
assert py_get_unit_from_dtype(np.dtype("m8[M]")) == 1
26+
assert py_get_unit_from_dtype(np.dtype("m8[W]")) == 2
27+
# B has been deprecated and removed -> no 3
28+
assert py_get_unit_from_dtype(np.dtype("m8[D]")) == 4
29+
assert py_get_unit_from_dtype(np.dtype("m8[h]")) == 5
30+
assert py_get_unit_from_dtype(np.dtype("m8[m]")) == 6
31+
assert py_get_unit_from_dtype(np.dtype("m8[s]")) == 7
32+
assert py_get_unit_from_dtype(np.dtype("m8[ms]")) == 8
33+
assert py_get_unit_from_dtype(np.dtype("m8[us]")) == 9
34+
assert py_get_unit_from_dtype(np.dtype("m8[ns]")) == 10
35+
assert py_get_unit_from_dtype(np.dtype("m8[ps]")) == 11
36+
assert py_get_unit_from_dtype(np.dtype("m8[fs]")) == 12
37+
assert py_get_unit_from_dtype(np.dtype("m8[as]")) == 13

0 commit comments

Comments
 (0)