|
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