|
5 | 5 | from textwrap import TextWrapper |
6 | 6 | import mmap |
7 | 7 | import time |
| 8 | +from typing import Any, Callable, Dict, Optional, Tuple, Union |
8 | 9 |
|
9 | 10 | import numpy as np |
10 | 11 | from asciitree import BoxStyle, LeftAligned |
11 | 12 | from asciitree.traversal import Traversal |
12 | 13 | from collections.abc import Iterable |
13 | | -from numcodecs.compat import ensure_text, ensure_ndarray_like |
| 14 | +from numcodecs.compat import ( |
| 15 | + ensure_text, |
| 16 | + ensure_ndarray_like, |
| 17 | + ensure_bytes, |
| 18 | + ensure_contiguous_ndarray_like |
| 19 | +) |
| 20 | +from numcodecs.ndarray_like import NDArrayLike |
14 | 21 | from numcodecs.registry import codec_registry |
15 | 22 | from numcodecs.blosc import cbuffer_sizes, cbuffer_metainfo |
16 | 23 |
|
17 | | -from typing import Any, Callable, Dict, Optional, Tuple, Union |
18 | | - |
19 | 24 |
|
20 | 25 | def flatten(arg: Iterable) -> Iterable: |
21 | 26 | for element in arg: |
@@ -696,3 +701,28 @@ def all_equal(value: Any, array: Any): |
696 | 701 | # using == raises warnings from numpy deprecated pattern, but |
697 | 702 | # using np.equal() raises type errors for structured dtypes... |
698 | 703 | return np.all(value == array) |
| 704 | + |
| 705 | + |
| 706 | +def ensure_contiguous_ndarray_or_bytes(buf) -> Union[NDArrayLike, bytes]: |
| 707 | + """Convenience function to coerce `buf` to ndarray-like array or bytes. |
| 708 | +
|
| 709 | + First check if `buf` can be zero-copy converted to a contiguous array. |
| 710 | + If not, `buf` will be copied to a newly allocated `bytes` object. |
| 711 | +
|
| 712 | + Parameters |
| 713 | + ---------- |
| 714 | + buf : ndarray-like, array-like, or bytes-like |
| 715 | + A numpy array like object such as numpy.ndarray, cupy.ndarray, or |
| 716 | + any object exporting a buffer interface. |
| 717 | +
|
| 718 | + Returns |
| 719 | + ------- |
| 720 | + arr : NDArrayLike or bytes |
| 721 | + A ndarray-like or bytes object |
| 722 | + """ |
| 723 | + |
| 724 | + try: |
| 725 | + return ensure_contiguous_ndarray_like(buf) |
| 726 | + except TypeError: |
| 727 | + # An error is raised if `buf` couldn't be zero-copy converted |
| 728 | + return ensure_bytes(buf) |
0 commit comments