Implements logic for `partition` and `rpartition`.
(
self,
*,
func: Callable,
dim: Hashable | None,
sep: str | bytes | Any | None,
)
| 2371 | return self._apply(func=func, func_args=(pat,), dtype=np.object_) |
| 2372 | |
| 2373 | def _partitioner( |
| 2374 | self, |
| 2375 | *, |
| 2376 | func: Callable, |
| 2377 | dim: Hashable | None, |
| 2378 | sep: str | bytes | Any | None, |
| 2379 | ) -> T_DataArray: |
| 2380 | """ |
| 2381 | Implements logic for `partition` and `rpartition`. |
| 2382 | """ |
| 2383 | sep = self._stringify(sep) |
| 2384 | |
| 2385 | if dim is None: |
| 2386 | listfunc = lambda x, isep: list(func(x, isep)) |
| 2387 | return self._apply(func=listfunc, func_args=(sep,), dtype=np.object_) |
| 2388 | |
| 2389 | # _apply breaks on an empty array in this case |
| 2390 | if not self._obj.size: |
| 2391 | return self._obj.copy().expand_dims({dim: 0}, axis=-1) |
| 2392 | |
| 2393 | arrfunc = lambda x, isep: np.array(func(x, isep), dtype=self._obj.dtype) |
| 2394 | |
| 2395 | # dtype MUST be object or strings can be truncated |
| 2396 | # See: https://github.com/numpy/numpy/issues/8352 |
| 2397 | return duck_array_ops.astype( |
| 2398 | self._apply( |
| 2399 | func=arrfunc, |
| 2400 | func_args=(sep,), |
| 2401 | dtype=np.object_, |
| 2402 | output_core_dims=[[dim]], |
| 2403 | output_sizes={dim: 3}, |
| 2404 | ), |
| 2405 | self._obj.dtype.kind, |
| 2406 | ) |
| 2407 | |
| 2408 | def partition( |
| 2409 | self, |
no test coverage detected