|
11 | 11 | import pandas.core.algorithms as algos
|
12 | 12 | import pandas.core.nanops as nanops
|
13 | 13 | from pandas.compat import zip
|
14 |
| -from pandas import to_timedelta |
15 |
| -from pandas import to_datetime |
| 14 | +from pandas import to_timedelta, to_datetime |
16 | 15 | import numpy as np
|
17 |
| -from pandas.types.common import (is_datetime64_dtype, is_timedelta64_dtype) |
| 16 | +from pandas.types.common import is_datetime64_dtype, is_timedelta64_dtype |
18 | 17 |
|
19 | 18 |
|
20 | 19 | def cut(x, bins, right=True, labels=None, retbins=False, precision=3,
|
@@ -85,14 +84,7 @@ def cut(x, bins, right=True, labels=None, retbins=False, precision=3,
|
85 | 84 | # NOTE: this binning code is changed a bit from histogram for var(x) == 0
|
86 | 85 | # for handling the cut for datetime and timedelta objects
|
87 | 86 |
|
88 |
| - dtype = None |
89 |
| - if is_timedelta64_dtype(x): |
90 |
| - x = x.view(np.int64) |
91 |
| - dtype = np.timedelta64 |
92 |
| - |
93 |
| - elif is_datetime64_dtype(x): |
94 |
| - x = x.view(np.int64) |
95 |
| - dtype = np.datetime64 |
| 87 | + original, x, dtype = _coerce_to_type(x) |
96 | 88 |
|
97 | 89 | if not np.iterable(bins):
|
98 | 90 | if is_scalar(bins) and bins < 1:
|
@@ -179,14 +171,7 @@ def qcut(x, q, labels=None, retbins=False, precision=3):
|
179 | 171 | >>> pd.qcut(range(5), 4, labels=False)
|
180 | 172 | array([0, 0, 1, 2, 3], dtype=int64)
|
181 | 173 | """
|
182 |
| - dtype = None |
183 |
| - if is_timedelta64_dtype(x): |
184 |
| - x = x.view(np.int64) |
185 |
| - dtype = np.timedelta64 |
186 |
| - |
187 |
| - elif is_datetime64_dtype(x): |
188 |
| - x = x.view(np.int64) |
189 |
| - dtype = np.datetime64 |
| 174 | + original, x, dtype = _coerce_to_type(x) |
190 | 175 |
|
191 | 176 | if is_integer(q):
|
192 | 177 | quantiles = np.linspace(0, 1, q + 1)
|
@@ -329,3 +314,16 @@ def _trim_zeros(x):
|
329 | 314 | if len(x) > 1 and x[-1] == '.':
|
330 | 315 | x = x[:-1]
|
331 | 316 | return x
|
| 317 | + |
| 318 | + |
| 319 | +def _coerce_to_type(x): |
| 320 | + dtype = None |
| 321 | + original = x |
| 322 | + if is_timedelta64_dtype(x): |
| 323 | + x = x.view(np.int64) |
| 324 | + dtype = np.timedelta64 |
| 325 | + |
| 326 | + elif is_datetime64_dtype(x): |
| 327 | + x = x.view(np.int64) |
| 328 | + dtype = np.datetime64 |
| 329 | + return original, x, dtype |
0 commit comments