diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/INSTALLER b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..5c69047b2eb8235994febeeae1da4a82365a240a --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/INSTALLER @@ -0,0 +1 @@ +uv \ No newline at end of file diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/METADATA b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..776376f5fdd8a31521a503f0061b507a603e03c8 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/METADATA @@ -0,0 +1,171 @@ +Metadata-Version: 2.4 +Name: zarr +Version: 3.1.6 +Summary: An implementation of chunked, compressed, N-dimensional arrays for Python +Project-URL: issues, https://github.com/zarr-developers/zarr-python/issues +Project-URL: changelog, https://zarr.readthedocs.io/en/stable/release-notes.html +Project-URL: Discussions, https://github.com/zarr-developers/zarr-python/discussions +Project-URL: documentation, https://zarr.readthedocs.io/ +Project-URL: homepage, https://github.com/zarr-developers/zarr-python +Author-email: Alistair Miles +Maintainer: jakirkham, Norman Rzepka, Ryan Abernathey, David Stansby, Deepak Cherian +Maintainer-email: Davis Bennett , Josh Moore , Joe Hamman , Juan Nunez-Iglesias , Martin Durant , Tom Augspurger +License-Expression: MIT +License-File: LICENSE.txt +Keywords: Python,compressed,ndimensional-arrays,zarr +Classifier: Development Status :: 6 - Mature +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Information Technology +Classifier: Intended Audience :: Science/Research +Classifier: Operating System :: Unix +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Requires-Python: >=3.11 +Requires-Dist: donfig>=0.8 +Requires-Dist: google-crc32c>=1.5 +Requires-Dist: numcodecs>=0.14 +Requires-Dist: numpy>=2.0 +Requires-Dist: packaging>=22.0 +Requires-Dist: typing-extensions>=4.12 +Provides-Extra: cli +Requires-Dist: typer; extra == 'cli' +Provides-Extra: gpu +Requires-Dist: cupy-cuda12x; extra == 'gpu' +Provides-Extra: optional +Requires-Dist: universal-pathlib; extra == 'optional' +Provides-Extra: remote +Requires-Dist: fsspec>=2023.10.0; extra == 'remote' +Requires-Dist: obstore>=0.5.1; extra == 'remote' +Description-Content-Type: text/markdown + +
+
+
+ +# Zarr + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Latest Release + + latest release + +
+ + latest release + +
Package Status + + status + +
License + + license + +
Build Status + + build status + +
Pre-commit Status + + pre-commit status + +
Coverage + + coverage + +
Downloads + + pypi downloads + +
Developer Chat + + + +
Funding + + CZI's Essential Open Source Software for Science + +
Citation + + DOI + +
+ +## What is it? + +Zarr is a Python package providing an implementation of compressed, chunked, N-dimensional arrays, designed for use in parallel computing. See the [documentation](https://zarr.readthedocs.io) for more information. + +## Main Features + +- [**Create**](https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#creating-an-array) N-dimensional arrays with any NumPy `dtype`. +- [**Chunk arrays**](https://zarr.readthedocs.io/en/stable/user-guide/performance.html#chunk-optimizations) along any dimension. +- [**Compress**](https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#compressors) and/or filter chunks using any NumCodecs codec. +- [**Store arrays**](https://zarr.readthedocs.io/en/stable/user-guide/storage.html) in memory, on disk, inside a zip file, on S3, etc... +- [**Read**](https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#reading-and-writing-data) an array [**concurrently**](https://zarr.readthedocs.io/en/stable/user-guide/performance.html#parallel-computing-and-synchronization) from multiple threads or processes. +- [**Write**](https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#reading-and-writing-data) to an array concurrently from multiple threads or processes. +- Organize arrays into hierarchies via [**groups**](https://zarr.readthedocs.io/en/stable/quickstart.html#hierarchical-groups). + +## Where to get it + +Zarr can be installed from PyPI using `pip`: + +```bash +pip install zarr +``` + +or via `conda`: + +```bash +conda install -c conda-forge zarr +``` + +For more details, including how to install from source, see the [installation documentation](https://zarr.readthedocs.io/en/stable/index.html#installation). diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/RECORD b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..554159cb9c621b0d0df8052e5c01c4186e46a51b --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/RECORD @@ -0,0 +1,103 @@ +../../../bin/zarr,sha256=PG0LRLGAhLssz-2QBC3m0yp4HOpbKfCfrnf235ik8hE,342 +zarr-3.1.6.dist-info/INSTALLER,sha256=5hhM4Q4mYTT9z6QB6PGpUAW81PGNFrYrdXMj4oM_6ak,2 +zarr-3.1.6.dist-info/METADATA,sha256=-diAtnIVxk9WPannPXucRNDlga_53MNQVUfW6VAPRUM,8608 +zarr-3.1.6.dist-info/RECORD,, +zarr-3.1.6.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr-3.1.6.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87 +zarr-3.1.6.dist-info/entry_points.txt,sha256=87EDiIFOzr9bpH-HNPSQRDKM-p9unAZrBvybRDNCEcw,75 +zarr-3.1.6.dist-info/licenses/LICENSE.txt,sha256=HXdW8z09og1TsWLC65vrqR7g3BrrRllyi23I8vdZQcs,1124 +zarr/__init__.py,sha256=s8jh56_lxQLeevIDxX93X5wsmTOrx17NWoQ6rjQd24A,4270 +zarr/_cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/_cli/cli.py,sha256=uTI0eBKOkh8Lanr1DyX2HxgcWeVbtpVnSrYGis5Hg2E,5852 +zarr/_compat.py,sha256=a6rDJz6Z7R4wcFgfkOw8BmgPb20CwO0tfE7fU1h4d6c,3589 +zarr/_version.py,sha256=LCICTIfLZ__LkFNkrPDD-bSal5fP2QP0iQ5beTR8e2Y,704 +zarr/abc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/abc/buffer.py,sha256=EHVr5Dwg759DzFEqwN2diAaNeUcXlF68lvyFBbaCJbA,196 +zarr/abc/codec.py,sha256=GI34j_VwNVPMaIkyLqWryhKMPKDfD7-D-SkIkrDw5Qw,16290 +zarr/abc/metadata.py,sha256=c1IkPTX3rEO_5pY9T05OH6uZbnhVvGHUV--GbMxpzoM,1405 +zarr/abc/numcodec.py,sha256=pJbdF-NqsJZEu_TvuEZj3QJManR6wXNRCaVOTfLBa40,2814 +zarr/abc/store.py,sha256=MkpqzxU7zb7cxrBAj_JGBX0lq-72k1Hrg0ajzaILYTo,23663 +zarr/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/api/asynchronous.py,sha256=CFNu9K7_tfs0HXdJz-ApqrNNVZ1069qNRiSKulnVNb0,49204 +zarr/api/synchronous.py,sha256=1MHbYAq6rGbThqa5HWBhCGZlwx34gD1PIi-CGEh3NFM,54557 +zarr/buffer/__init__.py,sha256=qJ_P2iFiWwwW1EnNJzJnRlv0JWgZQR09OCYDOeOr6E4,283 +zarr/buffer/cpu.py,sha256=ZMXfM3bgxk_Dar9R9BaNtt2oaF9D34p2_Bz_CHsgDYE,270 +zarr/buffer/gpu.py,sha256=6CjqB7BFb4oG1QkqF1O2VF8lT3Set2i4PDPhgbWYXpw,137 +zarr/codecs/__init__.py,sha256=QhbUyxk17cwe9Ag7Xc70xSuW35YyYrvG9QJ_idfrs78,3410 +zarr/codecs/_v2.py,sha256=fwymGyufsW5dd6Ba2y_EF0oLRLSfw50Yy1p9wy0lzZE,3649 +zarr/codecs/blosc.py,sha256=64fw9-OFH94Mcv8rNeymocRxTY5AMMkQGlerEshjyaY,11074 +zarr/codecs/bytes.py,sha256=d7iMb0G4hb71Ll2H4S3PgjMwSw8Sdy1v_ssXhTgGenU,4183 +zarr/codecs/crc32c_.py,sha256=2c5QCjLSbEgvbU3cWSrgM2HH9vxGkuHYL0afpsZzbg0,2572 +zarr/codecs/gzip.py,sha256=1TjO91IVemCIfQ5BwMexzcNcvXAJmQXzfI0Qnw_FZhM,2503 +zarr/codecs/numcodecs/__init__.py,sha256=9E4nPdDj10lgCEsqNlcfAg8X4MPmVk_VsnfOBM8d0bI,901 +zarr/codecs/numcodecs/_codecs.py,sha256=P13EXUHc-o0Kqoj2zLuyrS8Dm-czpR1pXA1yn5BIZJY,11953 +zarr/codecs/sharding.py,sha256=pDxta4FW4laLjsDDHwQS5E-rKf38H8msi4bkrLJzgZo,28162 +zarr/codecs/transpose.py,sha256=ZluSDxmHBzIySLOvSmvcusqGTKHO-lUkiErrJVdEe3I,4463 +zarr/codecs/vlen_utf8.py,sha256=CoOfIqma7THsLQOnlFSh60ulvk2vYDIMfe8sgnVa-KA,4536 +zarr/codecs/zstd.py,sha256=groE9suVC8Imuzg5eprg4HXIfCwGp96s40NvXG9TK6Y,3324 +zarr/convenience.py,sha256=ysKzQGUWaSey7I9LJMYhK3NSgJW0iRBoF8RI_FV0uYQ,800 +zarr/core/__init__.py,sha256=vPyRTNVnQXIblFLBgdbcs9CBGsL6xa3RmEGrQbHVXQY,283 +zarr/core/_info.py,sha256=yPWnlqPSryJXy2tCnyxiYt7spG4oSIuJJIMmCNw09eo,5015 +zarr/core/_tree.py,sha256=WTv5vafZOTu4YIzEJe_WtSH7e59nx6zPfvrxsVnmSLU,5074 +zarr/core/array.py,sha256=ff-XPTnzvGK3JNcn4uWf3MPtYzfInDY6Sncl9q-JTLM,222706 +zarr/core/array_spec.py,sha256=URY0qzEz13wxoYSEfER8ofUUFqqZGa0f8ajJLI6oz5s,3855 +zarr/core/attributes.py,sha256=wRNaIjvLMW8tgh1iw4dB1s08xvgEF4DMhW0a9HR81CE,1687 +zarr/core/buffer/__init__.py,sha256=jnvcnaKdaWfOpmNPOdf4adFeYnh2gTaK-dFI0dVzJxc,421 +zarr/core/buffer/core.py,sha256=IJxnBEY8RJsT6BxQ5LvyiXIhIsuym3MJKQuMqk3XzUs,18024 +zarr/core/buffer/cpu.py,sha256=k2uLEcE10_bb3UPpOp_AGvA4aPCDT-17luF5zDr69o4,7690 +zarr/core/buffer/gpu.py,sha256=4hCdSBtDzLWyyXeTehHSEqqjqQyrk6NAVOO6hT0D48Q,7777 +zarr/core/chunk_grids.py,sha256=H0bwdTOsQH5okmpBZEE1yS-hhZ4FVV8I9nOnPoBdo7U,10655 +zarr/core/chunk_key_encodings.py,sha256=s0NDx4bAjWHRE_Epk3EPJoZhQkHomasJNcrmX-o-dd0,4379 +zarr/core/codec_pipeline.py,sha256=aYadeiOGNNfhIbSpiAKp9FbaPqlODr7geQYdWI-y0Ng,22003 +zarr/core/common.py,sha256=AFi00SCRtsfRSxsqmjUWi9Z4s7NzYX8U5NO7qE8ikj4,7834 +zarr/core/config.py,sha256=lWPBGQm4xrwQd9Woye7Ok0RRE-mjdfLARvE4KSWD1m4,6283 +zarr/core/dtype/__init__.py,sha256=tO6uM5pxHie1tuEosYNRm5M9lMMi29gJeSQznoj0Dfc,9111 +zarr/core/dtype/common.py,sha256=HIbGlU1Ja_3G7FcF3wuNuJkebu_Bn9seLRWKRCehukc,8651 +zarr/core/dtype/npy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/core/dtype/npy/bool.py,sha256=HRUqa-DtUpTAm9nZcrOg7ikQ-tNy3AwHHn_yPbVy41c,8908 +zarr/core/dtype/npy/bytes.py,sha256=7h_uA75HyX7uJ5frcgZMeeNsjoHnreUMvpElFVNDE9c,39979 +zarr/core/dtype/npy/common.py,sha256=-mQSr44DuJ-Emp_enxPNIq3-E37Ia6Mj8WMHlPgwAu0,14314 +zarr/core/dtype/npy/complex.py,sha256=o0-2hlqQx7rB2UhvaVj0JdawEw7QUWCYFXJlTXIhVL4,12928 +zarr/core/dtype/npy/float.py,sha256=BCBZj9oLVjwpYZT3pb2aA-sBeWyfffHzOcXw_Xgu3GU,13631 +zarr/core/dtype/npy/int.py,sha256=PbUz-OEQkdCb6X_K4soF1wVYvU7BWLF6mhYsTgOdCFw,47134 +zarr/core/dtype/npy/string.py,sha256=IYsgr8asN8yYG_YOFbtLgk11iV8O27bZKG3dD4ZdkqM,24952 +zarr/core/dtype/npy/structured.py,sha256=y2k6fJd8G9iVq6oF0neyBeygCazdWAS-8Z01ayI-Z0Q,15332 +zarr/core/dtype/npy/time.py,sha256=8tTA2jMtLPMN1BHuibUjdjUmnboQ-4Z7Hy2r4O7iw8U,27921 +zarr/core/dtype/registry.py,sha256=m-edO29HLCAdPdWpAy_Eph5L_YTS0kMsw2E2gk17e2E,7396 +zarr/core/dtype/wrapper.py,sha256=xJvjMDWm0T_ArG3eD_qnzZ0IyYOm5m038EzovHu6-OU,9724 +zarr/core/group.py,sha256=vLoumjmvibKg-0nFDQ_yiNzHWcHbSqrsGgHKMlCNzEI,150283 +zarr/core/indexing.py,sha256=TynMb87bWmL6XJ31UzPn6PmpNitJxIyeL7U00yY-ST8,60057 +zarr/core/metadata/__init__.py,sha256=ej8tBNY67cbFAQdjQo8akiyFgdnNFXtN5SzaRD0cDk4,539 +zarr/core/metadata/common.py,sha256=r2w276Hnok9EWqb2D8uWKMqq7PhUKhc7PvQOX3W44SE,255 +zarr/core/metadata/io.py,sha256=2EiPzMtqv-AD5xEIiZqepMuV34yiqR3vxSdNESXHxro,3194 +zarr/core/metadata/v2.py,sha256=PJk-Tb2pl2cawRMbzNaHoA779LLAWKVWpe0FrKkxVy0,12557 +zarr/core/metadata/v3.py,sha256=wIEDF621hb9brc_aSCJH8UVPwWripHK265n8duG2RZ8,17250 +zarr/core/sync.py,sha256=YW6RRuupot1GbOBA7hXDfp3SsFRKyI8cKcTY2koPSSk,7319 +zarr/core/sync_group.py,sha256=EC-s7osWK6Csj0yCTwaTV814pwi92XLm-LD5JYnMAyY,6063 +zarr/creation.py,sha256=qkqT4ibcoE1ovVLvOD_m7ja_TVKJLhPsMc169M0A4Zo,802 +zarr/dtype.py,sha256=6UirY7ewXg4Oa7BSLKuVB_klwrUWI3EY7AKkzGaIH8Q,1828 +zarr/errors.py,sha256=9rGz46pr1VMQKCkWv65MRc2h5M0YlxWriHMEg17OlYU,3679 +zarr/experimental/__init__.py,sha256=d5z1ssdeTxAhWIUFQ9c5vWSd4FlBJBDVeLhAGfXegZE,89 +zarr/experimental/cache_store.py,sha256=_7rqe82js0ZZHZZoZXTMxAmdHeyGcyCxuywjIYbjw68,17282 +zarr/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/metadata/migrate_v3.py,sha256=5IqCWcqDZQmNbRKdeTOzyyRY51uu9wqXXsx4DcEqVAY,10453 +zarr/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +zarr/registry.py,sha256=rbJnN7QZL-Vb1pQTyjGAFEDTTkVzQPdi1pKgFJhvZmI,11490 +zarr/storage/__init__.py,sha256=u3WB2W2ufq-zveezvceyfMBtlbvsO8sQeBBmqNvIIis,1331 +zarr/storage/_common.py,sha256=8vEsKZCT1nA4dMMD30emCTR1jkMwbCvzeS0iy2GAALk,22615 +zarr/storage/_fsspec.py,sha256=FI9kJ-XC38ONObOACsJWQ68IURtqO5zj0OgDWZciZgw,14702 +zarr/storage/_local.py,sha256=NChyhZkaPs4NNzMXLrYDAHEbQd7vkotZsnKrLSCY07Y,20278 +zarr/storage/_logging.py,sha256=9rWvjpZLtvi0fs9Lu30YZ6EvwZqBNaucLO6WTpER1PU,7546 +zarr/storage/_memory.py,sha256=OseWGwGvT6S_Lvg8eGoo4twjtOsaeV37_NdLgrmPONo,17582 +zarr/storage/_obstore.py,sha256=8T17S05SrS80NOJUGauhpU8IPWo_FXyrmmpRh1ppr8w,17505 +zarr/storage/_utils.py,sha256=FkoUnXY7wEZTHiL-bdnkThGqQRPhSbQp8csEVc4o6Ng,4969 +zarr/storage/_wrapper.py,sha256=OBh6r1WNGeQqtYlxewrkN-JXJzoRvrkn40WOW3UG_ZQ,4876 +zarr/storage/_zip.py,sha256=Yc2RKoNFhOKkQq4XBNvvyB3cuyQUsjjQRAViy1hGMSY,9232 +zarr/testing/__init__.py,sha256=d1QaEs3zOnKTEQSDX1Z62HgsEEGQnjM7RjH3UU1Jkfg,436 +zarr/testing/buffer.py,sha256=r2E0FMSomfPaXocxs0t07xotgoClEdkpYbH70v6hLig,2330 +zarr/testing/conftest.py,sha256=IxBNMxVMC_PNze4KoSdnjWsMgrA8iU_aHBuiHJFq0lk,456 +zarr/testing/stateful.py,sha256=TFQYgsGlqNcGkpt9PR_py8H9croxdkmryiJoYwQ_Yhs,24665 +zarr/testing/store.py,sha256=Lo3gTsk33eIJUqZFHvcbVnZSXAobQ1c_QOeNbvHK9BM,27570 +zarr/testing/strategies.py,sha256=GlFnSvJ76Y7ayslJXP6zygGiXczlJvRPc-mfDhpMmSU,15509 +zarr/testing/utils.py,sha256=C0CMZsoDm4MfWWo00m2ZdH2JAmVZ4DrKvt3K5N7kGgs,1121 +zarr/types.py,sha256=lSktZQyOx7yn-BKrJZvxdpP-SDQkrFBfBxUhJDyb1Bk,664 diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/REQUESTED b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/WHEEL b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..b1b94fd58e7e9ed0ef3449473bc48de68afcc3fe --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: hatchling 1.29.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/entry_points.txt b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..04359de2fb8caae8afcaf16f179fc97096b23d26 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/entry_points.txt @@ -0,0 +1,5 @@ +[console_scripts] +zarr = zarr._cli.cli:app + +[pytest11] +zarr = zarr.testing diff --git a/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/licenses/LICENSE.txt b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/licenses/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..1e8da4d2424241744037944bdf0a54ff458b3524 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr-3.1.6.dist-info/licenses/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/.venv/lib/python3.11/site-packages/zarr/core/attributes.py b/.venv/lib/python3.11/site-packages/zarr/core/attributes.py new file mode 100644 index 0000000000000000000000000000000000000000..7f29e443655ae03c92bba26d91f791d0ab7a7781 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/attributes.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import TYPE_CHECKING + +from zarr.core.common import JSON + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr.core.group import Group + from zarr.types import AnyArray + + +class Attributes(MutableMapping[str, JSON]): + def __init__(self, obj: AnyArray | Group) -> None: + # key=".zattrs", read_only=False, cache=True, synchronizer=None + self._obj = obj + + def __getitem__(self, key: str) -> JSON: + return self._obj.metadata.attributes[key] + + def __setitem__(self, key: str, value: JSON) -> None: + new_attrs = dict(self._obj.metadata.attributes) + new_attrs[key] = value + self._obj = self._obj.update_attributes(new_attrs) + + def __delitem__(self, key: str) -> None: + new_attrs = dict(self._obj.metadata.attributes) + del new_attrs[key] + self.put(new_attrs) + + def __iter__(self) -> Iterator[str]: + return iter(self._obj.metadata.attributes) + + def __len__(self) -> int: + return len(self._obj.metadata.attributes) + + def put(self, d: dict[str, JSON]) -> None: + """ + Overwrite all attributes with the values from `d`. + + Equivalent to the following pseudo-code, but performed atomically. + + ```python + attrs = {"a": 1, "b": 2} + attrs.clear() + attrs.update({"a": "3", "c": 4}) + print(attrs) + #> {'a': '3', 'c': 4} + ``` + """ + self._obj.metadata.attributes.clear() + self._obj = self._obj.update_attributes(d) + + def asdict(self) -> dict[str, JSON]: + return dict(self._obj.metadata.attributes) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/__init__.py b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ebec61a372b67ae7dcc12c8855852cdea828c12a --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__init__.py @@ -0,0 +1,21 @@ +from zarr.core.buffer.core import ( + ArrayLike, + Buffer, + BufferPrototype, + NDArrayLike, + NDArrayLikeOrScalar, + NDBuffer, + default_buffer_prototype, +) +from zarr.core.buffer.cpu import numpy_buffer_prototype + +__all__ = [ + "ArrayLike", + "Buffer", + "BufferPrototype", + "NDArrayLike", + "NDArrayLikeOrScalar", + "NDBuffer", + "default_buffer_prototype", + "numpy_buffer_prototype", +] diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3da42caa9ca0e9a5bace04ca8098016296b8fec Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/core.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/core.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d3423bcc83e361c3a4d3ece0818e86fc74a275b Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/core.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/cpu.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/cpu.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9f81fa15df9ed69d757825bc5b93434e5fdc92a Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/cpu.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/gpu.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/gpu.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..89f96eb892ad91d635de25432bc9f0406eec8e34 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/buffer/__pycache__/gpu.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/core.py b/.venv/lib/python3.11/site-packages/zarr/core/buffer/core.py new file mode 100644 index 0000000000000000000000000000000000000000..9602a55258e5800765cb9cf676f5917c87d1ef09 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/buffer/core.py @@ -0,0 +1,586 @@ +from __future__ import annotations + +import sys +from abc import ABC, abstractmethod +from collections.abc import Iterable +from typing import ( + TYPE_CHECKING, + Any, + Literal, + NamedTuple, + Protocol, + SupportsIndex, + cast, + runtime_checkable, +) + +import numpy as np +import numpy.typing as npt + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + from typing import Self + + from zarr.codecs.bytes import Endian + from zarr.core.common import BytesLike + +# Everything here is imported into ``zarr.core.buffer`` namespace. +__all__: list[str] = [] + + +@runtime_checkable +class ArrayLike(Protocol): + """Protocol for the array-like type that underlie Buffer""" + + @property + def dtype(self) -> np.dtype[Any]: ... + + @property + def ndim(self) -> int: ... + + @property + def size(self) -> int: ... + + def __getitem__(self, key: slice) -> Self: ... + + def __setitem__(self, key: slice, value: Any) -> None: ... + + +@runtime_checkable +class NDArrayLike(Protocol): + """Protocol for the nd-array-like type that underlie NDBuffer""" + + @property + def dtype(self) -> np.dtype[Any]: ... + + @property + def ndim(self) -> int: ... + + @property + def size(self) -> int: ... + + @property + def shape(self) -> tuple[int, ...]: ... + + def __len__(self) -> int: ... + + def __getitem__(self, key: slice) -> Self: ... + + def __setitem__(self, key: slice, value: Any) -> None: ... + + def __array__(self) -> npt.NDArray[Any]: ... + + def reshape( + self, shape: tuple[int, ...] | Literal[-1], *, order: Literal["A", "C", "F"] = ... + ) -> Self: ... + + def view(self, dtype: npt.DTypeLike) -> Self: ... + + def astype( + self, + dtype: npt.DTypeLike, + order: Literal["K", "A", "C", "F"] = ..., + *, + copy: bool = ..., + ) -> Self: ... + + def fill(self, value: Any) -> None: ... + + def copy(self) -> Self: ... + + def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Self: ... + + def ravel(self, order: Literal["K", "A", "C", "F"] = ...) -> Self: ... + + def all(self) -> bool: ... + + def __eq__(self, other: object) -> Self: # type: ignore[override] + """Element-wise equal + + Notes + ----- + Type checkers such as mypy complains because the return type isn't a bool like + its supertype "object", which violates the Liskov substitution principle. + This is true, but since NumPy's ndarray is defined as an element-wise equal, + our hands are tied. + """ + + +ScalarType = int | float | complex | bytes | str | bool | np.generic +NDArrayLikeOrScalar = ScalarType | NDArrayLike + + +def check_item_key_is_1d_contiguous(key: Any) -> None: + """Raises error if `key` isn't a 1d contiguous slice""" + if not isinstance(key, slice): + raise TypeError( + f"Item key has incorrect type (expected slice, got {key.__class__.__name__})" + ) + if not (key.step is None or key.step == 1): + raise ValueError("slice must be contiguous") + + +class Buffer(ABC): + """A flat contiguous memory block + + We use Buffer throughout Zarr to represent a contiguous block of memory. + + A Buffer is backed by an underlying array-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + array-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + This buffer is untyped, so all indexing and sizes are in bytes. + + Parameters + ---------- + array_like + array-like object that must be 1-dim, contiguous, and byte dtype. + """ + + def __init__(self, array_like: ArrayLike) -> None: + if array_like.ndim != 1: + raise ValueError("array_like: only 1-dim allowed") + if array_like.dtype != np.dtype("B"): + raise ValueError("array_like: only byte dtype allowed") + self._data = array_like + + @classmethod + @abstractmethod + def create_zero_length(cls) -> Self: + """Create an empty buffer with length zero + + Returns + ------- + New empty 0-length buffer + """ + if cls is Buffer: + raise NotImplementedError("Cannot call abstract method on the abstract class 'Buffer'") + return cls( + cast("ArrayLike", None) + ) # This line will never be reached, but it satisfies the type checker + + @classmethod + def from_array_like(cls, array_like: ArrayLike) -> Self: + """Create a new buffer of an array-like object + + Parameters + ---------- + array_like + array-like object that must be 1-dim, contiguous, and byte dtype. + + Returns + ------- + New buffer representing `array_like` + """ + return cls(array_like) + + @classmethod + @abstractmethod + def from_buffer(cls, buffer: Buffer) -> Self: + """Create a new buffer of an existing Buffer + + This is useful if you want to ensure that an existing buffer is + of the correct subclass of Buffer. E.g., MemoryStore uses this + to return a buffer instance of the subclass specified by its + BufferPrototype argument. + + Typically, this only copies data if the data has to be moved between + memory types, such as from host to device memory. + + Parameters + ---------- + buffer + buffer object. + + Returns + ------- + A new buffer representing the content of the input buffer + + Notes + ----- + Subclasses of `Buffer` must override this method to implement + more optimal conversions that avoid copies where possible + """ + if cls is Buffer: + raise NotImplementedError("Cannot call abstract method on the abstract class 'Buffer'") + return cls( + cast("ArrayLike", None) + ) # This line will never be reached, but it satisfies the type checker + + @classmethod + @abstractmethod + def from_bytes(cls, bytes_like: BytesLike) -> Self: + """Create a new buffer of a bytes-like object (host memory) + + Parameters + ---------- + bytes_like + bytes-like object + + Returns + ------- + New buffer representing `bytes_like` + """ + if cls is Buffer: + raise NotImplementedError("Cannot call abstract method on the abstract class 'Buffer'") + return cls( + cast("ArrayLike", None) + ) # This line will never be reached, but it satisfies the type checker + + def as_array_like(self) -> ArrayLike: + """Returns the underlying array (host or device memory) of this buffer + + This will never copy data. + + Returns + ------- + The underlying 1d array such as a NumPy or CuPy array. + """ + return self._data + + @abstractmethod + def as_numpy_array(self) -> npt.NDArray[Any]: + """Returns the buffer as a NumPy array (host memory). + + Notes + ----- + Might have to copy data, consider using `.as_array_like()` instead. + + Returns + ------- + NumPy array of this buffer (might be a data copy) + """ + ... + + def as_buffer_like(self) -> BytesLike: + """Returns the buffer as an object that implements the Python buffer protocol. + + Notes + ----- + Might have to copy data, since the implementation uses `.as_numpy_array()`. + + Returns + ------- + An object that implements the Python buffer protocol + """ + return memoryview(self.as_numpy_array()) # type: ignore[arg-type] + + def to_bytes(self) -> bytes: + """Returns the buffer as `bytes` (host memory). + + Warnings + -------- + Will always copy data, only use this method for small buffers such as metadata + buffers. If possible, use `.as_numpy_array()` or `.as_array_like()` instead. + + Returns + ------- + `bytes` of this buffer (data copy) + """ + return bytes(self.as_numpy_array()) + + def __getitem__(self, key: slice) -> Self: + check_item_key_is_1d_contiguous(key) + return self.__class__(self._data.__getitem__(key)) + + def __setitem__(self, key: slice, value: Any) -> None: + check_item_key_is_1d_contiguous(key) + self._data.__setitem__(key, value) + + def __len__(self) -> int: + return self._data.size + + @abstractmethod + def combine(self, others: Iterable[Buffer]) -> Self: + """Concatenate many buffers""" + ... + + def __add__(self, other: Buffer) -> Self: + """Concatenate two buffers""" + return self.combine([other]) + + def __eq__(self, other: object) -> bool: + # Another Buffer class can override this to choose a more efficient path + return isinstance(other, Buffer) and np.array_equal( + self.as_numpy_array(), other.as_numpy_array() + ) + + +class NDBuffer: + """An n-dimensional memory block + + We use NDBuffer throughout Zarr to represent a n-dimensional memory block. + + An NDBuffer is backed by an underlying ndarray-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + ndarray-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer + is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However, + in order to use Python's type system to differentiate between the contiguous + Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the + two classes separate. + + Parameters + ---------- + array : ndarray_like + ndarray-like object that is convertible to a regular Numpy array. + """ + + def __init__(self, array: NDArrayLike) -> None: + self._data = array + + @classmethod + @abstractmethod + def create( + cls, + *, + shape: Iterable[int], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + fill_value: Any | None = None, + ) -> Self: + """Create a new buffer and its underlying ndarray-like object + + Parameters + ---------- + shape + The shape of the buffer and its underlying ndarray-like object + dtype + The datatype of the buffer and its underlying ndarray-like object + order + Whether to store multi-dimensional data in row-major (C-style) or + column-major (Fortran-style) order in memory. + fill_value + If not None, fill the new buffer with a scalar value. + + Returns + ------- + New buffer representing a new ndarray_like object + + Notes + ----- + A subclass can overwrite this method to create an ndarray-like object + other then the default Numpy array. + """ + if cls is NDBuffer: + raise NotImplementedError( + "Cannot call abstract method on the abstract class 'NDBuffer'" + ) + return cls( + cast("NDArrayLike", None) + ) # This line will never be reached, but it satisfies the type checker + + @classmethod + def empty( + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + ) -> Self: + """ + Create an empty buffer with the given shape, dtype, and order. + + This method can be faster than ``NDBuffer.create`` because it doesn't + have to initialize the memory used by the underlying ndarray-like + object. + + Parameters + ---------- + shape + The shape of the buffer and its underlying ndarray-like object + dtype + The datatype of the buffer and its underlying ndarray-like object + order + Whether to store multi-dimensional data in row-major (C-style) or + column-major (Fortran-style) order in memory. + + Returns + ------- + buffer + New buffer representing a new ndarray_like object with empty data. + + See Also + -------- + NDBuffer.create + Create a new buffer with some initial fill value. + """ + # Implementations should override this method if they have a faster way + # to allocate an empty buffer. + return cls.create(shape=shape, dtype=dtype, order=order) + + @classmethod + def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self: + """Create a new buffer of an ndarray-like object + + Parameters + ---------- + ndarray_like + ndarray-like object + + Returns + ------- + New buffer representing `ndarray_like` + """ + return cls(ndarray_like) + + @classmethod + @abstractmethod + def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self: + """Create a new buffer of Numpy array-like object + + Parameters + ---------- + array_like + Object that can be coerced into a Numpy array + + Returns + ------- + New buffer representing `array_like` + """ + if cls is NDBuffer: + raise NotImplementedError( + "Cannot call abstract method on the abstract class 'NDBuffer'" + ) + return cls( + cast("NDArrayLike", None) + ) # This line will never be reached, but it satisfies the type checker + + def as_ndarray_like(self) -> NDArrayLike: + """Returns the underlying array (host or device memory) of this buffer + + This will never copy data. + + Returns + ------- + The underlying array such as a NumPy or CuPy array. + """ + return self._data + + @abstractmethod + def as_numpy_array(self) -> npt.NDArray[Any]: + """Returns the buffer as a NumPy array (host memory). + + Warnings + -------- + Might have to copy data, consider using `.as_ndarray_like()` instead. + + Returns + ------- + NumPy array of this buffer (might be a data copy) + """ + ... + + def as_scalar(self) -> ScalarType: + """Returns the buffer as a scalar value""" + if self._data.size != 1: + raise ValueError("Buffer does not contain a single scalar value") + return cast("ScalarType", self.as_numpy_array()[()]) + + @property + def dtype(self) -> np.dtype[Any]: + return self._data.dtype + + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape + + @property + def byteorder(self) -> Endian: + from zarr.codecs.bytes import Endian + + if self.dtype.byteorder == "<": + return Endian.little + elif self.dtype.byteorder == ">": + return Endian.big + else: + return Endian(sys.byteorder) + + def reshape(self, newshape: tuple[int, ...] | Literal[-1]) -> Self: + return self.__class__(self._data.reshape(newshape)) + + def squeeze(self, axis: tuple[int, ...]) -> Self: + newshape = tuple(a for i, a in enumerate(self.shape) if i not in axis) + return self.__class__(self._data.reshape(newshape)) + + def astype(self, dtype: npt.DTypeLike, order: Literal["K", "A", "C", "F"] = "K") -> Self: + return self.__class__(self._data.astype(dtype=dtype, order=order)) + + @abstractmethod + def __getitem__(self, key: Any) -> Self: ... + + @abstractmethod + def __setitem__(self, key: Any, value: Any) -> None: ... + + def __len__(self) -> int: + return self._data.__len__() + + def __repr__(self) -> str: + return f"" + + def all_equal(self, other: Any, equal_nan: bool = True) -> bool: + """Compare to `other` using np.array_equal.""" + if other is None: + # Handle None fill_value for Zarr V2 + return False + # Handle positive and negative zero by comparing bit patterns: + if ( + np.asarray(other).dtype.kind == "f" + and other == 0.0 + and self._data.dtype.kind not in ("U", "S", "T", "O", "V") + ): + _data, other = np.broadcast_arrays(self._data, np.asarray(other, self._data.dtype)) + void_dtype = "V" + str(_data.dtype.itemsize) + return np.array_equal(_data.view(void_dtype), other.view(void_dtype)) + # use array_equal to obtain equal_nan=True functionality + # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value + # every single time we have to write data? + _data, other = np.broadcast_arrays(self._data, other) + return np.array_equal( + _data, + other, + equal_nan=equal_nan + if self._data.dtype.kind not in ("U", "S", "T", "O", "V") + else False, + ) + + def fill(self, value: Any) -> None: + self._data.fill(value) + + def copy(self) -> Self: + return self.__class__(self._data.copy()) + + def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Self: + return self.__class__(self._data.transpose(axes)) + + +class BufferPrototype(NamedTuple): + """Prototype of the Buffer and NDBuffer class + + The protocol must be pickable. + + Attributes + ---------- + buffer + The Buffer class to use when Zarr needs to create new Buffer. + nd_buffer + The NDBuffer class to use when Zarr needs to create new NDBuffer. + """ + + buffer: type[Buffer] + nd_buffer: type[NDBuffer] + + +# The default buffer prototype used throughout the Zarr codebase. +def default_buffer_prototype() -> BufferPrototype: + from zarr.registry import ( + get_buffer_class, + get_ndbuffer_class, + ) + + return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class()) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/cpu.py b/.venv/lib/python3.11/site-packages/zarr/core/buffer/cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..8994281b582a2947199cd3f3556fc3562dfd4fb5 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/buffer/cpu.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from typing import ( + TYPE_CHECKING, + Any, + Literal, +) + +import numpy as np +import numpy.typing as npt + +from zarr.core.buffer import core +from zarr.registry import ( + register_buffer, + register_ndbuffer, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + from typing import Self + + from zarr.core.buffer.core import ArrayLike, NDArrayLike + from zarr.core.common import BytesLike + + +class Buffer(core.Buffer): + """A flat contiguous memory block + + We use Buffer throughout Zarr to represent a contiguous block of memory. + + A Buffer is backed by an underlying array-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + array-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + This buffer is untyped, so all indexing and sizes are in bytes. + + Parameters + ---------- + array_like + array-like object that must be 1-dim, contiguous, and byte dtype. + """ + + def __init__(self, array_like: ArrayLike) -> None: + super().__init__(array_like) + + @classmethod + def create_zero_length(cls) -> Self: + return cls(np.array([], dtype="B")) + + @classmethod + def from_buffer(cls, buffer: core.Buffer) -> Self: + """Create a new buffer of an existing Buffer + + This is useful if you want to ensure that an existing buffer is + of the correct subclass of Buffer. E.g., MemoryStore uses this + to return a buffer instance of the subclass specified by its + BufferPrototype argument. + + Typically, this only copies data if the data has to be moved between + memory types, such as from host to device memory. + + Parameters + ---------- + buffer + buffer object. + + Returns + ------- + A new buffer representing the content of the input buffer + + Notes + ----- + Subclasses of `Buffer` must override this method to implement + more optimal conversions that avoid copies where possible + """ + return cls.from_array_like(buffer.as_numpy_array()) + + @classmethod + def from_bytes(cls, bytes_like: BytesLike) -> Self: + """Create a new buffer of a bytes-like object (host memory) + + Parameters + ---------- + bytes_like + bytes-like object + + Returns + ------- + New buffer representing `bytes_like` + """ + return cls.from_array_like(np.frombuffer(bytes_like, dtype="B")) + + def as_numpy_array(self) -> npt.NDArray[Any]: + """Returns the buffer as a NumPy array (host memory). + + Notes + ----- + Might have to copy data, consider using `.as_array_like()` instead. + + Returns + ------- + NumPy array of this buffer (might be a data copy) + """ + return np.asanyarray(self._data) + + def combine(self, others: Iterable[core.Buffer]) -> Self: + data = [np.asanyarray(self._data)] + for buf in others: + other_array = buf.as_array_like() + assert other_array.dtype == np.dtype("B") + data.append(np.asanyarray(other_array)) + return self.__class__(np.concatenate(data)) + + +class NDBuffer(core.NDBuffer): + """An n-dimensional memory block + + We use NDBuffer throughout Zarr to represent a n-dimensional memory block. + + An NDBuffer is backed by an underlying ndarray-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + ndarray-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer + is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However, + in order to use Python's type system to differentiate between the contiguous + Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the + two classes separate. + + Parameters + ---------- + array + ndarray-like object that is convertible to a regular Numpy array. + """ + + def __init__(self, array: NDArrayLike) -> None: + super().__init__(array) + + @classmethod + def create( + cls, + *, + shape: Iterable[int], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + fill_value: Any | None = None, + ) -> Self: + # np.zeros is much faster than np.full, and therefore using it when possible is better. + if fill_value is None or (isinstance(fill_value, int) and fill_value == 0): + return cls(np.zeros(shape=tuple(shape), dtype=dtype, order=order)) + else: + return cls(np.full(shape=tuple(shape), fill_value=fill_value, dtype=dtype, order=order)) + + @classmethod + def empty( + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + ) -> Self: + return cls(np.empty(shape=shape, dtype=dtype, order=order)) + + @classmethod + def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self: + return cls.from_ndarray_like(np.asanyarray(array_like)) + + def as_numpy_array(self) -> npt.NDArray[Any]: + """Returns the buffer as a NumPy array (host memory). + + Warnings + -------- + Might have to copy data, consider using `.as_ndarray_like()` instead. + + Returns + ------- + NumPy array of this buffer (might be a data copy) + """ + return np.asanyarray(self._data) + + def __getitem__(self, key: Any) -> Self: + return self.__class__(np.asanyarray(self._data.__getitem__(key))) + + def __setitem__(self, key: Any, value: Any) -> None: + if isinstance(value, NDBuffer): + value = value._data + self._data.__setitem__(key, value) + + +def as_numpy_array_wrapper( + func: Callable[[npt.NDArray[Any]], bytes], buf: core.Buffer, prototype: core.BufferPrototype +) -> core.Buffer: + """Converts the input of `func` to a numpy array and the output back to `Buffer`. + + This function is useful when calling a `func` that only support host memory such + as `GZip.decode` and `Blosc.decode`. In this case, use this wrapper to convert + the input `buf` to a Numpy array and convert the result back into a `Buffer`. + + Parameters + ---------- + func + The callable that will be called with the converted `buf` as input. + `func` must return bytes, which will be converted into a `Buffer` + before returned. + buf + The buffer that will be converted to a Numpy array before given as + input to `func`. + prototype + The prototype of the output buffer. + + Returns + ------- + The result of `func` converted to a `Buffer` + """ + return prototype.buffer.from_bytes(func(buf.as_numpy_array())) + + +# CPU buffer prototype using numpy arrays +buffer_prototype = core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer) +# default_buffer_prototype = buffer_prototype + + +# The numpy prototype used for E.g. when reading the shard index +def numpy_buffer_prototype() -> core.BufferPrototype: + return core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer) + + +register_buffer(Buffer, qualname="zarr.buffer.cpu.Buffer") +register_ndbuffer(NDBuffer, qualname="zarr.buffer.cpu.NDBuffer") + + +# backwards compatibility +register_buffer(Buffer, qualname="zarr.core.buffer.cpu.Buffer") +register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.cpu.NDBuffer") diff --git a/.venv/lib/python3.11/site-packages/zarr/core/buffer/gpu.py b/.venv/lib/python3.11/site-packages/zarr/core/buffer/gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..8672942364d62481c814fa2d287280ec2174eb9f --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/buffer/gpu.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import warnings +from typing import ( + TYPE_CHECKING, + Any, + Literal, + cast, +) + +import numpy as np +import numpy.typing as npt + +from zarr.core.buffer import core +from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike +from zarr.errors import ZarrUserWarning +from zarr.registry import ( + register_buffer, + register_ndbuffer, +) + +if TYPE_CHECKING: + from collections.abc import Iterable + from typing import Self + + from zarr.core.common import BytesLike + +try: + import cupy as cp +except ImportError: + cp = None + + +class Buffer(core.Buffer): + """A flat contiguous memory block on the GPU + + We use Buffer throughout Zarr to represent a contiguous block of memory. + + A Buffer is backed by an underlying array-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + array-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + This buffer is untyped, so all indexing and sizes are in bytes. + + Parameters + ---------- + array_like + array-like object that must be 1-dim, contiguous, and byte dtype. + """ + + def __init__(self, array_like: ArrayLike) -> None: + if cp is None: + raise ImportError( + "Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy." + ) + + if array_like.ndim != 1: + raise ValueError("array_like: only 1-dim allowed") + if array_like.dtype != np.dtype("B"): + raise ValueError("array_like: only byte dtype allowed") + + if not hasattr(array_like, "__cuda_array_interface__"): + # Slow copy based path for arrays that don't support the __cuda_array_interface__ + # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol + msg = ( + "Creating a zarr.buffer.gpu.Buffer with an array that does not support the " + "__cuda_array_interface__ for zero-copy transfers, " + "falling back to slow copy based path" + ) + warnings.warn( + msg, + category=ZarrUserWarning, + stacklevel=2, + ) + self._data = cp.asarray(array_like) + + @classmethod + def create_zero_length(cls) -> Self: + """Create an empty buffer with length zero + + Returns + ------- + New empty 0-length buffer + """ + return cls(cp.array([], dtype="B")) + + @classmethod + def from_buffer(cls, buffer: core.Buffer) -> Self: + """Create a GPU Buffer given an arbitrary Buffer + This will try to be zero-copy if `buffer` is already on the + GPU and will trigger a copy if not. + + Returns + ------- + New GPU Buffer constructed from `buffer` + """ + return cls(buffer.as_array_like()) + + @classmethod + def from_bytes(cls, bytes_like: BytesLike) -> Self: + return cls.from_array_like(cp.frombuffer(bytes_like, dtype="B")) + + def as_numpy_array(self) -> npt.NDArray[Any]: + return cast("npt.NDArray[Any]", cp.asnumpy(self._data)) + + def combine(self, others: Iterable[core.Buffer]) -> Self: + data = [cp.asanyarray(self._data)] + for other in others: + other_array = other.as_array_like() + assert other_array.dtype == np.dtype("B") + gpu_other = Buffer(other_array) + gpu_other_array = gpu_other.as_array_like() + data.append(cp.asanyarray(gpu_other_array)) + return self.__class__(cp.concatenate(data)) + + +class NDBuffer(core.NDBuffer): + """A n-dimensional memory block on the GPU + + We use NDBuffer throughout Zarr to represent a n-dimensional memory block. + + An NDBuffer is backed by an underlying ndarray-like instance that represents + the memory. The memory type is unspecified; can be regular host memory, + CUDA device memory, or something else. The only requirement is that the + ndarray-like instance can be copied/converted to a regular Numpy array + (host memory). + + Notes + ----- + The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer + is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However, + in order to use Python's type system to differentiate between the contiguous + Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the + two classes separate. + + Parameters + ---------- + array + ndarray-like object that is convertible to a regular Numpy array. + """ + + def __init__(self, array: NDArrayLike) -> None: + if cp is None: + raise ImportError( + "Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy." + ) + + # assert array.ndim > 0 + assert array.dtype != object + self._data = array + + if not hasattr(array, "__cuda_array_interface__"): + # Slow copy based path for arrays that don't support the __cuda_array_interface__ + # TODO: Add a fast zero-copy path for arrays that support the dlpack protocol + msg = ( + "Creating a zarr.buffer.gpu.NDBuffer with an array that does not support the " + "__cuda_array_interface__ for zero-copy transfers, " + "falling back to slow copy based path" + ) + warnings.warn( + msg, + stacklevel=2, + ) + self._data = cp.asarray(array) + + @classmethod + def create( + cls, + *, + shape: Iterable[int], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + fill_value: Any | None = None, + ) -> Self: + ret = cls(cp.empty(shape=tuple(shape), dtype=dtype, order=order)) + if fill_value is not None: + ret.fill(fill_value) + return ret + + @classmethod + def empty( + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + ) -> Self: + return cls(cp.empty(shape=shape, dtype=dtype, order=order)) + + @classmethod + def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self: + """Create a new buffer of Numpy array-like object + + Parameters + ---------- + array_like + Object that can be coerced into a Numpy array + + Returns + ------- + New buffer representing `array_like` + """ + return cls(cp.asarray(array_like)) + + def as_numpy_array(self) -> npt.NDArray[Any]: + """Returns the buffer as a NumPy array (host memory). + + Warnings + -------- + Might have to copy data, consider using `.as_ndarray_like()` instead. + + Returns + ------- + NumPy array of this buffer (might be a data copy) + """ + return cast("npt.NDArray[Any]", cp.asnumpy(self._data)) + + def __getitem__(self, key: Any) -> Self: + return self.__class__(self._data.__getitem__(key)) + + def __setitem__(self, key: Any, value: Any) -> None: + if isinstance(value, NDBuffer): + value = value._data + elif isinstance(value, core.NDBuffer): + gpu_value = NDBuffer(value.as_ndarray_like()) + value = gpu_value._data + self._data.__setitem__(key, value) + + +buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer) + +register_buffer(Buffer, qualname="zarr.buffer.gpu.Buffer") +register_ndbuffer(NDBuffer, qualname="zarr.buffer.gpu.NDBuffer") + +# backwards compatibility +register_buffer(Buffer, qualname="zarr.core.buffer.gpu.Buffer") +register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.gpu.NDBuffer") diff --git a/.venv/lib/python3.11/site-packages/zarr/core/chunk_grids.py b/.venv/lib/python3.11/site-packages/zarr/core/chunk_grids.py new file mode 100644 index 0000000000000000000000000000000000000000..c903eba013d0539b4d659dc1a0426c5e3c229c80 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/chunk_grids.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +import itertools +import math +import numbers +import operator +import warnings +from abc import abstractmethod +from dataclasses import dataclass +from functools import reduce +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np + +import zarr +from zarr.abc.metadata import Metadata +from zarr.core.common import ( + JSON, + NamedConfig, + ShapeLike, + ceildiv, + parse_named_configuration, + parse_shapelike, +) +from zarr.errors import ZarrUserWarning + +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import Self + + from zarr.core.array import ShardsLike + + +def _guess_chunks( + shape: tuple[int, ...] | int, + typesize: int, + *, + increment_bytes: int = 256 * 1024, + min_bytes: int = 128 * 1024, + max_bytes: int = 64 * 1024 * 1024, +) -> tuple[int, ...]: + """ + Iteratively guess an appropriate chunk layout for an array, given its shape and + the size of each element in bytes, and size constraints expressed in bytes. This logic is + adapted from h5py. + + Parameters + ---------- + shape : tuple[int, ...] + The chunk shape. + typesize : int + The size, in bytes, of each element of the chunk. + increment_bytes : int = 256 * 1024 + The number of bytes used to increment or decrement the target chunk size in bytes. + min_bytes : int = 128 * 1024 + The soft lower bound on the final chunk size in bytes. + max_bytes : int = 64 * 1024 * 1024 + The hard upper bound on the final chunk size in bytes. + + Returns + ------- + tuple[int, ...] + + """ + if min_bytes >= max_bytes: + raise ValueError(f"Cannot have more min_bytes ({min_bytes}) than max_bytes ({max_bytes})") + if isinstance(shape, int): + shape = (shape,) + + if typesize == 0: + return shape + + ndims = len(shape) + # require chunks to have non-zero length for all dimensions + chunks = np.maximum(np.array(shape, dtype="=f8"), 1) + + # Determine the optimal chunk size in bytes using a PyTables expression. + # This is kept as a float. + dset_size = np.prod(chunks) * typesize + target_size = increment_bytes * (2 ** np.log10(dset_size / (1024.0 * 1024))) + + if target_size > max_bytes: + target_size = max_bytes + elif target_size < min_bytes: + target_size = min_bytes + + idx = 0 + while True: + # Repeatedly loop over the axes, dividing them by 2. Stop when: + # 1a. We're smaller than the target chunk size, OR + # 1b. We're within 50% of the target chunk size, AND + # 2. The chunk is smaller than the maximum chunk size + + chunk_bytes = np.prod(chunks) * typesize + + if ( + chunk_bytes < target_size or abs(chunk_bytes - target_size) / target_size < 0.5 + ) and chunk_bytes < max_bytes: + break + + if np.prod(chunks) == 1: + break # Element size larger than max_bytes + + chunks[idx % ndims] = math.ceil(chunks[idx % ndims] / 2.0) + idx += 1 + + return tuple(int(x) for x in chunks) + + +def normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tuple[int, ...]: + """Convenience function to normalize the `chunks` argument for an array + with the given `shape`.""" + + # N.B., expect shape already normalized + + # handle auto-chunking + if chunks is None or chunks is True: + return _guess_chunks(shape, typesize) + + # handle no chunking + if chunks is False: + return shape + + # handle 1D convenience form + if isinstance(chunks, numbers.Integral): + chunks = tuple(int(chunks) for _ in shape) + + # handle dask-style chunks (iterable of iterables) + if all(isinstance(c, (tuple, list)) for c in chunks): + for i, c in enumerate(chunks): + if any(x != y for x, y in itertools.pairwise(c[:-1])) or (len(c) > 1 and c[-1] > c[0]): + raise ValueError( + f"Irregular chunk sizes in dimension {i}: {tuple(c)}. " + "Only uniform chunks (with an optional smaller final chunk) are supported." + ) + chunks = tuple(c[0] for c in chunks) + + # handle bad dimensionality + if len(chunks) > len(shape): + raise ValueError("too many dimensions in chunks") + + # handle underspecified chunks + if len(chunks) < len(shape): + # assume chunks across remaining dimensions + chunks += shape[len(chunks) :] + + # handle None or -1 in chunks + if -1 in chunks or None in chunks: + chunks = tuple( + s if c == -1 or c is None else int(c) for s, c in zip(shape, chunks, strict=False) + ) + + if not all(isinstance(c, numbers.Integral) for c in chunks): + raise TypeError("non integer value in chunks") + + return tuple(int(c) for c in chunks) + + +@dataclass(frozen=True) +class ChunkGrid(Metadata): + @classmethod + def from_dict(cls, data: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any]) -> ChunkGrid: + if isinstance(data, ChunkGrid): + return data + + name_parsed, _ = parse_named_configuration(data) + if name_parsed == "regular": + return RegularChunkGrid._from_dict(data) + raise ValueError(f"Unknown chunk grid. Got {name_parsed}.") + + @abstractmethod + def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: + pass + + @abstractmethod + def get_nchunks(self, array_shape: tuple[int, ...]) -> int: + pass + + +@dataclass(frozen=True) +class RegularChunkGrid(ChunkGrid): + chunk_shape: tuple[int, ...] + + def __init__(self, *, chunk_shape: ShapeLike) -> None: + chunk_shape_parsed = parse_shapelike(chunk_shape) + + object.__setattr__(self, "chunk_shape", chunk_shape_parsed) + + @classmethod + def _from_dict(cls, data: dict[str, JSON] | NamedConfig[str, Any]) -> Self: + _, configuration_parsed = parse_named_configuration(data, "regular") + + return cls(**configuration_parsed) # type: ignore[arg-type] + + def to_dict(self) -> dict[str, JSON]: + return {"name": "regular", "configuration": {"chunk_shape": tuple(self.chunk_shape)}} + + def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: + return itertools.product( + *(range(ceildiv(s, c)) for s, c in zip(array_shape, self.chunk_shape, strict=False)) + ) + + def get_nchunks(self, array_shape: tuple[int, ...]) -> int: + return reduce( + operator.mul, + itertools.starmap(ceildiv, zip(array_shape, self.chunk_shape, strict=True)), + 1, + ) + + +def _guess_num_chunks_per_axis_shard( + chunk_shape: tuple[int, ...], item_size: int, max_bytes: int, array_shape: tuple[int, ...] +) -> int: + """Generate the number of chunks per axis to hit a target max byte size for a shard. + + For example, for a (2,2,2) chunk size and item size 4, maximum bytes of 256 would return 2. + In other words the shard would be a (2,2,2) grid of (2,2,2) chunks + i.e., prod(chunk_shape) * (returned_val * len(chunk_shape)) * item_size = 256 bytes. + + Parameters + ---------- + chunk_shape + The shape of the (inner) chunks. + item_size + The item size of the data i.e., 2 for uint16. + max_bytes + The maximum number of bytes per shard to allow. + array_shape + The shape of the underlying array. + + Returns + ------- + The number of chunks per axis. + """ + bytes_per_chunk = np.prod(chunk_shape) * item_size + if max_bytes < bytes_per_chunk: + return 1 + num_axes = len(chunk_shape) + chunks_per_shard = 1 + # First check for byte size, second check to make sure we don't go bigger than the array shape + while (bytes_per_chunk * ((chunks_per_shard + 1) ** num_axes)) <= max_bytes and all( + c * (chunks_per_shard + 1) <= a for c, a in zip(chunk_shape, array_shape, strict=True) + ): + chunks_per_shard += 1 + return chunks_per_shard + + +def _auto_partition( + *, + array_shape: tuple[int, ...], + chunk_shape: tuple[int, ...] | Literal["auto"], + shard_shape: ShardsLike | None, + item_size: int, +) -> tuple[tuple[int, ...] | None, tuple[int, ...]]: + """ + Automatically determine the shard shape and chunk shape for an array, given the shape and dtype of the array. + If `shard_shape` is `None` and the chunk_shape is "auto", the chunks will be set heuristically based + on the dtype and shape of the array. + If `shard_shape` is "auto", then the shard shape will be set heuristically from the dtype and shape + of the array; if the `chunk_shape` is also "auto", then the chunks will be set heuristically as well, + given the dtype and shard shape. Otherwise, the chunks will be returned as-is. + """ + if shard_shape is None: + _shards_out: None | tuple[int, ...] = None + if chunk_shape == "auto": + _chunks_out = _guess_chunks(array_shape, item_size) + else: + _chunks_out = chunk_shape + else: + if chunk_shape == "auto": + # aim for a 1MiB chunk + _chunks_out = _guess_chunks(array_shape, item_size, max_bytes=1048576) + else: + _chunks_out = chunk_shape + + if shard_shape == "auto": + warnings.warn( + "Automatic shard shape inference is experimental and may change without notice.", + ZarrUserWarning, + stacklevel=2, + ) + _shards_out = () + target_shard_size_bytes = zarr.config.get("array.target_shard_size_bytes", None) + num_chunks_per_shard_axis = ( + _guess_num_chunks_per_axis_shard( + chunk_shape=_chunks_out, + item_size=item_size, + max_bytes=target_shard_size_bytes, + array_shape=array_shape, + ) + if (has_auto_shard := (target_shard_size_bytes is not None)) + else 2 + ) + for a_shape, c_shape in zip(array_shape, _chunks_out, strict=True): + # The previous heuristic was `a_shape // c_shape > 8` and now, with target_shard_size_bytes, we only check that the shard size is less than the array size. + can_shard_axis = a_shape // c_shape > 8 if not has_auto_shard else True + if can_shard_axis: + _shards_out += (c_shape * num_chunks_per_shard_axis,) + else: + _shards_out += (c_shape,) + elif isinstance(shard_shape, dict): + _shards_out = tuple(shard_shape["shape"]) + else: + _shards_out = shard_shape + + return _shards_out, _chunks_out diff --git a/.venv/lib/python3.11/site-packages/zarr/core/chunk_key_encodings.py b/.venv/lib/python3.11/site-packages/zarr/core/chunk_key_encodings.py new file mode 100644 index 0000000000000000000000000000000000000000..9eef80656da8aa74bc805e85a975a9692c09a13a --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/chunk_key_encodings.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypedDict, cast + +if TYPE_CHECKING: + from typing import NotRequired, Self + +from zarr.abc.metadata import Metadata +from zarr.core.common import ( + JSON, + NamedConfig, + parse_named_configuration, +) +from zarr.registry import get_chunk_key_encoding_class, register_chunk_key_encoding + +SeparatorLiteral = Literal[".", "/"] + + +def parse_separator(data: JSON) -> SeparatorLiteral: + if data not in (".", "/"): + raise ValueError(f"Expected an '.' or '/' separator. Got {data} instead.") + return cast("SeparatorLiteral", data) + + +class ChunkKeyEncodingParams(TypedDict): + name: Literal["v2", "default"] + separator: NotRequired[SeparatorLiteral] + + +@dataclass(frozen=True) +class ChunkKeyEncoding(ABC, Metadata): + """ + Defines how chunk coordinates are mapped to store keys. + + Subclasses must define a class variable `name` and implement `encode_chunk_key`. + """ + + name: ClassVar[str] + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, config_parsed = parse_named_configuration(data, require_configuration=False) + return cls(**config_parsed or {}) + + def to_dict(self) -> dict[str, JSON]: + return {"name": self.name, "configuration": super().to_dict()} + + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: + """ + Optional: decode a chunk key string into chunk coordinates. + Not required for normal operation; override if needed for testing or debugging. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not implement decode_chunk_key.") + + @abstractmethod + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + """ + Encode chunk coordinates into a chunk key string. + Must be implemented by subclasses. + """ + + +ChunkKeyEncodingLike: TypeAlias = ( + dict[str, JSON] | ChunkKeyEncodingParams | ChunkKeyEncoding | NamedConfig[str, Any] +) + + +@dataclass(frozen=True) +class DefaultChunkKeyEncoding(ChunkKeyEncoding): + name: ClassVar[Literal["default"]] = "default" + separator: SeparatorLiteral = "/" + + def __post_init__(self) -> None: + separator_parsed = parse_separator(self.separator) + object.__setattr__(self, "separator", separator_parsed) + + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: + if chunk_key == "c": + return () + return tuple(map(int, chunk_key[1:].split(self.separator))) + + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + return self.separator.join(map(str, ("c",) + chunk_coords)) + + +@dataclass(frozen=True) +class V2ChunkKeyEncoding(ChunkKeyEncoding): + name: ClassVar[Literal["v2"]] = "v2" + separator: SeparatorLiteral = "." + + def __post_init__(self) -> None: + separator_parsed = parse_separator(self.separator) + object.__setattr__(self, "separator", separator_parsed) + + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: + return tuple(map(int, chunk_key.split(self.separator))) + + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + chunk_identifier = self.separator.join(map(str, chunk_coords)) + return "0" if chunk_identifier == "" else chunk_identifier + + +def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: + """ + Take an implicit specification of a chunk key encoding and parse it into a ChunkKeyEncoding object. + """ + if isinstance(data, ChunkKeyEncoding): + return data + + # handle ChunkKeyEncodingParams + if "name" in data and "separator" in data: + data = {"name": data["name"], "configuration": {"separator": data["separator"]}} # type: ignore[typeddict-item] + + # Now must be a named config + data = cast("dict[str, JSON]", data) + + name_parsed, _ = parse_named_configuration(data, require_configuration=False) + try: + chunk_key_encoding = get_chunk_key_encoding_class(name_parsed).from_dict(data) + except KeyError as e: + raise ValueError(f"Unknown chunk key encoding: {e.args[0]!r}") from e + + return chunk_key_encoding + + +register_chunk_key_encoding("default", DefaultChunkKeyEncoding) +register_chunk_key_encoding("v2", V2ChunkKeyEncoding) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/codec_pipeline.py b/.venv/lib/python3.11/site-packages/zarr/core/codec_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..eed49556d3b0d97949477f0ea531aa4a0b2b72be --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/codec_pipeline.py @@ -0,0 +1,562 @@ +from __future__ import annotations + +from dataclasses import dataclass +from itertools import islice, pairwise +from typing import TYPE_CHECKING, Any, TypeVar +from warnings import warn + +from zarr.abc.codec import ( + ArrayArrayCodec, + ArrayBytesCodec, + ArrayBytesCodecPartialDecodeMixin, + ArrayBytesCodecPartialEncodeMixin, + BytesBytesCodec, + Codec, + CodecPipeline, +) +from zarr.core.common import concurrent_map +from zarr.core.config import config +from zarr.core.indexing import SelectorTuple, is_scalar +from zarr.errors import ZarrUserWarning +from zarr.registry import register_pipeline + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + from typing import Self + + from zarr.abc.store import ByteGetter, ByteSetter + from zarr.core.array_spec import ArraySpec + from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer + from zarr.core.chunk_grids import ChunkGrid + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + +T = TypeVar("T") +U = TypeVar("U") + + +def _unzip2(iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]: + out0: list[T] = [] + out1: list[U] = [] + for item0, item1 in iterable: + out0.append(item0) + out1.append(item1) + return (out0, out1) + + +def batched(iterable: Iterable[T], n: int) -> Iterable[tuple[T, ...]]: + if n < 1: + raise ValueError("n must be at least one") + it = iter(iterable) + while batch := tuple(islice(it, n)): + yield batch + + +def resolve_batched(codec: Codec, chunk_specs: Iterable[ArraySpec]) -> Iterable[ArraySpec]: + return [codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs] + + +def fill_value_or_default(chunk_spec: ArraySpec) -> Any: + fill_value = chunk_spec.fill_value + if fill_value is None: + # Zarr V2 allowed `fill_value` to be null in the metadata. + # Zarr V3 requires it to be set. This has already been + # validated when decoding the metadata, but we support reading + # Zarr V2 data and need to support the case where fill_value + # is None. + return chunk_spec.dtype.default_scalar() + else: + return fill_value + + +@dataclass(frozen=True) +class BatchedCodecPipeline(CodecPipeline): + """Default codec pipeline. + + This batched codec pipeline divides the chunk batches into batches of a configurable + batch size ("mini-batch"). Fetching, decoding, encoding and storing are performed in + lock step for each mini-batch. Multiple mini-batches are processing concurrently. + """ + + array_array_codecs: tuple[ArrayArrayCodec, ...] + array_bytes_codec: ArrayBytesCodec + bytes_bytes_codecs: tuple[BytesBytesCodec, ...] + batch_size: int + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + return type(self).from_codecs(c.evolve_from_array_spec(array_spec=array_spec) for c in self) + + @classmethod + def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) -> Self: + array_array_codecs, array_bytes_codec, bytes_bytes_codecs = codecs_from_list(codecs) + + return cls( + array_array_codecs=array_array_codecs, + array_bytes_codec=array_bytes_codec, + bytes_bytes_codecs=bytes_bytes_codecs, + batch_size=batch_size or config.get("codec_pipeline.batch_size"), + ) + + @property + def supports_partial_decode(self) -> bool: + """Determines whether the codec pipeline supports partial decoding. + + Currently, only codec pipelines with a single ArrayBytesCodec that supports + partial decoding can support partial decoding. This limitation is due to the fact + that ArrayArrayCodecs can change the slice selection leading to non-contiguous + slices and BytesBytesCodecs can change the chunk bytes in a way that slice + selections cannot be attributed to byte ranges anymore which renders partial + decoding infeasible. + + This limitation may softened in the future.""" + return (len(self.array_array_codecs) + len(self.bytes_bytes_codecs)) == 0 and isinstance( + self.array_bytes_codec, ArrayBytesCodecPartialDecodeMixin + ) + + @property + def supports_partial_encode(self) -> bool: + """Determines whether the codec pipeline supports partial encoding. + + Currently, only codec pipelines with a single ArrayBytesCodec that supports + partial encoding can support partial encoding. This limitation is due to the fact + that ArrayArrayCodecs can change the slice selection leading to non-contiguous + slices and BytesBytesCodecs can change the chunk bytes in a way that slice + selections cannot be attributed to byte ranges anymore which renders partial + encoding infeasible. + + This limitation may softened in the future.""" + return (len(self.array_array_codecs) + len(self.bytes_bytes_codecs)) == 0 and isinstance( + self.array_bytes_codec, ArrayBytesCodecPartialEncodeMixin + ) + + def __iter__(self) -> Iterator[Codec]: + yield from self.array_array_codecs + yield self.array_bytes_codec + yield from self.bytes_bytes_codecs + + def validate( + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGrid, + ) -> None: + for codec in self: + codec.validate(shape=shape, dtype=dtype, chunk_grid=chunk_grid) + + def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int: + for codec in self: + byte_length = codec.compute_encoded_size(byte_length, array_spec) + array_spec = codec.resolve_metadata(array_spec) + return byte_length + + def _codecs_with_resolved_metadata_batched( + self, chunk_specs: Iterable[ArraySpec] + ) -> tuple[ + list[tuple[ArrayArrayCodec, list[ArraySpec]]], + tuple[ArrayBytesCodec, list[ArraySpec]], + list[tuple[BytesBytesCodec, list[ArraySpec]]], + ]: + aa_codecs_with_spec: list[tuple[ArrayArrayCodec, list[ArraySpec]]] = [] + chunk_specs = list(chunk_specs) + for aa_codec in self.array_array_codecs: + aa_codecs_with_spec.append((aa_codec, chunk_specs)) + chunk_specs = [aa_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs] + + ab_codec_with_spec = (self.array_bytes_codec, chunk_specs) + chunk_specs = [ + self.array_bytes_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs + ] + + bb_codecs_with_spec: list[tuple[BytesBytesCodec, list[ArraySpec]]] = [] + for bb_codec in self.bytes_bytes_codecs: + bb_codecs_with_spec.append((bb_codec, chunk_specs)) + chunk_specs = [bb_codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs] + + return (aa_codecs_with_spec, ab_codec_with_spec, bb_codecs_with_spec) + + async def decode_batch( + self, + chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[NDBuffer | None]: + chunk_bytes_batch: Iterable[Buffer | None] + chunk_bytes_batch, chunk_specs = _unzip2(chunk_bytes_and_specs) + ( + aa_codecs_with_spec, + ab_codec_with_spec, + bb_codecs_with_spec, + ) = self._codecs_with_resolved_metadata_batched(chunk_specs) + + for bb_codec, chunk_spec_batch in bb_codecs_with_spec[::-1]: + chunk_bytes_batch = await bb_codec.decode( + zip(chunk_bytes_batch, chunk_spec_batch, strict=False) + ) + + ab_codec, chunk_spec_batch = ab_codec_with_spec + chunk_array_batch = await ab_codec.decode( + zip(chunk_bytes_batch, chunk_spec_batch, strict=False) + ) + + for aa_codec, chunk_spec_batch in aa_codecs_with_spec[::-1]: + chunk_array_batch = await aa_codec.decode( + zip(chunk_array_batch, chunk_spec_batch, strict=False) + ) + + return chunk_array_batch + + async def decode_partial_batch( + self, + batch_info: Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]], + ) -> Iterable[NDBuffer | None]: + assert self.supports_partial_decode + assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialDecodeMixin) + return await self.array_bytes_codec.decode_partial(batch_info) + + async def encode_batch( + self, + chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + chunk_array_batch: Iterable[NDBuffer | None] + chunk_specs: Iterable[ArraySpec] + chunk_array_batch, chunk_specs = _unzip2(chunk_arrays_and_specs) + + for aa_codec in self.array_array_codecs: + chunk_array_batch = await aa_codec.encode( + zip(chunk_array_batch, chunk_specs, strict=False) + ) + chunk_specs = resolve_batched(aa_codec, chunk_specs) + + chunk_bytes_batch = await self.array_bytes_codec.encode( + zip(chunk_array_batch, chunk_specs, strict=False) + ) + chunk_specs = resolve_batched(self.array_bytes_codec, chunk_specs) + + for bb_codec in self.bytes_bytes_codecs: + chunk_bytes_batch = await bb_codec.encode( + zip(chunk_bytes_batch, chunk_specs, strict=False) + ) + chunk_specs = resolve_batched(bb_codec, chunk_specs) + + return chunk_bytes_batch + + async def encode_partial_batch( + self, + batch_info: Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]], + ) -> None: + assert self.supports_partial_encode + assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialEncodeMixin) + await self.array_bytes_codec.encode_partial(batch_info) + + async def read_batch( + self, + batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + out: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> None: + if self.supports_partial_decode: + chunk_array_batch = await self.decode_partial_batch( + [ + (byte_getter, chunk_selection, chunk_spec) + for byte_getter, chunk_spec, chunk_selection, *_ in batch_info + ] + ) + for chunk_array, (_, chunk_spec, _, out_selection, _) in zip( + chunk_array_batch, batch_info, strict=False + ): + if chunk_array is not None: + if drop_axes: + chunk_array = chunk_array.squeeze(axis=drop_axes) + out[out_selection] = chunk_array + else: + out[out_selection] = fill_value_or_default(chunk_spec) + else: + chunk_bytes_batch = await concurrent_map( + [(byte_getter, array_spec.prototype) for byte_getter, array_spec, *_ in batch_info], + lambda byte_getter, prototype: byte_getter.get(prototype), + config.get("async.concurrency"), + ) + chunk_array_batch = await self.decode_batch( + [ + (chunk_bytes, chunk_spec) + for chunk_bytes, (_, chunk_spec, *_) in zip( + chunk_bytes_batch, batch_info, strict=False + ) + ], + ) + for chunk_array, (_, chunk_spec, chunk_selection, out_selection, _) in zip( + chunk_array_batch, batch_info, strict=False + ): + if chunk_array is not None: + tmp = chunk_array[chunk_selection] + if drop_axes: + tmp = tmp.squeeze(axis=drop_axes) + out[out_selection] = tmp + else: + out[out_selection] = fill_value_or_default(chunk_spec) + + def _merge_chunk_array( + self, + existing_chunk_array: NDBuffer | None, + value: NDBuffer, + out_selection: SelectorTuple, + chunk_spec: ArraySpec, + chunk_selection: SelectorTuple, + is_complete_chunk: bool, + drop_axes: tuple[int, ...], + ) -> NDBuffer: + if ( + is_complete_chunk + and value.shape == chunk_spec.shape + # Guard that this is not a partial chunk at the end with is_complete_chunk=True + and value[out_selection].shape == chunk_spec.shape + ): + return value + if existing_chunk_array is None: + chunk_array = chunk_spec.prototype.nd_buffer.create( + shape=chunk_spec.shape, + dtype=chunk_spec.dtype.to_native_dtype(), + order=chunk_spec.order, + fill_value=fill_value_or_default(chunk_spec), + ) + else: + chunk_array = existing_chunk_array.copy() # make a writable copy + if chunk_selection == () or is_scalar( + value.as_ndarray_like(), chunk_spec.dtype.to_native_dtype() + ): + chunk_value = value + else: + chunk_value = value[out_selection] + # handle missing singleton dimensions + if drop_axes: + item = tuple( + None # equivalent to np.newaxis + if idx in drop_axes + else slice(None) + for idx in range(chunk_spec.ndim) + ) + chunk_value = chunk_value[item] + chunk_array[chunk_selection] = chunk_value + return chunk_array + + async def write_batch( + self, + batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + value: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> None: + if self.supports_partial_encode: + # Pass scalar values as is + if len(value.shape) == 0: + await self.encode_partial_batch( + [ + (byte_setter, value, chunk_selection, chunk_spec) + for byte_setter, chunk_spec, chunk_selection, out_selection, _ in batch_info + ], + ) + else: + await self.encode_partial_batch( + [ + (byte_setter, value[out_selection], chunk_selection, chunk_spec) + for byte_setter, chunk_spec, chunk_selection, out_selection, _ in batch_info + ], + ) + + else: + # Read existing bytes if not total slice + async def _read_key( + byte_setter: ByteSetter | None, prototype: BufferPrototype + ) -> Buffer | None: + if byte_setter is None: + return None + return await byte_setter.get(prototype=prototype) + + chunk_bytes_batch: Iterable[Buffer | None] + chunk_bytes_batch = await concurrent_map( + [ + ( + None if is_complete_chunk else byte_setter, + chunk_spec.prototype, + ) + for byte_setter, chunk_spec, chunk_selection, _, is_complete_chunk in batch_info + ], + _read_key, + config.get("async.concurrency"), + ) + chunk_array_decoded = await self.decode_batch( + [ + (chunk_bytes, chunk_spec) + for chunk_bytes, (_, chunk_spec, *_) in zip( + chunk_bytes_batch, batch_info, strict=False + ) + ], + ) + + chunk_array_merged = [ + self._merge_chunk_array( + chunk_array, + value, + out_selection, + chunk_spec, + chunk_selection, + is_complete_chunk, + drop_axes, + ) + for chunk_array, ( + _, + chunk_spec, + chunk_selection, + out_selection, + is_complete_chunk, + ) in zip(chunk_array_decoded, batch_info, strict=False) + ] + chunk_array_batch: list[NDBuffer | None] = [] + for chunk_array, (_, chunk_spec, *_) in zip( + chunk_array_merged, batch_info, strict=False + ): + if chunk_array is None: + chunk_array_batch.append(None) # type: ignore[unreachable] + else: + if not chunk_spec.config.write_empty_chunks and chunk_array.all_equal( + fill_value_or_default(chunk_spec) + ): + chunk_array_batch.append(None) + else: + chunk_array_batch.append(chunk_array) + + chunk_bytes_batch = await self.encode_batch( + [ + (chunk_array, chunk_spec) + for chunk_array, (_, chunk_spec, *_) in zip( + chunk_array_batch, batch_info, strict=False + ) + ], + ) + + async def _write_key(byte_setter: ByteSetter, chunk_bytes: Buffer | None) -> None: + if chunk_bytes is None: + await byte_setter.delete() + else: + await byte_setter.set(chunk_bytes) + + await concurrent_map( + [ + (byte_setter, chunk_bytes) + for chunk_bytes, (byte_setter, *_) in zip( + chunk_bytes_batch, batch_info, strict=False + ) + ], + _write_key, + config.get("async.concurrency"), + ) + + async def decode( + self, + chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[NDBuffer | None]: + output: list[NDBuffer | None] = [] + for batch_info in batched(chunk_bytes_and_specs, self.batch_size): + output.extend(await self.decode_batch(batch_info)) + return output + + async def encode( + self, + chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + output: list[Buffer | None] = [] + for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size): + output.extend(await self.encode_batch(single_batch_info)) + return output + + async def read( + self, + batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + out: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> None: + await concurrent_map( + [ + (single_batch_info, out, drop_axes) + for single_batch_info in batched(batch_info, self.batch_size) + ], + self.read_batch, + config.get("async.concurrency"), + ) + + async def write( + self, + batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + value: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> None: + await concurrent_map( + [ + (single_batch_info, value, drop_axes) + for single_batch_info in batched(batch_info, self.batch_size) + ], + self.write_batch, + config.get("async.concurrency"), + ) + + +def codecs_from_list( + codecs: Iterable[Codec], +) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]: + from zarr.codecs.sharding import ShardingCodec + + array_array: tuple[ArrayArrayCodec, ...] = () + array_bytes_maybe: ArrayBytesCodec | None = None + bytes_bytes: tuple[BytesBytesCodec, ...] = () + + if any(isinstance(codec, ShardingCodec) for codec in codecs) and len(tuple(codecs)) > 1: + warn( + "Combining a `sharding_indexed` codec disables partial reads and " + "writes, which may lead to inefficient performance.", + category=ZarrUserWarning, + stacklevel=3, + ) + + for prev_codec, cur_codec in pairwise((None, *codecs)): + if isinstance(cur_codec, ArrayArrayCodec): + if isinstance(prev_codec, ArrayBytesCodec | BytesBytesCodec): + msg = ( + f"Invalid codec order. ArrayArrayCodec {cur_codec}" + "must be preceded by another ArrayArrayCodec. " + f"Got {type(prev_codec)} instead." + ) + raise TypeError(msg) + array_array += (cur_codec,) + + elif isinstance(cur_codec, ArrayBytesCodec): + if isinstance(prev_codec, BytesBytesCodec): + msg = ( + f"Invalid codec order. ArrayBytes codec {cur_codec}" + f" must be preceded by an ArrayArrayCodec. Got {type(prev_codec)} instead." + ) + raise TypeError(msg) + + if array_bytes_maybe is not None: + msg = ( + f"Got two instances of ArrayBytesCodec: {array_bytes_maybe} and {cur_codec}. " + "Only one array-to-bytes codec is allowed." + ) + raise ValueError(msg) + + array_bytes_maybe = cur_codec + + elif isinstance(cur_codec, BytesBytesCodec): + if isinstance(prev_codec, ArrayArrayCodec): + msg = ( + f"Invalid codec order. BytesBytesCodec {cur_codec}" + "must be preceded by either another BytesBytesCodec, or an ArrayBytesCodec. " + f"Got {type(prev_codec)} instead." + ) + bytes_bytes += (cur_codec,) + else: + raise TypeError + + if array_bytes_maybe is None: + raise ValueError("Required ArrayBytesCodec was not found.") + else: + return array_array, array_bytes_maybe, bytes_bytes + + +register_pipeline(BatchedCodecPipeline) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/common.py b/.venv/lib/python3.11/site-packages/zarr/core/common.py new file mode 100644 index 0000000000000000000000000000000000000000..275d062ebaea36a69c8a692ee1fc753c467ef9d5 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/common.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import asyncio +import functools +import math +import operator +import warnings +from collections.abc import Iterable, Mapping, Sequence +from enum import Enum +from itertools import starmap +from typing import ( + TYPE_CHECKING, + Any, + Final, + Generic, + Literal, + NotRequired, + TypedDict, + TypeVar, + cast, + overload, +) + +import numpy as np +from typing_extensions import ReadOnly + +from zarr.core.config import config as zarr_config +from zarr.errors import ZarrRuntimeWarning + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Iterator + + +ZARR_JSON = "zarr.json" +ZARRAY_JSON = ".zarray" +ZGROUP_JSON = ".zgroup" +ZATTRS_JSON = ".zattrs" +ZMETADATA_V2_JSON = ".zmetadata" + +BytesLike = bytes | bytearray | memoryview +ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any] +# For backwards compatibility +ChunkCoords = tuple[int, ...] +ZarrFormat = Literal[2, 3] +NodeType = Literal["array", "group"] +JSON = str | int | float | bool | Mapping[str, "JSON"] | Sequence["JSON"] | None +MemoryOrder = Literal["C", "F"] +AccessModeLiteral = Literal["r", "r+", "a", "w", "w-"] +ANY_ACCESS_MODE: Final = "r", "r+", "a", "w", "w-" +DimensionNames = Iterable[str | None] | None + +TName = TypeVar("TName", bound=str) +TConfig = TypeVar("TConfig", bound=Mapping[str, object]) + + +class NamedConfig(TypedDict, Generic[TName, TConfig]): + """ + A typed dictionary representing an object with a name and configuration, where the configuration + is an optional mapping of string keys to values, e.g. another typed dictionary or a JSON object. + + This class is generic with two type parameters: the type of the name (``TName``) and the type of + the configuration (``TConfig``). + """ + + name: ReadOnly[TName] + """The name of the object.""" + + configuration: NotRequired[ReadOnly[TConfig]] + """The configuration of the object. Not required.""" + + +class NamedRequiredConfig(TypedDict, Generic[TName, TConfig]): + """ + A typed dictionary representing an object with a name and configuration, where the configuration + is a mapping of string keys to values, e.g. another typed dictionary or a JSON object. + + This class is generic with two type parameters: the type of the name (``TName``) and the type of + the configuration (``TConfig``). + """ + + name: ReadOnly[TName] + """The name of the object.""" + + configuration: ReadOnly[TConfig] + """The configuration of the object.""" + + +def product(tup: tuple[int, ...]) -> int: + return functools.reduce(operator.mul, tup, 1) + + +def ceildiv(a: float, b: float) -> int: + if a == 0: + return 0 + return math.ceil(a / b) + + +T = TypeVar("T", bound=tuple[Any, ...]) +V = TypeVar("V") + + +async def concurrent_map( + items: Iterable[T], + func: Callable[..., Awaitable[V]], + limit: int | None = None, +) -> list[V]: + if limit is None: + return await asyncio.gather(*list(starmap(func, items))) + + else: + sem = asyncio.Semaphore(limit) + + async def run(item: tuple[Any]) -> V: + async with sem: + return await func(*item) + + return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items]) + + +E = TypeVar("E", bound=Enum) + + +def enum_names(enum: type[E]) -> Iterator[str]: + for item in enum: + yield item.name + + +def parse_enum(data: object, cls: type[E]) -> E: + if isinstance(data, cls): + return data + if not isinstance(data, str): + raise TypeError(f"Expected str, got {type(data)}") + if data in enum_names(cls): + return cls(data) + raise ValueError(f"Value must be one of {list(enum_names(cls))!r}. Got {data} instead.") + + +def parse_name(data: JSON, expected: str | None = None) -> str: + if isinstance(data, str): + if expected is None or data == expected: + return data + raise ValueError(f"Expected '{expected}'. Got {data} instead.") + else: + raise TypeError(f"Expected a string, got an instance of {type(data)}.") + + +def parse_configuration(data: JSON) -> JSON: + if not isinstance(data, dict): + raise TypeError(f"Expected dict, got {type(data)}") + return data + + +@overload +def parse_named_configuration( + data: JSON | NamedConfig[str, Any], expected_name: str | None = None +) -> tuple[str, dict[str, JSON]]: ... + + +@overload +def parse_named_configuration( + data: JSON | NamedConfig[str, Any], + expected_name: str | None = None, + *, + require_configuration: bool = True, +) -> tuple[str, dict[str, JSON] | None]: ... + + +def parse_named_configuration( + data: JSON | NamedConfig[str, Any], + expected_name: str | None = None, + *, + require_configuration: bool = True, +) -> tuple[str, JSON | None]: + if not isinstance(data, dict): + raise TypeError(f"Expected dict, got {type(data)}") + if "name" not in data: + raise ValueError(f"Named configuration does not have a 'name' key. Got {data}.") + name_parsed = parse_name(data["name"], expected_name) + if "configuration" in data: + configuration_parsed = parse_configuration(data["configuration"]) + elif require_configuration: + raise ValueError(f"Named configuration does not have a 'configuration' key. Got {data}.") + else: + configuration_parsed = None + return name_parsed, configuration_parsed + + +def parse_shapelike(data: ShapeLike) -> tuple[int, ...]: + """ + Parse a shape-like input into an explicit shape. + """ + if isinstance(data, int | np.integer): + if data < 0: + raise ValueError(f"Expected a non-negative integer. Got {data} instead") + return (int(data),) + try: + data_tuple = tuple(data) + except TypeError as e: + msg = f"Expected an integer or an iterable of integers. Got {data} instead." + raise TypeError(msg) from e + + if not all(isinstance(v, int | np.integer) for v in data_tuple): + msg = f"Expected an iterable of integers. Got {data} instead." + raise TypeError(msg) + if not all(v > -1 for v in data_tuple): + msg = f"Expected all values to be non-negative. Got {data} instead." + raise ValueError(msg) + + # cast NumPy scalars to plain python ints + return tuple(int(x) for x in data_tuple) + + +def parse_fill_value(data: Any) -> Any: + # todo: real validation + return data + + +def parse_order(data: Any) -> Literal["C", "F"]: + if data in ("C", "F"): + return cast("Literal['C', 'F']", data) + raise ValueError(f"Expected one of ('C', 'F'), got {data} instead.") + + +def parse_bool(data: Any) -> bool: + if isinstance(data, bool): + return data + raise ValueError(f"Expected bool, got {data} instead.") + + +def _warn_write_empty_chunks_kwarg() -> None: + # TODO: link to docs page on array configuration in this message + msg = ( + "The `write_empty_chunks` keyword argument is deprecated and will be removed in future versions. " + "To control whether empty chunks are written to storage, either use the `config` keyword " + "argument, as in `config={'write_empty_chunks': True}`," + "or change the global 'array.write_empty_chunks' configuration variable." + ) + warnings.warn(msg, ZarrRuntimeWarning, stacklevel=2) + + +def _warn_order_kwarg() -> None: + # TODO: link to docs page on array configuration in this message + msg = ( + "The `order` keyword argument has no effect for Zarr format 3 arrays. " + "To control the memory layout of the array, either use the `config` keyword " + "argument, as in `config={'order': 'C'}`," + "or change the global 'array.order' configuration variable." + ) + warnings.warn(msg, ZarrRuntimeWarning, stacklevel=2) + + +def _default_zarr_format() -> ZarrFormat: + """Return the default zarr_version""" + return cast("ZarrFormat", int(zarr_config.get("default_zarr_format", 3))) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/config.py b/.venv/lib/python3.11/site-packages/zarr/core/config.py new file mode 100644 index 0000000000000000000000000000000000000000..f8f8ea4f5f1b813b195366b7965dfd0443400d95 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/config.py @@ -0,0 +1,153 @@ +""" +The config module is responsible for managing the configuration of zarr and is based on the Donfig python library. +For selecting custom implementations of codecs, pipelines, buffers and ndbuffers, first register the implementations +in the registry and then select them in the config. + +Example: + An implementation of the bytes codec in a class ``your.module.NewBytesCodec`` requires the value of ``codecs.bytes`` + to be ``your.module.NewBytesCodec``. Donfig can be configured programmatically, by environment variables, or from + YAML files in standard locations. + + ```python + from your.module import NewBytesCodec + from zarr.core.config import register_codec, config + + register_codec("bytes", NewBytesCodec) + config.set({"codecs.bytes": "your.module.NewBytesCodec"}) + ``` + + Instead of setting the value programmatically with ``config.set``, you can also set the value with an environment + variable. The environment variable ``ZARR_CODECS__BYTES`` can be set to ``your.module.NewBytesCodec``. The double + underscore ``__`` is used to indicate nested access. + + ```bash + export ZARR_CODECS__BYTES="your.module.NewBytesCodec" + ``` + +For more information, see the Donfig documentation at https://github.com/pytroll/donfig. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, cast + +from donfig import Config as DConfig + +if TYPE_CHECKING: + from donfig.config_obj import ConfigSet + + +class BadConfigError(ValueError): + _msg = "bad Config: %r" + + +class Config(DConfig): # type: ignore[misc] + """The Config will collect configuration from config files and environment variables + + Example environment variables: + Grabs environment variables of the form "ZARR_FOO__BAR_BAZ=123" and + turns these into config variables of the form ``{"foo": {"bar-baz": 123}}`` + It transforms the key and value in the following way: + + - Lower-cases the key text + - Treats ``__`` (double-underscore) as nested access + - Calls ``ast.literal_eval`` on the value + + """ + + def reset(self) -> None: + self.clear() + self.refresh() + + def enable_gpu(self) -> ConfigSet: + """ + Configure Zarr to use GPUs where possible. + """ + return self.set( + {"buffer": "zarr.buffer.gpu.Buffer", "ndbuffer": "zarr.buffer.gpu.NDBuffer"} + ) + + +# these keys were removed from the config as part of the 3.1.0 release. +# these deprecations should be removed in 3.1.1 or thereabouts. +deprecations = { + "array.v2_default_compressor.numeric": None, + "array.v2_default_compressor.string": None, + "array.v2_default_compressor.bytes": None, + "array.v2_default_filters.string": None, + "array.v2_default_filters.bytes": None, + "array.v3_default_filters.numeric": None, + "array.v3_default_filters.raw": None, + "array.v3_default_filters.bytes": None, + "array.v3_default_serializer.numeric": None, + "array.v3_default_serializer.string": None, + "array.v3_default_serializer.bytes": None, + "array.v3_default_compressors.string": None, + "array.v3_default_compressors.bytes": None, + "array.v3_default_compressors": None, +} + +# The default configuration for zarr +config = Config( + "zarr", + defaults=[ + { + "default_zarr_format": 3, + "array": { + "order": "C", + "write_empty_chunks": False, + "target_shard_size_bytes": None, + }, + "async": {"concurrency": 10, "timeout": None}, + "threading": {"max_workers": None}, + "json_indent": 2, + "codec_pipeline": { + "path": "zarr.core.codec_pipeline.BatchedCodecPipeline", + "batch_size": 1, + }, + "codecs": { + "blosc": "zarr.codecs.blosc.BloscCodec", + "gzip": "zarr.codecs.gzip.GzipCodec", + "zstd": "zarr.codecs.zstd.ZstdCodec", + "bytes": "zarr.codecs.bytes.BytesCodec", + "endian": "zarr.codecs.bytes.BytesCodec", # compatibility with earlier versions of ZEP1 + "crc32c": "zarr.codecs.crc32c_.Crc32cCodec", + "sharding_indexed": "zarr.codecs.sharding.ShardingCodec", + "transpose": "zarr.codecs.transpose.TransposeCodec", + "vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec", + "vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec", + "numcodecs.bz2": "zarr.codecs.numcodecs.BZ2", + "numcodecs.crc32": "zarr.codecs.numcodecs.CRC32", + "numcodecs.crc32c": "zarr.codecs.numcodecs.CRC32C", + "numcodecs.lz4": "zarr.codecs.numcodecs.LZ4", + "numcodecs.lzma": "zarr.codecs.numcodecs.LZMA", + "numcodecs.zfpy": "zarr.codecs.numcodecs.ZFPY", + "numcodecs.adler32": "zarr.codecs.numcodecs.Adler32", + "numcodecs.astype": "zarr.codecs.numcodecs.AsType", + "numcodecs.bitround": "zarr.codecs.numcodecs.BitRound", + "numcodecs.blosc": "zarr.codecs.numcodecs.Blosc", + "numcodecs.delta": "zarr.codecs.numcodecs.Delta", + "numcodecs.fixedscaleoffset": "zarr.codecs.numcodecs.FixedScaleOffset", + "numcodecs.fletcher32": "zarr.codecs.numcodecs.Fletcher32", + "numcodecs.gzip": "zarr.codecs.numcodecs.GZip", + "numcodecs.jenkins_lookup3": "zarr.codecs.numcodecs.JenkinsLookup3", + "numcodecs.pcodec": "zarr.codecs.numcodecs.PCodec", + "numcodecs.packbits": "zarr.codecs.numcodecs.PackBits", + "numcodecs.shuffle": "zarr.codecs.numcodecs.Shuffle", + "numcodecs.quantize": "zarr.codecs.numcodecs.Quantize", + "numcodecs.zlib": "zarr.codecs.numcodecs.Zlib", + "numcodecs.zstd": "zarr.codecs.numcodecs.Zstd", + }, + "buffer": "zarr.buffer.cpu.Buffer", + "ndbuffer": "zarr.buffer.cpu.NDBuffer", + } + ], + deprecations=deprecations, +) + + +def parse_indexing_order(data: Any) -> Literal["C", "F"]: + if data in ("C", "F"): + return cast("Literal['C', 'F']", data) + msg = f"Expected one of ('C', 'F'), got {data} instead." + raise ValueError(msg) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/__init__.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1049a2063f12df41fd4bb840d2a884a852fde875 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__init__.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Final, TypeAlias + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeJSON, +) +from zarr.core.dtype.npy.bool import Bool +from zarr.core.dtype.npy.bytes import ( + NullTerminatedBytes, + NullterminatedBytesJSON_V2, + NullTerminatedBytesJSON_V3, + RawBytes, + RawBytesJSON_V2, + RawBytesJSON_V3, + VariableLengthBytes, + VariableLengthBytesJSON_V2, +) +from zarr.core.dtype.npy.complex import Complex64, Complex128 +from zarr.core.dtype.npy.float import Float16, Float32, Float64 +from zarr.core.dtype.npy.int import Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64 +from zarr.core.dtype.npy.structured import Structured, StructuredJSON_V2, StructuredJSON_V3 +from zarr.core.dtype.npy.time import ( + DateTime64, + DateTime64JSON_V2, + DateTime64JSON_V3, + TimeDelta64, + TimeDelta64JSON_V2, + TimeDelta64JSON_V3, +) + +if TYPE_CHECKING: + from zarr.core.common import ZarrFormat + +from collections.abc import Mapping + +import numpy as np +import numpy.typing as npt + +from zarr.core.common import JSON +from zarr.core.dtype.npy.string import ( + FixedLengthUTF32, + FixedLengthUTF32JSON_V2, + FixedLengthUTF32JSON_V3, + VariableLengthUTF8, + VariableLengthUTF8JSON_V2, +) +from zarr.core.dtype.registry import DataTypeRegistry +from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + +__all__ = [ + "Bool", + "Complex64", + "Complex128", + "DataTypeRegistry", + "DataTypeValidationError", + "DateTime64", + "DateTime64JSON_V2", + "DateTime64JSON_V3", + "FixedLengthUTF32", + "FixedLengthUTF32JSON_V2", + "FixedLengthUTF32JSON_V3", + "Float16", + "Float32", + "Float64", + "Int8", + "Int16", + "Int32", + "Int64", + "NullTerminatedBytes", + "NullTerminatedBytesJSON_V3", + "NullterminatedBytesJSON_V2", + "RawBytes", + "RawBytesJSON_V2", + "RawBytesJSON_V3", + "Structured", + "StructuredJSON_V2", + "StructuredJSON_V3", + "TBaseDType", + "TBaseScalar", + "TimeDelta64", + "TimeDelta64JSON_V2", + "TimeDelta64JSON_V3", + "UInt8", + "UInt16", + "UInt32", + "UInt64", + "VariableLengthBytes", + "VariableLengthBytesJSON_V2", + "VariableLengthUTF8", + "VariableLengthUTF8JSON_V2", + "ZDType", + "data_type_registry", + "parse_data_type", + "parse_dtype", +] + +data_type_registry = DataTypeRegistry() + +IntegerDType = Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 +INTEGER_DTYPE: Final = Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64 + +FloatDType = Float16 | Float32 | Float64 +FLOAT_DTYPE: Final = Float16, Float32, Float64 + +ComplexFloatDType = Complex64 | Complex128 +COMPLEX_FLOAT_DTYPE: Final = Complex64, Complex128 + +StringDType = FixedLengthUTF32 | VariableLengthUTF8 +STRING_DTYPE: Final = FixedLengthUTF32, VariableLengthUTF8 + +TimeDType = DateTime64 | TimeDelta64 +TIME_DTYPE: Final = DateTime64, TimeDelta64 + +BytesDType = RawBytes | NullTerminatedBytes | VariableLengthBytes +BYTES_DTYPE: Final = RawBytes, NullTerminatedBytes, VariableLengthBytes + +AnyDType = ( + Bool + | IntegerDType + | FloatDType + | ComplexFloatDType + | StringDType + | BytesDType + | Structured + | TimeDType + | VariableLengthBytes +) +# mypy has trouble inferring the type of variablelengthstring dtype, because its class definition +# depends on the installed numpy version. That's why the type: ignore statement is needed here. +ANY_DTYPE: Final = ( + Bool, + *INTEGER_DTYPE, + *FLOAT_DTYPE, + *COMPLEX_FLOAT_DTYPE, + *STRING_DTYPE, + *BYTES_DTYPE, + Structured, + *TIME_DTYPE, + VariableLengthBytes, +) + +# These are aliases for variable-length UTF-8 strings +# We handle them when a user requests a data type instead of using NumPy's dtype inferece because +# the default NumPy behavior -- to inspect the user-provided array data and choose +# an appropriately sized U dtype -- is unworkable for Zarr. +VLEN_UTF8_ALIAS: Final = ("str", str, "string") + +# This type models inputs that can be coerced to a ZDType +ZDTypeLike: TypeAlias = npt.DTypeLike | ZDType[TBaseDType, TBaseScalar] | Mapping[str, JSON] | str + +for dtype in ANY_DTYPE: + # mypy does not know that all the elements of ANY_DTYPE are subclasses of ZDType + data_type_registry.register(dtype._zarr_v3_name, dtype) # type: ignore[arg-type] + + +# TODO: find a better name for this function +def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, TBaseScalar]: + """ + Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype. + """ + if not isinstance(dtype, np.dtype): + na_dtype: np.dtype[np.generic] + if isinstance(dtype, list): + # this is a valid _VoidDTypeLike check + na_dtype = np.dtype([tuple(d) for d in dtype]) + else: + na_dtype = np.dtype(dtype) + else: + na_dtype = dtype + return data_type_registry.match_dtype(dtype=na_dtype) + + +def get_data_type_from_json( + dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat +) -> ZDType[TBaseDType, TBaseScalar]: + """ + Given a JSON representation of a data type and a Zarr format version, + attempt to create a ZDType instance from the registered ZDType classes. + """ + return data_type_registry.match_json(dtype_spec, zarr_format=zarr_format) + + +def parse_data_type( + dtype_spec: ZDTypeLike, + *, + zarr_format: ZarrFormat, +) -> ZDType[TBaseDType, TBaseScalar]: + """ + Interpret the input as a ZDType. + + This function wraps ``parse_dtype``. The only difference is the function name. This function may + be deprecated in a future version of Zarr Python in favor of ``parse_dtype``. + + Parameters + ---------- + dtype_spec : ZDTypeLike + The input to be interpreted as a ZDType. This could be a ZDType, which will be returned + directly, or a JSON representation of a ZDType, or a native dtype, or a python object that + can be converted into a native dtype. + zarr_format : ZarrFormat + The Zarr format version. This parameter is required because this function will attempt to + parse the JSON representation of a data type, and the JSON representation of data types + varies between Zarr 2 and Zarr 3. + + Returns + ------- + ZDType[TBaseDType, TBaseScalar] + The ZDType corresponding to the input. + + Examples + -------- + ```python + from zarr.dtype import parse_data_type + import numpy as np + parse_data_type("int32", zarr_format=2) + # Int32(endianness='little') + parse_data_type(np.dtype('S10'), zarr_format=2) + # NullTerminatedBytes(length=10) + parse_data_type({"name": "numpy.datetime64", "configuration": {"unit": "s", "scale_factor": 10}}, zarr_format=3) + # DateTime64(endianness='little', scale_factor=10, unit='s') + ``` + """ + return parse_dtype(dtype_spec, zarr_format=zarr_format) + + +def parse_dtype( + dtype_spec: ZDTypeLike, + *, + zarr_format: ZarrFormat, +) -> ZDType[TBaseDType, TBaseScalar]: + """ + Convert the input as a ZDType. + + Parameters + ---------- + dtype_spec : ZDTypeLike + The input to be converted to a ZDType. This could be a ZDType, which will be returned + directly, or a JSON representation of a ZDType, or a numpy dtype, or a python object that + can be converted into a native dtype. + zarr_format : ZarrFormat + The Zarr format version. This parameter is required because this function will attempt to + parse the JSON representation of a data type, and the JSON representation of data types + varies between Zarr 2 and Zarr 3. + + Returns + ------- + ZDType[TBaseDType, TBaseScalar] + The ZDType corresponding to the input. + + Examples + -------- + ```python + from zarr.dtype import parse_dtype + import numpy as np + parse_dtype("int32", zarr_format=2) + # Int32(endianness='little') + parse_dtype(np.dtype('S10'), zarr_format=2) + # NullTerminatedBytes(length=10) + parse_dtype({"name": "numpy.datetime64", "configuration": {"unit": "s", "scale_factor": 10}}, zarr_format=3) + # DateTime64(endianness='little', scale_factor=10, unit='s') + ``` + """ + if isinstance(dtype_spec, ZDType): + return dtype_spec + # First attempt to interpret the input as JSON + if isinstance(dtype_spec, Mapping | str | Sequence): + try: + return get_data_type_from_json(dtype_spec, zarr_format=zarr_format) # type: ignore[arg-type] + except ValueError: + # no data type matched this JSON-like input + pass + if dtype_spec in VLEN_UTF8_ALIAS: + # If the dtype request is one of the aliases for variable-length UTF-8 strings, + # return that dtype. + return VariableLengthUTF8() # type: ignore[return-value] + # otherwise, we have either a numpy dtype string, or a zarr v3 dtype string, and in either case + # we can create a native dtype from it, and do the dtype inference from that + return get_data_type_from_native_dtype(dtype_spec) # type: ignore[arg-type] diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7abcff771003ebfe065ac563fd2e240c310d4ec Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/common.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/common.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24c76179494f0369e291827de57759fe5011ce1f Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/common.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/registry.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/registry.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86c71e6d6335cbbf47d6b6fbf5c6fd8ab0cba36e Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/registry.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/wrapper.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/wrapper.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93530753c1ac061729e79878354661395b27720e Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/__pycache__/wrapper.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/common.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/common.py new file mode 100644 index 0000000000000000000000000000000000000000..6b70f595ba547c17e88965155f783f9198e28ebe --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/common.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import warnings +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import ( + ClassVar, + Final, + Generic, + Literal, + TypedDict, + TypeGuard, + TypeVar, +) + +from typing_extensions import ReadOnly + +from zarr.core.common import NamedConfig +from zarr.errors import UnstableSpecificationWarning + +EndiannessStr = Literal["little", "big"] +ENDIANNESS_STR: Final = "little", "big" + +SpecialFloatStrings = Literal["NaN", "Infinity", "-Infinity"] +SPECIAL_FLOAT_STRINGS: Final = ("NaN", "Infinity", "-Infinity") + +JSONFloatV2 = float | SpecialFloatStrings +JSONFloatV3 = float | SpecialFloatStrings | str + +ObjectCodecID = Literal["vlen-utf8", "vlen-bytes", "vlen-array", "pickle", "json2", "msgpack2"] +# These are the ids of the known object codecs for zarr v2. +OBJECT_CODEC_IDS: Final = ("vlen-utf8", "vlen-bytes", "vlen-array", "pickle", "json2", "msgpack2") + +# This is a wider type than our standard JSON type because we need +# to work with typeddict objects which are assignable to Mapping[str, object] +DTypeJSON = str | int | float | Sequence["DTypeJSON"] | None | Mapping[str, object] + +# The DTypeJSON_V2 type exists because ZDType.from_json takes a single argument, which must contain +# all the information necessary to decode the data type. Zarr v2 supports multiple distinct +# data types that all used the "|O" data type identifier. These data types can only be +# discriminated on the basis of their "object codec", i.e. a special data type specific +# compressor or filter. So to figure out what data type a zarr v2 array has, we need the +# data type identifier from metadata, as well as an object codec id if the data type identifier +# is "|O". +# So we will pack the name of the dtype alongside the name of the object codec id, if applicable, +# in a single dict, and pass that to the data type inference logic. +# These type variables have a very wide bound because the individual zdtype +# classes can perform a very specific type check. + +# This is the JSON representation of a structured dtype in zarr v2 +StructuredName_V2 = Sequence["str | StructuredName_V2"] + +# This models the type of the name a dtype might have in zarr v2 array metadata +DTypeName_V2 = StructuredName_V2 | str + +TDTypeNameV2_co = TypeVar("TDTypeNameV2_co", bound=DTypeName_V2, covariant=True) +TObjectCodecID_co = TypeVar("TObjectCodecID_co", bound=None | str, covariant=True) + + +class DTypeConfig_V2(TypedDict, Generic[TDTypeNameV2_co, TObjectCodecID_co]): + name: ReadOnly[TDTypeNameV2_co] + object_codec_id: ReadOnly[TObjectCodecID_co] + + +DTypeSpec_V2 = DTypeConfig_V2[DTypeName_V2, None | str] + + +def check_structured_dtype_v2_inner(data: object) -> TypeGuard[StructuredName_V2]: + """ + A type guard for the inner elements of a structured dtype. This is a recursive check because + the type is itself recursive. + + This check ensures that all the elements are 2-element sequences beginning with a string + and ending with either another string or another 2-element sequence beginning with a string and + ending with another instance of that type. + """ + if isinstance(data, (str, Mapping)): + return False + if not isinstance(data, Sequence): + return False + if len(data) != 2: + return False + if not (isinstance(data[0], str)): + return False + if isinstance(data[-1], str): + return True + elif isinstance(data[-1], Sequence): + return check_structured_dtype_v2_inner(data[-1]) + return False + + +def check_structured_dtype_name_v2(data: Sequence[object]) -> TypeGuard[StructuredName_V2]: + """ + Check that all the elements of a sequence are valid zarr v2 structured dtype identifiers + """ + return all(check_structured_dtype_v2_inner(d) for d in data) + + +def check_dtype_name_v2(data: object) -> TypeGuard[DTypeName_V2]: + """ + Type guard for narrowing the type of a python object to a valid zarr v2 dtype name. + """ + if isinstance(data, str): + return True + elif isinstance(data, Sequence): + return check_structured_dtype_name_v2(data) + return False + + +def check_dtype_spec_v2(data: object) -> TypeGuard[DTypeSpec_V2]: + """ + Type guard for narrowing a python object to an instance of DTypeSpec_V2 + """ + if not isinstance(data, Mapping): + return False + if set(data.keys()) != {"name", "object_codec_id"}: + return False + if not check_dtype_name_v2(data["name"]): + return False + return isinstance(data["object_codec_id"], str | None) + + +# By comparison, The JSON representation of a dtype in zarr v3 is much simpler. +# It's either a string, or a structured dict +DTypeSpec_V3 = str | NamedConfig[str, Mapping[str, object]] + + +def check_dtype_spec_v3(data: object) -> TypeGuard[DTypeSpec_V3]: + """ + Type guard for narrowing the type of a python object to an instance of + DTypeSpec_V3, i.e either a string or a dict with a "name" field that's a string and a + "configuration" field that's a mapping with string keys. + """ + if isinstance(data, str) or ( # noqa: SIM103 + isinstance(data, Mapping) + and set(data.keys()) == {"name", "configuration"} + and isinstance(data["configuration"], Mapping) + and all(isinstance(k, str) for k in data["configuration"]) + ): + return True + return False + + +def unpack_dtype_json(data: DTypeSpec_V2 | DTypeSpec_V3) -> DTypeJSON: + """ + Return the array metadata form of the dtype JSON representation. For the Zarr V3 form of dtype + metadata, this is a no-op. For the Zarr V2 form of dtype metadata, this unpacks the dtype name. + """ + if isinstance(data, Mapping) and set(data.keys()) == {"name", "object_codec_id"}: + return data["name"] + return data + + +class DataTypeValidationError(ValueError): ... + + +class ScalarTypeValidationError(ValueError): ... + + +@dataclass(frozen=True, kw_only=True) +class HasLength: + """ + A mix-in class for data types with a length attribute, such as fixed-size collections + of unicode strings, or bytes. + + Attributes + ---------- + length : int + The length of the scalars belonging to this data type. Note that this class does not assign + a unit to the length. Child classes may assign units. + """ + + length: int + + +@dataclass(frozen=True, kw_only=True) +class HasEndianness: + """ + A mix-in class for data types with an endianness attribute + """ + + endianness: EndiannessStr = "little" + + +@dataclass(frozen=True, kw_only=True) +class HasItemSize: + """ + A mix-in class for data types with an item size attribute. + This mix-in bears a property ``item_size``, which denotes the size of each element of the data + type, in bytes. + """ + + @property + def item_size(self) -> int: + raise NotImplementedError + + +@dataclass(frozen=True, kw_only=True) +class HasObjectCodec: + """ + A mix-in class for data types that require an object codec id. + This class bears the property ``object_codec_id``, which is the string name of an object + codec that is required to encode and decode the data type. + + In zarr-python 2.x certain data types like variable-length strings or variable-length arrays + used the catch-all numpy "object" data type for their in-memory representation. But these data + types cannot be stored as numpy object data types, because the object data type does not define + a fixed memory layout. So these data types required a special codec, called an "object codec", + that effectively defined a compact representation for the data type, which was used to encode + and decode the data type. + + Zarr-python 2.x would not allow the creation of arrays with the "object" data type if an object + codec was not specified, and thus the name of the object codec is effectively part of the data + type model. + """ + + object_codec_id: ClassVar[str] + + +def v3_unstable_dtype_warning(dtype: object) -> None: + """ + Emit this warning when a data type does not have a stable zarr v3 spec + """ + msg = ( + f"The data type ({dtype}) does not have a Zarr V3 specification. " + "That means that the representation of arrays saved with this data type may change without " + "warning in a future version of Zarr Python. " + "Arrays stored with this data type may be unreadable by other Zarr libraries. " + "Use this data type at your own risk! " + "Check https://github.com/zarr-developers/zarr-extensions/tree/main/data-types for the " + "status of data type specifications for Zarr V3." + ) + warnings.warn(msg, category=UnstableSpecificationWarning, stacklevel=2) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__init__.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fee23a649a7d04ddbd528b1324338accbd2be852 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bool.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bool.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04d4ae7d29d7904b384dbb0a509623806e3bc256 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bool.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bytes.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bytes.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b29c4a33cb739ee1497252ab87d18d19d1ff0bbf Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/bytes.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/common.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/common.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d44c3411756eaeba6bb03be96339e8b053fdd5e8 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/common.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/complex.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/complex.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d33d6ecf278d01942f660f1aa5e3767494cd6a53 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/complex.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/float.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/float.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d5d62b8b0af52edb501d10b069245cdf87faeda Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/float.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/int.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/int.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d3cbc4a905102b981b5b3c1c81a8ddc7e2025a70 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/int.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/string.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/string.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a41a723cdf68e622544859fedf9fa482f36114ac Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/string.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/structured.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/structured.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ff08f3ae8d89272f29af65e8acda67a3775d91c Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/structured.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/time.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/time.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd9f633c4aac77255976ec347ed493c865e60092 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/__pycache__/time.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bool.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bool.py new file mode 100644 index 0000000000000000000000000000000000000000..3e7f5b72f0037ec117ed7c3224efe000728864b4 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bool.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar, Literal, Self, TypeGuard, overload + +import numpy as np + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasItemSize, + check_dtype_spec_v2, +) +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + + +@dataclass(frozen=True, kw_only=True, slots=True) +class Bool(ZDType[np.dtypes.BoolDType, np.bool_], HasItemSize): + """ + A Zarr data type for arrays containing booleans. + + Wraps the [`np.dtypes.BoolDType`][numpy.dtypes.BoolDType] data type. Scalars for this data type are instances of + [`np.bool_`][numpy.bool_]. + + Attributes + ---------- + + _zarr_v3_name : Literal["bool"] = "bool" + The Zarr v3 name of the dtype. + _zarr_v2_name : ``Literal["|b1"]`` = ``"|b1"`` + The Zarr v2 name of the dtype, which is also a string representation + of the boolean dtype used by NumPy. + dtype_cls : ClassVar[type[np.dtypes.BoolDType]] = np.dtypes.BoolDType + The NumPy dtype class. + + References + ---------- + This class implements the boolean data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding)and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + _zarr_v3_name: ClassVar[Literal["bool"]] = "bool" + _zarr_v2_name: ClassVar[Literal["|b1"]] = "|b1" + dtype_cls = np.dtypes.BoolDType + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of Bool from an instance of np.dtypes.BoolDType. + + Parameters + ---------- + dtype : TBaseDType + The NumPy boolean dtype instance to convert. + + Returns + ------- + Bool + An instance of Bool. + + Raises + ------ + DataTypeValidationError + If the provided dtype is not compatible with this ZDType. + """ + if cls._check_native_dtype(dtype): + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self: Self) -> np.dtypes.BoolDType: + """ + Create a NumPy boolean dtype instance from this ZDType. + + Returns + ------- + np.dtypes.BoolDType + The NumPy boolean dtype. + """ + return self.dtype_cls() + + @classmethod + def _check_json_v2( + cls, + data: DTypeJSON, + ) -> TypeGuard[DTypeConfig_V2[Literal["|b1"], None]]: + """ + Check that the input is a valid JSON representation of a Bool. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + ``TypeGuard[DTypeConfig_V2[Literal["|b1"], None]]`` + True if the input is a valid JSON representation, False otherwise. + """ + return ( + check_dtype_spec_v2(data) + and data["name"] == cls._zarr_v2_name + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["bool"]]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + bool + True if the input is a valid JSON representation, False otherwise. + """ + return data == cls._zarr_v3_name + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of Bool from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Bool + An instance of Bool. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v2(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v2_name!r}" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls: type[Self], data: DTypeJSON) -> Self: + """ + Create an instance of Bool from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Bool + An instance of Bool. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal["|b1"], None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> Literal["bool"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal["|b1"], None] | Literal["bool"]: + """ + Serialize this Bool instance to JSON. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version (2 or 3). + + Returns + ------- + ``DTypeConfig_V2[Literal["|b1"], None] | Literal["bool"]`` + The JSON representation of the Bool instance. + + Raises + ------ + ValueError + If the zarr_format is not 2 or 3. + """ + if zarr_format == 2: + return {"name": self._zarr_v2_name, "object_codec_id": None} + elif zarr_format == 3: + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> bool: + """ + Check if the input can be cast to a boolean scalar. + + Parameters + ---------- + data : object + The data to check. + + Returns + ------- + bool + True if the input can be cast to a boolean scalar, False otherwise. + """ + return True + + def cast_scalar(self, data: object) -> np.bool_: + """ + Cast the input to a numpy boolean scalar. + + Parameters + ---------- + data : object + The data to cast. + + Returns + ------- + bool : np.bool_ + The numpy boolean scalar. + + Raises + ------ + TypeError + If the input cannot be converted to a numpy boolean. + """ + if self._check_scalar(data): + return np.bool_(data) + msg = ( # pragma: no cover + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) # pragma: no cover + + def default_scalar(self) -> np.bool_: + """ + Get the default value for the boolean dtype. + + Returns + ------- + bool : np.bool_ + The default value. + """ + return np.False_ + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> bool: + """ + Convert a scalar to a python bool. + + Parameters + ---------- + data : object + The value to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + bool + The JSON-serializable format. + """ + return bool(data) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.bool_: + """ + Read a JSON-serializable value as a numpy boolean scalar. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + bool : np.bool_ + The numpy boolean scalar. + + Raises + ------ + TypeError + If the input is not a valid boolean type. + """ + if self._check_scalar(data): + return np.bool_(data) + raise TypeError(f"Invalid type: {data}. Expected a boolean.") # pragma: no cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 1 diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bytes.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bytes.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf5985d6947b198fa002e73bde3b08a566c3d48 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/bytes.py @@ -0,0 +1,1279 @@ +from __future__ import annotations + +import base64 +import re +from dataclasses import dataclass +from typing import ClassVar, Literal, Self, TypedDict, TypeGuard, cast, overload + +import numpy as np + +from zarr.core.common import JSON, NamedConfig, ZarrFormat +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasItemSize, + HasLength, + HasObjectCodec, + check_dtype_spec_v2, + v3_unstable_dtype_warning, +) +from zarr.core.dtype.npy.common import check_json_str +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +BytesLike = np.bytes_ | str | bytes | int + + +class FixedLengthBytesConfig(TypedDict): + """ + A configuration for a data type that takes a ``length_bytes`` parameter. + + Attributes + ---------- + + length_bytes : int + The length in bytes of the data associated with this configuration. + + Examples + -------- + ```python + { + "length_bytes": 12 + } + ``` + """ + + length_bytes: int + + +class NullterminatedBytesJSON_V2(DTypeConfig_V2[str, None]): + """ + A wrapper around the JSON representation of the ``NullTerminatedBytes`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + + Examples + -------- + ```python + { + "name": "|S10", + "object_codec_id": None + } + ``` + """ + + +class NullTerminatedBytesJSON_V3( + NamedConfig[Literal["null_terminated_bytes"], FixedLengthBytesConfig] +): + """ + The JSON representation of the ``NullTerminatedBytes`` data type in Zarr V3. + + References + ---------- + This representation is not currently defined in an external specification. + + + Examples + -------- + ```python + { + "name": "null_terminated_bytes", + "configuration": { + "length_bytes": 12 + } + } + ``` + + """ + + +class RawBytesJSON_V2(DTypeConfig_V2[str, None]): + """ + A wrapper around the JSON representation of the ``RawBytes`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + + Examples + -------- + ```python + + { + "name": "|V10", + "object_codec_id": None + } + ``` + """ + + +class RawBytesJSON_V3(NamedConfig[Literal["raw_bytes"], FixedLengthBytesConfig]): + """ + The JSON representation of the ``RawBytes`` data type in Zarr V3. + + References + ---------- + This representation is not currently defined in an external specification. + + + Examples + -------- + ```python + { + "name": "raw_bytes", + "configuration": { + "length_bytes": 12 + } + } + ``` + """ + + +class VariableLengthBytesJSON_V2(DTypeConfig_V2[Literal["|O"], Literal["vlen-bytes"]]): + """ + A wrapper around the JSON representation of the ``VariableLengthBytes`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. The ``object_codec_id`` field is always ``"vlen-bytes"`` + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + Examples + -------- + ```python + { + "name": "|O", + "object_codec_id": "vlen-bytes" + } + ``` + """ + + +@dataclass(frozen=True, kw_only=True) +class NullTerminatedBytes(ZDType[np.dtypes.BytesDType[int], np.bytes_], HasLength, HasItemSize): + """ + A Zarr data type for arrays containing fixed-length null-terminated byte sequences. + + Wraps the [`np.dtypes.BytesDType`][numpy.dtypes.BytesDType] data type. Scalars for this data type are instances of + [`np.bytes_`][numpy.bytes_]. + + This data type is parametrized by an integral length which specifies size in bytes of each + scalar. Because this data type uses null-terminated semantics, indexing into + NumPy arrays with this data type may return fewer than ``length`` bytes. + + Attributes + ---------- + dtype_cls: ClassVar[type[np.dtypes.BytesDType[int]]] = np.dtypes.BytesDType + The NumPy data type wrapped by this ZDType. + _zarr_v3_name : ClassVar[Literal["null_terminated_bytes"]] + length : int + The length of the bytes. + + Notes + ----- + This data type is designed for compatibility with NumPy arrays that use the NumPy ``bytes`` data type. + It may not be desirable for usage outside of that context. If compatibility + with the NumPy ``bytes`` data type is not essential, consider using the ``RawBytes`` + or ``VariableLengthBytes`` data types instead. + """ + + dtype_cls = np.dtypes.BytesDType + _zarr_v3_name: ClassVar[Literal["null_terminated_bytes"]] = "null_terminated_bytes" + + def __post_init__(self) -> None: + """ + We don't allow instances of this class with length less than 1 because there is no way such + a data type can contain actual data. + """ + if self.length < 1: + raise ValueError(f"length must be >= 1, got {self.length}.") + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of NullTerminatedBytes from an instance of np.dtypes.BytesDType. + + This method checks if the provided data type is an instance of np.dtypes.BytesDType. + If so, it returns a new instance of NullTerminatedBytes with a length equal to the + length of input data type. + + Parameters + ---------- + dtype : TBaseDType + The native dtype to convert. + + Returns + ------- + NullTerminatedBytes + An instance of NullTerminatedBytes with the specified length. + + Raises + ------ + DataTypeValidationError + If the dtype is not compatible with NullTerminatedBytes. + """ + + if cls._check_native_dtype(dtype): + return cls(length=dtype.itemsize) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.BytesDType[int]: + """ + Create a NumPy bytes dtype from this NullTerminatedBytes ZDType. + + Returns + ------- + np.dtypes.BytesDType[int] + A NumPy data type object representing null-terminated bytes with a specified length. + """ + + return self.dtype_cls(self.length) + + @classmethod + def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[NullterminatedBytesJSON_V2]: + """ + Check that the input is a valid JSON representation of NullTerminatedBytes in Zarr V2. + + The input data must be a mapping that contains a "name" key that matches the pattern + "|S" and an "object_codec_id" key that is None. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + bool + True if the input data is a valid representation, False otherwise. + """ + + return ( + check_dtype_spec_v2(data) + and isinstance(data["name"], str) + and re.match(r"^\|S\d+$", data["name"]) is not None + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[NullTerminatedBytesJSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[NullTerminatedBytesJSON_V3] + True if the input is a valid representation of this class in Zarr V3, False + otherwise. + """ + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and isinstance(data["configuration"], dict) + and "length_bytes" in data["configuration"] + and isinstance(data["configuration"]["length_bytes"], int) + ) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from Zarr V2-flavored JSON. + + This method checks if the input data is a valid representation of + this class in Zarr V2. If so, it returns a new instance of + this class with a ``length`` as specified in the input data. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class. + """ + + if cls._check_json_v2(data): + name = data["name"] + return cls(length=int(name[2:])) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a string like '|S1', '|S2', etc" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from Zarr V3-flavored JSON. + + This method checks if the input data is a valid representation of + this class in Zarr V3. If so, it returns a new instance of + this class with a ``length`` as specified in the input data. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls(length=data["configuration"]["length_bytes"]) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> NullterminatedBytesJSON_V2: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> NullTerminatedBytesJSON_V3: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[str, None] | NullTerminatedBytesJSON_V3: + """ + Generate a JSON representation of this data type. + + Parameters + ---------- + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + NullterminatedBytesJSON_V2 | NullTerminatedBytesJSON_V3 + The JSON-serializable representation of the data type + """ + if zarr_format == 2: + return {"name": self.to_native_dtype().str, "object_codec_id": None} + elif zarr_format == 3: + v3_unstable_dtype_warning(self) + return { + "name": self._zarr_v3_name, + "configuration": {"length_bytes": self.length}, + } + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[BytesLike]: + """ + Check if the provided data is of type BytesLike. + + This method is used to verify if the input data can be considered as a + scalar of bytes-like type, which includes NumPy bytes, strings, bytes, + and integers. + + Parameters + ---------- + data : object + The data to check. + + Returns + ------- + TypeGuard[BytesLike] + True if the data is bytes-like, False otherwise. + """ + + return isinstance(data, BytesLike) + + def _cast_scalar_unchecked(self, data: BytesLike) -> np.bytes_: + """ + Cast the provided scalar data to [`np.bytes_`][numpy.bytes_], truncating if necessary. + + Parameters + ---------- + data : BytesLike + The data to cast. + + Returns + ------- + bytes : [`np.bytes_`][numpy.bytes_] + The casted data as a NumPy bytes scalar. + + Notes + ----- + This method does not perform any type checking. + The input data must be bytes-like. + """ + if isinstance(data, int): + return self.to_native_dtype().type(str(data)[: self.length]) + else: + return self.to_native_dtype().type(data[: self.length]) + + def cast_scalar(self, data: object) -> np.bytes_: + """ + Attempt to cast a given object to a NumPy bytes scalar. + + This method first checks if the provided data is a valid scalar that can be + converted to a NumPy bytes scalar. If the check succeeds, the unchecked casting + operation is performed. If the data is not valid, a TypeError is raised. + + Parameters + ---------- + data : object + The data to be cast to a NumPy bytes scalar. + + Returns + ------- + bytes : [`np.bytes_`][numpy.bytes_] + The data cast as a NumPy bytes scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a NumPy bytes scalar. + """ + + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> np.bytes_: + """ + Return a default scalar value, which for this data type is an empty byte string. + + Returns + ------- + bytes : [`np.bytes_`][numpy.bytes_] + The default scalar value. + """ + return np.bytes_(b"") + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert a scalar to a JSON-serializable string representation. + + This method encodes the given scalar as a NumPy bytes scalar and then + encodes the bytes as a base64-encoded string. + + Parameters + ---------- + data : object + The scalar to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + str + A string representation of the scalar. + """ + as_bytes = self.cast_scalar(data) + return base64.standard_b64encode(as_bytes).decode("ascii") + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.bytes_: + """ + Read a JSON-serializable value as [`np.bytes_`][numpy.bytes_]. + + Parameters + ---------- + data : JSON + The JSON-serializable base64-encoded string. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + bytes : [`np.bytes_`][numpy.bytes_] + The NumPy bytes scalar obtained from decoding the base64 string. + + Raises + ------ + TypeError + If the input data is not a base64-encoded string. + """ + + if check_json_str(data): + return self.to_native_dtype().type(base64.standard_b64decode(data.encode("ascii"))) + raise TypeError( + f"Invalid type: {data}. Expected a base64-encoded string." + ) # pragma: no cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return self.length + + +@dataclass(frozen=True, kw_only=True) +class RawBytes(ZDType[np.dtypes.VoidDType[int], np.void], HasLength, HasItemSize): + """ + A Zarr data type for arrays containing fixed-length sequences of raw bytes. + + Wraps the NumPy ``void`` data type. Scalars for this data type are instances of [`np.void`][numpy.void]. + + This data type is parametrized by an integral length which specifies size in bytes of each + scalar belonging to this data type. + + Attributes + ---------- + dtype_cls: ClassVar[type[np.dtypes.VoidDType[int]]] = np.dtypes.VoidDtype + The NumPy data type wrapped by this ZDType. + _zarr_v3_name : ClassVar[Literal["raw_bytes"]] + length : int + The length of the bytes. + + Notes + ----- + Although the NumPy "Void" data type is used to create "structured" data types in NumPy, this + class does not support structured data types. + + See the ``Structured`` data type for this functionality. + + """ + + # np.dtypes.VoidDType is specified in an odd way in NumPy + # it cannot be used to create instances of the dtype + # so we have to tell mypy to ignore this here + dtype_cls = np.dtypes.VoidDType # type: ignore[assignment] + _zarr_v3_name: ClassVar[Literal["raw_bytes"]] = "raw_bytes" + + def __post_init__(self) -> None: + """ + We don't allow instances of this class with length less than 1 because there is no way such + a data type can contain actual data. + """ + if self.length < 1: + raise ValueError(f"length must be >= 1, got {self.length}.") + + @classmethod + def _check_native_dtype( + cls: type[Self], dtype: TBaseDType + ) -> TypeGuard[np.dtypes.VoidDType[int]]: + """ + Check that the input is a NumPy void dtype with no fields. + + + Numpy void dtype comes in two forms: + * If the ``fields`` attribute is ``None``, then the dtype represents N raw bytes. + * If the ``fields`` attribute is not ``None``, then the dtype represents a structured dtype, + + In this check we ensure that ``fields`` is ``None``. + + Parameters + ---------- + dtype : TDBaseDType + The dtype to check. + + Returns + ------- + Bool + True if the dtype is an instance of np.dtypes.VoidDType with no fields, False otherwise. + """ + return cls.dtype_cls is type(dtype) and dtype.fields is None + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of RawBytes from an instance of np.dtypes.VoidDType. + + This method checks if the provided data type is compatible with RawBytes. The input + must be an instance of np.dtypes.VoidDType, and have no fields. If the input is compatible, + this method returns an instance of RawBytes with the specified length. + + + Parameters + ---------- + dtype : TBaseDType + The native dtype to convert. + + Returns + ------- + RawBytes + An instance of RawBytes with the specified length. + + Raises + ------ + DataTypeValidationError + If the dtype is not compatible with RawBytes. + """ + + if cls._check_native_dtype(dtype): + return cls(length=dtype.itemsize) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.VoidDType[int]: + """ + Create a NumPy void dtype from this RawBytes ZDType. + + Returns + ------- + np.dtypes.VoidDType[int] + A NumPy data type object representing raw bytes with a specified length. + """ + # Numpy does not allow creating a void type + # by invoking np.dtypes.VoidDType directly + return cast("np.dtypes.VoidDType[int]", np.dtype(f"V{self.length}")) + + @classmethod + def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[RawBytesJSON_V2]: + """ + Check that the input is a valid representation of this class in Zarr V2. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + True if the input is a valid representation of this class in Zarr V3, False otherwise. + + """ + return ( + check_dtype_spec_v2(data) + and isinstance(data["name"], str) + and re.match(r"^\|V\d+$", data["name"]) is not None + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[RawBytesJSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[RawBytesJSON_V3] + True if the input is a valid representation of this class in Zarr V3, False + otherwise. + """ + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and isinstance(data["configuration"], dict) + and set(data["configuration"].keys()) == {"length_bytes"} + and isinstance(data["configuration"]["length_bytes"], int) + ) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of RawBytes from Zarr V2-flavored JSON. + + This method checks if the input data is a valid representation of + RawBytes in Zarr V2. If so, it returns a new instance of + RawBytes with a ``length`` as specified in the input data. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class. + """ + if cls._check_json_v2(data): + name = data["name"] + return cls(length=int(name[2:])) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a string like '|V1', '|V2', etc" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of RawBytes from Zarr V3-flavored JSON. + + This method checks if the input data is a valid representation of + RawBytes in Zarr V3. If so, it returns a new instance of + RawBytes with a ``length`` as specified in the input data. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + RawBytes + An instance of RawBytes. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls(length=data["configuration"]["length_bytes"]) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> RawBytesJSON_V2: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> RawBytesJSON_V3: ... + + def to_json(self, zarr_format: ZarrFormat) -> RawBytesJSON_V2 | RawBytesJSON_V3: + """ + Generate a JSON representation of this data type. + + Parameters + ---------- + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + RawBytesJSON_V2 | RawBytesJSON_V3 + The JSON-serializable representation of the data type. + """ + if zarr_format == 2: + return {"name": self.to_native_dtype().str, "object_codec_id": None} + elif zarr_format == 3: + v3_unstable_dtype_warning(self) + return {"name": self._zarr_v3_name, "configuration": {"length_bytes": self.length}} + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[np.bytes_ | str | bytes | np.void]: + """ + Check if the provided data can be cast to np.void. + + This method is used to verify if the input data can be considered as a + scalar of bytes-like type, which includes np.bytes_, np.void, strings, and bytes objects. + + Parameters + ---------- + data : object + The data to check. + + Returns + ------- + TypeGuard[np.bytes_ | str | bytes | np.void] + True if the data is void-scalar-like, False otherwise. + """ + return isinstance(data, np.bytes_ | str | bytes | np.void) + + def _cast_scalar_unchecked(self, data: object) -> np.void: + """ + Cast the provided scalar data to np.void. + + Parameters + ---------- + data : BytesLike + The data to cast. + + Returns + ------- + np.void + The casted data as a NumPy void scalar. + + Notes + ----- + This method does not perform any type checking. + The input data must be castable to np.void. + """ + native_dtype = self.to_native_dtype() + # Without the second argument, NumPy will return a void scalar for dtype V1. + # The second argument ensures that, if native_dtype is something like V10, + # the result will actually be a V10 scalar. + return native_dtype.type(data, native_dtype) + + def cast_scalar(self, data: object) -> np.void: + """ + Attempt to cast a given object to a NumPy void scalar. + + This method first checks if the provided data is a valid scalar that can be + converted to a NumPy void scalar. If the check succeeds, the unchecked casting + operation is performed. If the data is not valid, a TypeError is raised. + + Parameters + ---------- + data : object + The data to be cast to a NumPy void scalar. + + Returns + ------- + np.void + The data cast as a NumPy void scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a NumPy void scalar. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> np.void: + """ + Return the default scalar value for this data type. + + The default scalar is a NumPy void scalar of the same length as the data type, + filled with zero bytes. + + Returns + ------- + np.void + The default scalar value. + """ + return self.to_native_dtype().type(("\x00" * self.length).encode("ascii")) + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert a scalar to a JSON-serializable string representation. + + This method converts the given scalar to bytes and then + encodes the bytes as a base64-encoded string. + + Parameters + ---------- + data : object + The scalar to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + str + A string representation of the scalar. + """ + as_bytes = self.cast_scalar(data) + return base64.standard_b64encode(as_bytes.tobytes()).decode("ascii") + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.void: + """ + Read a JSON-serializable value as an np.void. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + np.void + The NumPy void scalar. + + Raises + ------ + TypeError + If the data is not a string, or if the string is not a valid base64 encoding. + """ + if check_json_str(data): + return self.to_native_dtype().type(base64.standard_b64decode(data)) + raise TypeError(f"Invalid type: {data}. Expected a string.") # pragma: no cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return self.length + + +@dataclass(frozen=True, kw_only=True) +class VariableLengthBytes(ZDType[np.dtypes.ObjectDType, bytes], HasObjectCodec): + """ + A Zarr data type for arrays containing variable-length sequences of bytes. + + Wraps the NumPy "object" data type. Scalars for this data type are instances of ``bytes``. + + Attributes + ---------- + dtype_cls: ClassVar[type[np.dtypes.ObjectDType]] = np.dtypes.ObjectDType + The NumPy data type wrapped by this ZDType. + _zarr_v3_name: ClassVar[Literal["variable_length_bytes"]] = "variable_length_bytes" + The name of this data type in Zarr V3. + object_codec_id: ClassVar[Literal["vlen-bytes"]] = "vlen-bytes" + The object codec ID for this data type. + + Notes + ----- + Because this data type uses the NumPy "object" data type, it does not guarantee a compact memory + representation of array data. Therefore a "vlen-bytes" codec is needed to ensure that the array + data can be persisted to storage. + """ + + dtype_cls = np.dtypes.ObjectDType + _zarr_v3_name: ClassVar[Literal["variable_length_bytes"]] = "variable_length_bytes" + object_codec_id: ClassVar[Literal["vlen-bytes"]] = "vlen-bytes" + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of VariableLengthBytes from an instance of np.dtypes.ObjectDType. + + This method checks if the provided data type is an instance of np.dtypes.ObjectDType. + If so, it returns an instance of VariableLengthBytes. + + Parameters + ---------- + dtype : TBaseDType + The native dtype to convert. + + Returns + ------- + VariableLengthBytes + An instance of VariableLengthBytes. + + Raises + ------ + DataTypeValidationError + If the dtype is not compatible with VariableLengthBytes. + """ + if cls._check_native_dtype(dtype): + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.ObjectDType: + """ + Create a NumPy object dtype from this VariableLengthBytes ZDType. + + Returns + ------- + np.dtypes.ObjectDType + A NumPy data type object representing variable-length bytes. + """ + return self.dtype_cls() + + @classmethod + def _check_json_v2( + cls, + data: DTypeJSON, + ) -> TypeGuard[VariableLengthBytesJSON_V2]: + """ + Check that the input is a valid JSON representation of a NumPy O dtype, and that the + object codec id is appropriate for variable-length bytes strings. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + True if the input is a valid representation of this class in Zarr V2, False + otherwise. + """ + # Check that the input is a valid JSON representation of a Zarr v2 data type spec. + if not check_dtype_spec_v2(data): + return False + + # Check that the object codec id is appropriate for variable-length bytes strings. + if data["name"] != "|O": + return False + return data["object_codec_id"] == cls.object_codec_id + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["variable_length_bytes"]]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[Literal["variable_length_bytes"]] + True if the input is a valid representation of this class in Zarr V3, False otherwise. + """ + + return data in (cls._zarr_v3_name, "bytes") + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this VariableLengthBytes from Zarr V2-flavored JSON. + + This method checks if the input data is a valid representation of this class + in Zarr V2. If so, it returns a new instance this class. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class class. + """ + + if cls._check_json_v2(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string '|O' and an object_codec_id of {cls.object_codec_id}" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of VariableLengthBytes from Zarr V3-flavored JSON. + + This method checks if the input data is a valid representation of + VariableLengthBytes in Zarr V3. If so, it returns a new instance of + VariableLengthBytes. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + VariableLengthBytes + An instance of VariableLengthBytes. + + Raises + ------ + DataTypeValidationError + If the input data is not a valid representation of this class. + """ + + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> VariableLengthBytesJSON_V2: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> Literal["variable_length_bytes"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> VariableLengthBytesJSON_V2 | Literal["variable_length_bytes"]: + """ + Convert the variable-length bytes data type to a JSON-serializable form. + + Parameters + ---------- + zarr_format : ZarrFormat + The zarr format version. Accepted values are 2 and 3. + + Returns + ------- + ``DTypeConfig_V2[Literal["|O"], Literal["vlen-bytes"]] | Literal["variable_length_bytes"]`` + The JSON-serializable representation of the variable-length bytes data type. + For zarr_format 2, returns a dictionary with "name" and "object_codec_id". + For zarr_format 3, returns a string identifier "variable_length_bytes". + + Raises + ------ + ValueError + If zarr_format is not 2 or 3. + """ + + if zarr_format == 2: + return {"name": "|O", "object_codec_id": self.object_codec_id} + elif zarr_format == 3: + v3_unstable_dtype_warning(self) + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def default_scalar(self) -> bytes: + """ + Return the default scalar value for the variable-length bytes data type. + + Returns + ------- + bytes + The default scalar value, which is an empty byte string. + """ + + return b"" + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert a scalar to a JSON-serializable string representation. + + This method encodes the given scalar as bytes and then + encodes the bytes as a base64-encoded string. + + Parameters + ---------- + data : object + The scalar to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + str + A string representation of the scalar. + """ + return base64.standard_b64encode(data).decode("ascii") # type: ignore[arg-type] + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> bytes: + """ + Decode a base64-encoded JSON string to bytes. + + Parameters + ---------- + data : JSON + The JSON-serializable base64-encoded string. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + bytes + The decoded bytes from the base64 string. + + Raises + ------ + TypeError + If the input data is not a base64-encoded string. + """ + + if check_json_str(data): + return base64.standard_b64decode(data.encode("ascii")) + raise TypeError(f"Invalid type: {data}. Expected a string.") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[BytesLike]: + """ + Check if the provided data is of type BytesLike. + + This method is used to verify if the input data can be considered as a + scalar of bytes-like type, which includes NumPy bytes, strings, bytes, + and integers. + + Parameters + ---------- + data : object + The data to check. + + Returns + ------- + TypeGuard[BytesLike] + True if the data is bytes-like, False otherwise. + """ + return isinstance(data, BytesLike) + + def _cast_scalar_unchecked(self, data: BytesLike) -> bytes: + """ + Cast the provided scalar data to bytes. + + Parameters + ---------- + data : BytesLike + The data to cast. + + Returns + ------- + bytes + The casted data as bytes. + + Notes + ----- + This method does not perform any type checking. + The input data must be bytes-like. + """ + if isinstance(data, str): + return bytes(data, encoding="utf-8") + return bytes(data) + + def cast_scalar(self, data: object) -> bytes: + """ + Attempt to cast a given object to a bytes scalar. + + This method first checks if the provided data is a valid scalar that can be + converted to a bytes scalar. If the check succeeds, the unchecked casting + operation is performed. If the data is not valid, a TypeError is raised. + + Parameters + ---------- + data : object + The data to be cast to a bytes scalar. + + Returns + ------- + bytes + The data cast as a bytes scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a bytes scalar. + """ + + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/common.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/common.py new file mode 100644 index 0000000000000000000000000000000000000000..107b3bd12d1921948b570a68076051c9bd7c6789 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/common.py @@ -0,0 +1,581 @@ +from __future__ import annotations + +import base64 +import struct +import sys +from collections.abc import Sequence +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NewType, + SupportsComplex, + SupportsFloat, + SupportsIndex, + SupportsInt, + TypeGuard, + TypeVar, +) + +import numpy as np + +from zarr.core.dtype.common import ( + ENDIANNESS_STR, + SPECIAL_FLOAT_STRINGS, + EndiannessStr, + JSONFloatV2, + JSONFloatV3, +) + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + +IntLike = SupportsInt | SupportsIndex | bytes | str +FloatLike = SupportsIndex | SupportsFloat | bytes | str +ComplexLike = SupportsFloat | SupportsIndex | SupportsComplex | bytes | str | None +DateTimeUnit = Literal[ + "Y", "M", "W", "D", "h", "m", "s", "ms", "us", "μs", "ns", "ps", "fs", "as", "generic" +] +DATETIME_UNIT: Final = ( + "Y", + "M", + "W", + "D", + "h", + "m", + "s", + "ms", + "us", + "μs", + "ns", + "ps", + "fs", + "as", + "generic", +) + +IntishFloat = NewType("IntishFloat", float) +"""A type for floats that represent integers, like 1.0 (but not 1.1).""" + +IntishStr = NewType("IntishStr", str) +"""A type for strings that represent integers, like "0" or "42".""" + +FloatishStr = NewType("FloatishStr", str) +"""A type for strings that represent floats, like "3.14" or "-2.5".""" + +NumpyEndiannessStr = Literal[">", "<", "="] +NUMPY_ENDIANNESS_STR: Final = ">", "<", "=" + +TFloatDType_co = TypeVar( + "TFloatDType_co", + bound=np.dtypes.Float16DType | np.dtypes.Float32DType | np.dtypes.Float64DType, + covariant=True, +) +TFloatScalar_co = TypeVar( + "TFloatScalar_co", bound=np.float16 | np.float32 | np.float64, covariant=True +) + +TComplexDType_co = TypeVar( + "TComplexDType_co", bound=np.dtypes.Complex64DType | np.dtypes.Complex128DType, covariant=True +) +TComplexScalar_co = TypeVar("TComplexScalar_co", bound=np.complex64 | np.complex128, covariant=True) + + +def endianness_from_numpy_str(endianness: NumpyEndiannessStr) -> EndiannessStr: + """ + Convert a numpy endianness string literal to a human-readable literal value. + + Parameters + ---------- + endianness : Literal[">", "<", "="] + The numpy string representation of the endianness. + + Returns + ------- + Endianness + The human-readable representation of the endianness. + + Raises + ------ + ValueError + If the endianness is invalid. + """ + match endianness: + case "=": + # Use the local system endianness + return sys.byteorder + case "<": + return "little" + case ">": + return "big" + raise ValueError(f"Invalid endianness: {endianness!r}. Expected one of {NUMPY_ENDIANNESS_STR}") + + +def endianness_to_numpy_str(endianness: EndiannessStr) -> NumpyEndiannessStr: + """ + Convert an endianness literal to its numpy string representation. + + Parameters + ---------- + endianness : Endianness + The endianness to convert. + + Returns + ------- + Literal[">", "<"] + The numpy string representation of the endianness. + + Raises + ------ + ValueError + If the endianness is invalid. + """ + match endianness: + case "little": + return "<" + case "big": + return ">" + raise ValueError( + f"Invalid endianness: {endianness!r}. Expected one of {ENDIANNESS_STR} or None" + ) + + +def get_endianness_from_numpy_dtype(dtype: np.dtype[np.generic]) -> EndiannessStr: + """ + Gets the endianness from a numpy dtype that has an endianness. This function will + raise a ValueError if the numpy data type does not have a concrete endianness. + """ + endianness = dtype.byteorder + if dtype.byteorder in NUMPY_ENDIANNESS_STR: + return endianness_from_numpy_str(endianness) # type: ignore [arg-type] + raise ValueError(f"The dtype {dtype} has an unsupported endianness: {endianness}") + + +def float_from_json_v2(data: JSONFloatV2) -> float: + """ + Convert a JSON float to a float (Zarr v2). + + Parameters + ---------- + data : JSONFloat + The JSON float to convert. + + Returns + ------- + float + The float value. + """ + match data: + case "NaN": + return float("nan") + case "Infinity": + return float("inf") + case "-Infinity": + return float("-inf") + case _: + return float(data) + + +def float_from_json_v3(data: JSONFloatV3) -> float: + """ + Convert a JSON float to a float (v3). + + Parameters + ---------- + data : JSONFloat + The JSON float to convert. + + Returns + ------- + float + The float value. + + Notes + ----- + Zarr V3 allows floats to be stored as hex strings. To quote the spec: + "...for float32, "NaN" is equivalent to "0x7fc00000". + This representation is the only way to specify a NaN value other than the specific NaN value + denoted by "NaN"." + """ + + if isinstance(data, str): + if data in SPECIAL_FLOAT_STRINGS: + return float_from_json_v2(data) # type: ignore[arg-type] + if not data.startswith("0x"): + msg = ( + f"Invalid float value: {data!r}. Expected a string starting with the hex prefix" + " '0x', or one of 'NaN', 'Infinity', or '-Infinity'." + ) + raise ValueError(msg) + if len(data[2:]) == 4: + dtype_code = ">e" + elif len(data[2:]) == 8: + dtype_code = ">f" + elif len(data[2:]) == 16: + dtype_code = ">d" + else: + msg = ( + f"Invalid hexadecimal float value: {data!r}. " + "Expected the '0x' prefix to be followed by 4, 8, or 16 numeral characters" + ) + raise ValueError(msg) + return float(struct.unpack(dtype_code, bytes.fromhex(data[2:]))[0]) + return float_from_json_v2(data) + + +def bytes_from_json(data: str, *, zarr_format: ZarrFormat) -> bytes: + """ + Convert a JSON string to bytes + + Parameters + ---------- + data : str + The JSON string to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + bytes + The bytes. + """ + if zarr_format == 2: + return base64.b64decode(data.encode("ascii")) + # TODO: differentiate these as needed. This is a spec question. + if zarr_format == 3: + return base64.b64decode(data.encode("ascii")) + raise ValueError(f"Invalid zarr format: {zarr_format}. Expected 2 or 3.") # pragma: no cover + + +def bytes_to_json(data: bytes, zarr_format: ZarrFormat) -> str: + """ + Convert bytes to JSON. + + Parameters + ---------- + data : bytes + The bytes to store. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + str + The bytes encoded as ascii using the base64 alphabet. + """ + # TODO: decide if we are going to make this implementation zarr format-specific + return base64.b64encode(data).decode("ascii") + + +def float_to_json_v2(data: float | np.floating[Any]) -> JSONFloatV2: + """ + Convert a float to JSON (v2). + + Parameters + ---------- + data : float or np.floating + The float value to convert. + + Returns + ------- + JSONFloat + The JSON representation of the float. + """ + if np.isnan(data): + return "NaN" + elif np.isinf(data): + return "Infinity" if data > 0 else "-Infinity" + return float(data) + + +def float_to_json_v3(data: float | np.floating[Any]) -> JSONFloatV3: + """ + Convert a float to JSON (v3). + + Parameters + ---------- + data : float or np.floating + The float value to convert. + + Returns + ------- + JSONFloat + The JSON representation of the float. + """ + # v3 can in principle handle distinct NaN values, but numpy does not represent these explicitly + # so we just reuse the v2 routine here + return float_to_json_v2(data) + + +def complex_float_to_json_v3( + data: complex | np.complexfloating[Any, Any], +) -> tuple[JSONFloatV3, JSONFloatV3]: + """ + Convert a complex number to JSON as defined by the Zarr V3 spec. + + Parameters + ---------- + data : complex or np.complexfloating + The complex value to convert. + + Returns + ------- + tuple[JSONFloat, JSONFloat] + The JSON representation of the complex number. + """ + return float_to_json_v3(data.real), float_to_json_v3(data.imag) + + +def complex_float_to_json_v2( + data: complex | np.complexfloating[Any, Any], +) -> tuple[JSONFloatV2, JSONFloatV2]: + """ + Convert a complex number to JSON as defined by the Zarr V2 spec. + + Parameters + ---------- + data : complex | np.complexfloating + The complex value to convert. + + Returns + ------- + tuple[JSONFloat, JSONFloat] + The JSON representation of the complex number. + """ + return float_to_json_v2(data.real), float_to_json_v2(data.imag) + + +def complex_float_from_json_v2(data: tuple[JSONFloatV2, JSONFloatV2]) -> complex: + """ + Convert a JSON complex float to a complex number (v2). + + Parameters + ---------- + data : tuple[JSONFloat, JSONFloat] + The JSON complex float to convert. + + Returns + ------- + np.complexfloating + The complex number. + """ + return complex(float_from_json_v2(data[0]), float_from_json_v2(data[1])) + + +def complex_float_from_json_v3(data: tuple[JSONFloatV3, JSONFloatV3]) -> complex: + """ + Convert a JSON complex float to a complex number (v3). + + Parameters + ---------- + data : tuple[JSONFloat, JSONFloat] + The JSON complex float to convert. + + Returns + ------- + np.complexfloating + The complex number. + """ + return complex(float_from_json_v3(data[0]), float_from_json_v3(data[1])) + + +def check_json_float_v2(data: JSON) -> TypeGuard[JSONFloatV2]: + """ + Check if a JSON value represents a float (v2). + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a float, False otherwise. + """ + return data in ("NaN", "Infinity", "-Infinity") or isinstance(data, float | int) + + +def check_json_float_v3(data: JSON) -> TypeGuard[JSONFloatV3]: + """ + Check if a JSON value represents a float (v3). + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a float, False otherwise. + """ + return check_json_float_v2(data) or (isinstance(data, str) and data.startswith("0x")) + + +def check_json_complex_float_v2(data: JSON) -> TypeGuard[tuple[JSONFloatV2, JSONFloatV2]]: + """ + Check if a JSON value represents a complex float, as per the behavior of zarr-python 2.x + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a complex float, False otherwise. + """ + return ( + not isinstance(data, str) + and isinstance(data, Sequence) + and len(data) == 2 + and check_json_float_v2(data[0]) + and check_json_float_v2(data[1]) + ) + + +def check_json_complex_float_v3(data: JSON) -> TypeGuard[tuple[JSONFloatV3, JSONFloatV3]]: + """ + Check if a JSON value represents a complex float, as per the zarr v3 spec + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a complex float, False otherwise. + """ + return ( + not isinstance(data, str) + and isinstance(data, Sequence) + and len(data) == 2 + and check_json_float_v3(data[0]) + and check_json_float_v3(data[1]) + ) + + +def check_json_int(data: JSON) -> TypeGuard[int]: + """ + Check if a JSON value is an integer. + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is an integer, False otherwise. + """ + return bool(isinstance(data, int)) + + +def check_json_intish_float(data: JSON) -> TypeGuard[IntishFloat]: + """ + Check if a JSON value is an "intish float", i.e. a float that represents an integer, like 0.0. + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is an intish float, False otherwise. + """ + return isinstance(data, float) and data.is_integer() + + +def check_json_intish_str(data: JSON) -> TypeGuard[IntishStr]: + """ + Check if a JSON value is a string that represents an integer, like "0", "42", or "-5". + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + bool + True if the data is a string representing an integer, False otherwise. + """ + if not isinstance(data, str): + return False + + try: + int(data) + except ValueError: + return False + else: + return True + + +def check_json_floatish_str(data: JSON) -> TypeGuard[FloatishStr]: + """ + Check if a JSON value is a string that represents a float, like "3.14", "-2.5", or "0.0". + + Note: This function is intended to be used AFTER check_json_float_v2/v3, so it only + handles regular string representations that those functions don't cover. + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + bool + True if the data is a string representing a regular float, False otherwise. + """ + if not isinstance(data, str): + return False + + try: + float(data) + except ValueError: + return False + else: + return True + + +def check_json_str(data: JSON) -> TypeGuard[str]: + """ + Check if a JSON value is a string. + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a string, False otherwise. + """ + return bool(isinstance(data, str)) + + +def check_json_bool(data: JSON) -> TypeGuard[bool]: + """ + Check if a JSON value is a boolean. + + Parameters + ---------- + data : JSON + The JSON value to check. + + Returns + ------- + Bool + True if the data is a boolean, False otherwise. + """ + return isinstance(data, bool) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/complex.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/complex.py new file mode 100644 index 0000000000000000000000000000000000000000..99abee5e24ed060a0642cff64e5f6102ad9e4ffb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/complex.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + ClassVar, + Literal, + Self, + TypeGuard, + overload, +) + +import numpy as np + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasEndianness, + HasItemSize, + check_dtype_spec_v2, +) +from zarr.core.dtype.npy.common import ( + ComplexLike, + TComplexDType_co, + TComplexScalar_co, + check_json_complex_float_v2, + check_json_complex_float_v3, + complex_float_from_json_v2, + complex_float_from_json_v3, + complex_float_to_json_v2, + complex_float_to_json_v3, + endianness_to_numpy_str, + get_endianness_from_numpy_dtype, +) +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + + +@dataclass(frozen=True) +class BaseComplex(ZDType[TComplexDType_co, TComplexScalar_co], HasEndianness, HasItemSize): + """ + A base class for Zarr data types that wrap NumPy complex float data types. + """ + + # This attribute holds the possible zarr v2 JSON names for the data type + _zarr_v2_names: ClassVar[tuple[str, ...]] + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of this data type from a NumPy complex dtype. + + Parameters + ---------- + dtype : TBaseDType + The native dtype to convert. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the dtype is not compatible with this data type. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> TComplexDType_co: + """ + Convert this class to a NumPy complex dtype with the appropriate byte order. + + Returns + ------- + TComplexDType_co + A NumPy data type object representing the complex data type with the specified byte order. + """ + + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) # type: ignore[return-value] + + @classmethod + def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[DTypeConfig_V2[str, None]]: + """ + Check that the input is a valid JSON representation of this data type. + + The input data must be a mapping that contains a "name" key that is one of + the strings from cls._zarr_v2_names and an "object_codec_id" key that is None. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + bool + True if the input is a valid JSON representation, False otherwise. + """ + return ( + check_dtype_spec_v2(data) + and data["name"] in cls._zarr_v2_names + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[str]: + """ + Check that the input is a valid JSON representation of this data type in Zarr V3. + + This method verifies that the provided data matches the expected Zarr V3 + representation, which is the string specified by the class-level attribute _zarr_v3_name. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[str] + True if the input is a valid representation of this class in Zarr V3, False otherwise. + """ + + return data == cls._zarr_v3_name + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this class. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v2(data): + # Going via numpy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected {cls._zarr_v3_name}." + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[str, None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> str: ... + + def to_json(self, zarr_format: ZarrFormat) -> DTypeConfig_V2[str, None] | str: + """ + Serialize this object to a JSON-serializable representation. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version. Supported values are 2 and 3. + + Returns + ------- + DTypeConfig_V2[str, None] | str + If ``zarr_format`` is 2, a dictionary with ``"name"`` and ``"object_codec_id"`` keys is + returned. + If ``zarr_format`` is 3, a string representation of the complex data type is returned. + + Raises + ------ + ValueError + If `zarr_format` is not 2 or 3. + """ + + if zarr_format == 2: + return {"name": self.to_native_dtype().str, "object_codec_id": None} + elif zarr_format == 3: + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[ComplexLike]: + """ + Check that the input is a scalar complex value. + + Parameters + ---------- + data : object + The value to check. + + Returns + ------- + TypeGuard[ComplexLike] + True if the input is a scalar complex value, False otherwise. + """ + return isinstance(data, ComplexLike) + + def _cast_scalar_unchecked(self, data: ComplexLike) -> TComplexScalar_co: + """ + Cast the provided scalar data to the native scalar type of this class. + + Parameters + ---------- + data : ComplexLike + The data to cast. + + Returns + ------- + TComplexScalar_co + The casted data as a numpy complex scalar. + + Notes + ----- + This method does not perform any type checking. + The input data must be a scalar complex value. + """ + return self.to_native_dtype().type(data) # type: ignore[return-value] + + def cast_scalar(self, data: object) -> TComplexScalar_co: + """ + Attempt to cast a given object to a numpy complex scalar. + + Parameters + ---------- + data : object + The data to be cast to a numpy complex scalar. + + Returns + ------- + TComplexScalar_co + The data cast as a numpy complex scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a numpy complex scalar. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> TComplexScalar_co: + """ + Get the default value, which is 0 cast to this dtype + + Returns + ------- + Int scalar + The default value. + """ + return self._cast_scalar_unchecked(0) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TComplexScalar_co: + """ + Read a JSON-serializable value as a numpy float. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + TScalar_co + The numpy float. + """ + if zarr_format == 2: + if check_json_complex_float_v2(data): + return self._cast_scalar_unchecked(complex_float_from_json_v2(data)) + raise TypeError( + f"Invalid type: {data}. Expected a float or a special string encoding of a float." + ) + elif zarr_format == 3: + if check_json_complex_float_v3(data): + return self._cast_scalar_unchecked(complex_float_from_json_v3(data)) + raise TypeError( + f"Invalid type: {data}. Expected a float or a special string encoding of a float." + ) + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> JSON: + """ + Convert an object to a JSON-serializable float. + + Parameters + ---------- + data : _BaseScalar + The value to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + JSON + The JSON-serializable form of the complex number, which is a list of two floats, + each of which is encoding according to a zarr-format-specific encoding. + """ + if zarr_format == 2: + return complex_float_to_json_v2(self.cast_scalar(data)) + elif zarr_format == 3: + return complex_float_to_json_v3(self.cast_scalar(data)) + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + +@dataclass(frozen=True, kw_only=True) +class Complex64(BaseComplex[np.dtypes.Complex64DType, np.complex64]): + """ + A Zarr data type for arrays containing 64 bit complex floats. + + Wraps the [`np.dtypes.Complex64DType`][numpy.dtypes.Complex64DType] data type. Scalars for this data type + are instances of [`np.complex64`][numpy.complex64]. + + Attributes + ---------- + dtype_cls : Type[np.dtypes.Complex64DType] + The numpy dtype class for this data type. + _zarr_v3_name : ClassVar[Literal["complex64"]] + The name of this data type in Zarr V3. + _zarr_v2_names : ClassVar[tuple[Literal[">c8"], Literal["c8"], Literal["c8", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 8 + + +@dataclass(frozen=True, kw_only=True) +class Complex128(BaseComplex[np.dtypes.Complex128DType, np.complex128], HasEndianness): + """ + A Zarr data type for arrays containing 64 bit complex floats. + + Wraps the [`np.dtypes.Complex128DType`][numpy.dtypes.Complex128DType] data type. Scalars for this data type + are instances of [`np.complex128`][numpy.complex128]. + + Attributes + ---------- + dtype_cls : Type[np.dtypes.Complex128DType] + The numpy dtype class for this data type. + _zarr_v3_name : ClassVar[Literal["complex128"]] + The name of this data type in Zarr V3. + _zarr_v2_names : ClassVar[tuple[Literal[">c16"], Literal["c16"], Literal["c16", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 16 diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/float.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/float.py new file mode 100644 index 0000000000000000000000000000000000000000..2a23cb429d2b90402b505e2eab5073227ffab1bb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/float.py @@ -0,0 +1,435 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar, Literal, Self, TypeGuard, overload + +import numpy as np + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasEndianness, + HasItemSize, + check_dtype_spec_v2, +) +from zarr.core.dtype.npy.common import ( + FloatLike, + TFloatDType_co, + TFloatScalar_co, + check_json_float_v2, + check_json_float_v3, + check_json_floatish_str, + endianness_to_numpy_str, + float_from_json_v2, + float_from_json_v3, + float_to_json_v2, + float_to_json_v3, + get_endianness_from_numpy_dtype, +) +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + + +@dataclass(frozen=True) +class BaseFloat(ZDType[TFloatDType_co, TFloatScalar_co], HasEndianness, HasItemSize): + """ + A base class for Zarr data types that wrap NumPy float data types. + """ + + # This attribute holds the possible zarr v2 JSON names for the data type + _zarr_v2_names: ClassVar[tuple[str, ...]] + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of this ZDType from a NumPy data type. + + Parameters + ---------- + dtype : TBaseDType + The NumPy data type. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> TFloatDType_co: + """ + Convert the wrapped data type to a NumPy data type. + + Returns + ------- + TFloatDType_co + The NumPy data type. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) # type: ignore[return-value] + + @classmethod + def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[DTypeConfig_V2[str, None]]: + """ + Check that the input is a valid JSON representation of this data type. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + TypeGuard[DTypeConfig_V2[str, None]] + True if the input is a valid JSON representation of this data type, False otherwise. + """ + return ( + check_dtype_spec_v2(data) + and data["name"] in cls._zarr_v2_names + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[str]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + TypeGuard[str] + True if the input is a valid JSON representation of this class, False otherwise. + """ + return data == cls._zarr_v3_name + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this ZDType from Zarr v2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this ZDType from Zarr v3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected {cls._zarr_v3_name}." + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[str, None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> str: ... + + def to_json(self, zarr_format: ZarrFormat) -> DTypeConfig_V2[str, None] | str: + """ + Convert the wrapped data type to a JSON-serializable form. + + Parameters + ---------- + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + DTypeConfig_V2[str, None] or str + The JSON-serializable representation of the wrapped data type. + + Raises + ------ + ValueError + If zarr_format is not 2 or 3. + """ + if zarr_format == 2: + return {"name": self.to_native_dtype().str, "object_codec_id": None} + elif zarr_format == 3: + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[FloatLike]: + """ + Check that the input is a valid scalar value. + + Parameters + ---------- + data : object + The input to check. + + Returns + ------- + TypeGuard[FloatLike] + True if the input is a valid scalar value, False otherwise. + """ + if isinstance(data, str): + # Only accept strings that are valid float representations (e.g. "NaN", "inf"). + # Plain strings that cannot be converted should return False so that cast_scalar + # raises TypeError rather than a confusing ValueError. + try: + self.to_native_dtype().type(data) + except (ValueError, OverflowError): + return False + else: + return True + return isinstance(data, FloatLike) + + def _cast_scalar_unchecked(self, data: FloatLike) -> TFloatScalar_co: + """ + Cast a scalar value to a NumPy float scalar. + + Parameters + ---------- + data : FloatLike + The scalar value to cast. + + Returns + ------- + TFloatScalar_co + The NumPy float scalar. + """ + return self.to_native_dtype().type(data) # type: ignore[return-value] + + def cast_scalar(self, data: object) -> TFloatScalar_co: + """ + Cast a scalar value to a NumPy float scalar. + + Parameters + ---------- + data : object + The scalar value to cast. + + Returns + ------- + TFloatScalar_co + The NumPy float scalar. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> TFloatScalar_co: + """ + Get the default value, which is 0 cast to this zdtype. + + Returns + ------- + TFloatScalar_co + The default value. + """ + return self._cast_scalar_unchecked(0) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TFloatScalar_co: + """ + Read a JSON-serializable value as a NumPy float scalar. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + TFloatScalar_co + The NumPy float scalar. + """ + if zarr_format == 2: + if check_json_float_v2(data): + return self._cast_scalar_unchecked(float_from_json_v2(data)) + elif check_json_floatish_str(data): + return self._cast_scalar_unchecked(float(data)) + else: + raise TypeError( + f"Invalid type: {data}. Expected a float or a special string encoding of a float." + ) + elif zarr_format == 3: + if check_json_float_v3(data): + return self._cast_scalar_unchecked(float_from_json_v3(data)) + elif check_json_floatish_str(data): + return self._cast_scalar_unchecked(float(data)) + else: + raise TypeError( + f"Invalid type: {data}. Expected a float or a special string encoding of a float." + ) + else: + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> float | str: + """ + Convert an object to a JSON-serializable float. + + Parameters + ---------- + data : _BaseScalar + The value to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + JSON + The JSON-serializable form of the float, which is potentially a number or a string. + See the zarr specifications for details on the JSON encoding for floats. + """ + if zarr_format == 2: + return float_to_json_v2(self.cast_scalar(data)) + elif zarr_format == 3: + return float_to_json_v3(self.cast_scalar(data)) + else: + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + +@dataclass(frozen=True, kw_only=True) +class Float16(BaseFloat[np.dtypes.Float16DType, np.float16]): + """ + A Zarr data type for arrays containing 16-bit floating point numbers. + + Wraps the [`np.dtypes.Float16DType`][numpy.dtypes.Float16DType] data type. Scalars for this data type are instances + of [`np.float16`][numpy.float16]. + + Attributes + ---------- + dtype_cls : Type[np.dtypes.Float16DType] + The NumPy dtype class for this data type. + + References + ---------- + This class implements the float16 data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Float16DType + _zarr_v3_name = "float16" + _zarr_v2_names: ClassVar[tuple[Literal[">f2"], Literal["f2", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 2 + + +@dataclass(frozen=True, kw_only=True) +class Float32(BaseFloat[np.dtypes.Float32DType, np.float32]): + """ + A Zarr data type for arrays containing 32-bit floating point numbers. + + Wraps the [`np.dtypes.Float32DType`][numpy.dtypes.Float32DType] data type. Scalars for this data type are instances + of [`np.float32`][numpy.float32]. + + Attributes + ---------- + dtype_cls : Type[np.dtypes.Float32DType] + The NumPy dtype class for this data type. + + References + ---------- + This class implements the float32 data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Float32DType + _zarr_v3_name = "float32" + _zarr_v2_names: ClassVar[tuple[Literal[">f4"], Literal["f4", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 4 + + +@dataclass(frozen=True, kw_only=True) +class Float64(BaseFloat[np.dtypes.Float64DType, np.float64]): + """ + A Zarr data type for arrays containing 64-bit floating point numbers. + + Wraps the [`np.dtypes.Float64DType`][numpy.dtypes.Float64DType] data type. Scalars for this data type are instances + of [`np.float64`][numpy.float64]. + + Attributes + ---------- + dtype_cls : Type[np.dtypes.Float64DType] + The NumPy dtype class for this data type. + + References + ---------- + This class implements the float64 data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Float64DType + _zarr_v3_name = "float64" + _zarr_v2_names: ClassVar[tuple[Literal[">f8"], Literal["f8", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 8 diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/int.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/int.py new file mode 100644 index 0000000000000000000000000000000000000000..f71f535abbfdc18ab193a7e7bb3c2e7317154508 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/int.py @@ -0,0 +1,1521 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + ClassVar, + Literal, + Self, + SupportsIndex, + SupportsInt, + TypeGuard, + TypeVar, + overload, +) + +import numpy as np + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasEndianness, + HasItemSize, + check_dtype_spec_v2, +) +from zarr.core.dtype.npy.common import ( + check_json_int, + check_json_intish_float, + check_json_intish_str, + endianness_to_numpy_str, + get_endianness_from_numpy_dtype, +) +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + +_NumpyIntDType = ( + np.dtypes.Int8DType + | np.dtypes.Int16DType + | np.dtypes.Int32DType + | np.dtypes.Int64DType + | np.dtypes.UInt8DType + | np.dtypes.UInt16DType + | np.dtypes.UInt32DType + | np.dtypes.UInt64DType +) +_NumpyIntScalar = ( + np.int8 | np.int16 | np.int32 | np.int64 | np.uint8 | np.uint16 | np.uint32 | np.uint64 +) +TIntDType_co = TypeVar("TIntDType_co", bound=_NumpyIntDType, covariant=True) +TIntScalar_co = TypeVar("TIntScalar_co", bound=_NumpyIntScalar, covariant=True) +IntLike = SupportsInt | SupportsIndex | bytes | str + + +@dataclass(frozen=True) +class BaseInt(ZDType[TIntDType_co, TIntScalar_co], HasItemSize): + """ + A base class for integer data types in Zarr. + + This class provides methods for serialization and deserialization of integer types + in both Zarr v2 and v3 formats, as well as methods for checking and casting scalars. + """ + + _zarr_v2_names: ClassVar[tuple[str, ...]] + + @classmethod + def _check_json_v2(cls, data: object) -> TypeGuard[DTypeConfig_V2[str, None]]: + """ + Check that the input is a valid JSON representation of this integer data type in Zarr V2. + + This method verifies that the provided data matches the expected Zarr V2 representation + for this data type. The input data must be a mapping that contains a "name" key that is + one of the strings from cls._zarr_v2_names and an "object_codec_id" key that is None. + + Parameters + ---------- + data : object + The JSON data to check. + + Returns + ------- + TypeGuard[DTypeConfig_V2[str, None]] + True if the input is a valid representation of this class in Zarr V2, + False otherwise. + """ + + return ( + check_dtype_spec_v2(data) + and data["name"] in cls._zarr_v2_names + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: object) -> TypeGuard[str]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : object + The JSON data to check. + + Returns + ------- + TypeGuard[str] + True if the input is a valid representation of this class in Zarr v3, + False otherwise. + """ + return data == cls._zarr_v3_name + + def _check_scalar(self, data: object) -> TypeGuard[IntLike]: + """ + Check if the input object is of an IntLike type. + + This method verifies whether the provided data can be considered as an integer-like + value, which includes objects supporting integer conversion. + + Parameters + ---------- + data : object + The data to check. + + Returns + ------- + TypeGuard[IntLike] + True if the data is IntLike, False otherwise. + """ + + return isinstance(data, IntLike) + + def _cast_scalar_unchecked(self, data: IntLike) -> TIntScalar_co: + """ + Casts a given scalar value to the native integer scalar type without type checking. + + Parameters + ---------- + data : IntLike + The scalar value to cast. + + Returns + ------- + TIntScalar_co + The casted integer scalar of the native dtype. + """ + + return self.to_native_dtype().type(data) # type: ignore[return-value] + + def cast_scalar(self, data: object) -> TIntScalar_co: + """ + Attempt to cast a given object to a NumPy integer scalar. + + Parameters + ---------- + data : object + The data to be cast to a NumPy integer scalar. + + Returns + ------- + TIntScalar_co + The data cast as a NumPy integer scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a NumPy integer scalar. + """ + + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> TIntScalar_co: + """ + Get the default value, which is 0 cast to this dtype. + + Returns + ------- + TIntScalar_co + The default value. + """ + return self._cast_scalar_unchecked(0) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> TIntScalar_co: + """ + Read a JSON-serializable value as a NumPy int scalar. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The Zarr format version. + + Returns + ------- + TIntScalar_co + The NumPy int scalar. + + Raises + ------ + TypeError + If the input is not a valid integer type. + """ + if check_json_int(data): + return self._cast_scalar_unchecked(data) + if check_json_intish_float(data): + return self._cast_scalar_unchecked(int(data)) + + if check_json_intish_str(data): + return self._cast_scalar_unchecked(int(data)) + + raise TypeError(f"Invalid type: {data}. Expected an integer.") + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> int: + """ + Convert an object to a JSON serializable scalar. For the integer data types, + the JSON form is a plain integer. + + Parameters + ---------- + data : object + The value to convert. + zarr_format : ZarrFormat + The Zarr format version. + + Returns + ------- + int + The JSON-serializable form of the scalar. + """ + return int(self.cast_scalar(data)) + + +@dataclass(frozen=True, kw_only=True) +class Int8(BaseInt[np.dtypes.Int8DType, np.int8]): + """ + A Zarr data type for arrays containing 8-bit signed integers. + + Wraps the [`np.dtypes.Int8DType`][numpy.dtypes.Int8DType] data type. Scalars for this data type are + instances of [`np.int8`][numpy.int8]. + + Attributes + ---------- + dtype_cls : np.dtypes.Int8DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 8-bit signed integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Int8DType + _zarr_v3_name: ClassVar[Literal["int8"]] = "int8" + _zarr_v2_names: ClassVar[tuple[Literal["|i1"]]] = ("|i1",) + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an Int8 from an np.dtype('int8') instance. + + Parameters + ---------- + dtype : TBaseDType + The np.dtype('int8') instance. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not a valid representation of this class Int8. + """ + if cls._check_native_dtype(dtype): + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self: Self) -> np.dtypes.Int8DType: + """ + Convert the Int8 instance to an np.dtype('int8') instance. + + Returns + ------- + np.dtypes.Int8DType + The np.dtype('int8') instance. + """ + return self.dtype_cls() + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an Int8 from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class Int8. + """ + if cls._check_json_v2(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v2_names[0]!r}" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an Int8 from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class Int8. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal["|i1"], None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> Literal["int8"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal["|i1"], None] | Literal["int8"]: + """ + Convert the data type to a JSON-serializable form. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version. + + Returns + ------- + ``DTypeConfig_V2[Literal["|i1"], None] | Literal["int8"]`` + The JSON-serializable representation of the data type. + + Raises + ------ + ValueError + If the zarr_format is not 2 or 3. + """ + if zarr_format == 2: + return {"name": self._zarr_v2_names[0], "object_codec_id": None} + elif zarr_format == 3: + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 1 + + +@dataclass(frozen=True, kw_only=True) +class UInt8(BaseInt[np.dtypes.UInt8DType, np.uint8]): + """ + A Zarr data type for arrays containing 8-bit unsigned integers. + + Wraps the [`np.dtypes.UInt8DType`][numpy.dtypes.UInt8DType] data type. Scalars for this data type are instances of [`np.uint8`][numpy.uint8]. + + Attributes + ---------- + dtype_cls : np.dtypes.UInt8DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 8-bit unsigned integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.UInt8DType + _zarr_v3_name: ClassVar[Literal["uint8"]] = "uint8" + _zarr_v2_names: ClassVar[tuple[Literal["|u1"]]] = ("|u1",) + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create a UInt8 from an np.dtype('uint8') instance. + """ + if cls._check_native_dtype(dtype): + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self: Self) -> np.dtypes.UInt8DType: + """ + Create a NumPy unsigned 8-bit integer dtype instance from this UInt8 ZDType. + + Returns + ------- + np.dtypes.UInt8DType + The NumPy unsigned 8-bit integer dtype. + """ + + return self.dtype_cls() + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + + if cls._check_json_v2(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v2_names[0]!r}" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal["|u1"], None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> Literal["uint8"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal["|u1"], None] | Literal["uint8"]: + """ + Convert the data type to a JSON-serializable form. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version. Supported values are 2 and 3. + + Returns + ------- + ``DTypeConfig_V2[Literal["|u1"], None] | Literal["uint8"]`` + The JSON-serializable representation of the data type. + + Raises + ------ + ValueError + If `zarr_format` is not 2 or 3. + """ + if zarr_format == 2: + # For Zarr format version 2, return a dictionary with the name and object codec ID. + return {"name": self._zarr_v2_names[0], "object_codec_id": None} + elif zarr_format == 3: + # For Zarr format version 3, return the v3 name as a string. + return self._zarr_v3_name + # Raise an error if the zarr_format is neither 2 nor 3. + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 1 + + +@dataclass(frozen=True, kw_only=True) +class Int16(BaseInt[np.dtypes.Int16DType, np.int16], HasEndianness): + """ + A Zarr data type for arrays containing 16-bit signed integers. + + Wraps the [`np.dtypes.Int16DType`][numpy.dtypes.Int16DType] data type. Scalars for this data type are instances of + [`np.int16`][numpy.int16]. + + Attributes + ---------- + dtype_cls : np.dtypes.Int16DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 16-bit signed integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Int16DType + _zarr_v3_name: ClassVar[Literal["int16"]] = "int16" + _zarr_v2_names: ClassVar[tuple[Literal[">i2"], Literal["i2", " Self: + """ + Create an instance of this data type from an np.dtype('int16') instance. + + Parameters + ---------- + dtype : np.dtype + The instance of np.dtype('int16') to create from. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not an instance of np.dtype('int16'). + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.Int16DType: + """ + Convert the data type to an np.dtype('int16') instance. + + Returns + ------- + np.dtype + The np.dtype('int16') instance. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names!r}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">i2", " Literal["int16"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">i2", "i2", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 2 + + +@dataclass(frozen=True, kw_only=True) +class UInt16(BaseInt[np.dtypes.UInt16DType, np.uint16], HasEndianness): + """ + A Zarr data type for arrays containing 16-bit unsigned integers. + + Wraps the [`np.dtypes.UInt16DType`][numpy.dtypes.UInt16DType] data type. Scalars for this data type are instances of + [`np.uint16`][numpy.uint16]. + + Attributes + ---------- + dtype_cls : np.dtypes.UInt16DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the unsigned 16-bit unsigned integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.UInt16DType + _zarr_v3_name: ClassVar[Literal["uint16"]] = "uint16" + _zarr_v2_names: ClassVar[tuple[Literal[">u2"], Literal["u2", " Self: + """ + Create an instance of this data type from an np.dtype('uint16') instance. + + Parameters + ---------- + dtype : np.dtype + The NumPy data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not an instance of np.dtype('uint16'). + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.UInt16DType: + """ + Convert the data type to an np.dtype('uint16') instance. + + Returns + ------- + np.dtype + The np.dtype('uint16') instance. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of UInt16. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of UInt16. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">u2", " Literal["uint16"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">u2", "u2", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 2 + + +@dataclass(frozen=True, kw_only=True) +class Int32(BaseInt[np.dtypes.Int32DType, np.int32], HasEndianness): + """ + A Zarr data type for arrays containing 32-bit signed integers. + + Wraps the [`np.dtypes.Int32DType`][numpy.dtypes.Int32DType] data type. Scalars for this data type are instances of + [`np.int32`][numpy.int32]. + + Attributes + ---------- + dtype_cls : np.dtypes.Int32DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 32-bit signed integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Int32DType + _zarr_v3_name: ClassVar[Literal["int32"]] = "int32" + _zarr_v2_names: ClassVar[tuple[Literal[">i4"], Literal["i4", " TypeGuard[np.dtypes.Int32DType]: + """ + A type guard that checks if the input is assignable to the type of ``cls.dtype_class`` + + This method is overridden for this particular data type because of a Windows-specific issue + where np.dtype('i') creates an instance of ``np.dtypes.IntDType``, rather than an + instance of ``np.dtypes.Int32DType``, even though both represent 32-bit signed integers. + + Parameters + ---------- + dtype : TDType + The dtype to check. + + Returns + ------- + Bool + True if the dtype matches, False otherwise. + """ + return super()._check_native_dtype(dtype) or dtype == np.dtypes.Int32DType() + + @classmethod + def from_native_dtype(cls: type[Self], dtype: TBaseDType) -> Self: + """ + Create an Int32 from an np.dtype('int32') instance. + + Parameters + ---------- + dtype : TBaseDType + The np.dtype('int32') instance. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class Int32. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self: Self) -> np.dtypes.Int32DType: + """ + Convert the Int32 instance to an np.dtype('int32') instance. + + Returns + ------- + np.dtypes.Int32DType + The np.dtype('int32') instance. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an Int32 from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class Int32. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names!r}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an Int32 from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class Int32. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">i4", " Literal["int32"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">i4", "i4", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 4 + + +@dataclass(frozen=True, kw_only=True) +class UInt32(BaseInt[np.dtypes.UInt32DType, np.uint32], HasEndianness): + """ + A Zarr data type for arrays containing 32-bit unsigned integers. + + Wraps the [`np.dtypes.UInt32DType`][numpy.dtypes.UInt32DType] data type. Scalars for this data type are instances of + [`np.uint32`][numpy.uint32]. + + Attributes + ---------- + dtype_cls : np.dtypes.UInt32DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 32-bit unsigned integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.UInt32DType + _zarr_v3_name: ClassVar[Literal["uint32"]] = "uint32" + _zarr_v2_names: ClassVar[tuple[Literal[">u4"], Literal["u4", " Self: + """ + Create a UInt32 from an np.dtype('uint32') instance. + + Parameters + ---------- + dtype : TBaseDType + The NumPy data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not a valid representation of this class 32-bit unsigned + integer. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.UInt32DType: + """ + Create a NumPy unsigned 32-bit integer dtype instance from this UInt32 ZDType. + + Returns + ------- + np.dtypes.UInt32DType + The NumPy unsigned 32-bit integer dtype. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class 32-bit unsigned + integer. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class 32-bit unsigned + integer. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">u4", " Literal["uint32"]: ... + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">u4", "u4", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 4 + + +@dataclass(frozen=True, kw_only=True) +class Int64(BaseInt[np.dtypes.Int64DType, np.int64], HasEndianness): + """ + A Zarr data type for arrays containing 64-bit signed integers. + + Wraps the [`np.dtypes.Int64DType`][numpy.dtypes.Int64DType] data type. Scalars for this data type are instances of + [`np.int64`][numpy.int64]. + + Attributes + ---------- + dtype_cls : np.dtypes.Int64DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the 64-bit signed integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.Int64DType + _zarr_v3_name: ClassVar[Literal["int64"]] = "int64" + _zarr_v2_names: ClassVar[tuple[Literal[">i8"], Literal["i8", " Self: + """ + Create an Int64 from an np.dtype('int64') instance. + + Parameters + ---------- + dtype : TBaseDType + The NumPy data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not a valid representation of this class 64-bit signed + integer. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.Int64DType: + """ + Create a NumPy signed 64-bit integer dtype instance from this Int64 ZDType. + + Returns + ------- + np.dtypes.Int64DType + The NumPy signed 64-bit integer dtype. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class 64-bit signed + integer. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class 64-bit signed + integer. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">i8", " Literal["int64"]: ... + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">i8", "i8", " int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 8 + + +@dataclass(frozen=True, kw_only=True) +class UInt64(BaseInt[np.dtypes.UInt64DType, np.uint64], HasEndianness): + """ + A Zarr data type for arrays containing 64-bit unsigned integers. + + Wraps the [`np.dtypes.UInt64DType`][numpy.dtypes.UInt64DType] data type. Scalars for this data type + are instances of [`np.uint64`][numpy.uint64]. + + Attributes + ---------- + dtype_cls: np.dtypes.UInt64DType + The class of the underlying NumPy dtype. + + References + ---------- + This class implements the unsigned 64-bit integer data type defined in Zarr V2 and V3. + + See the [Zarr V2](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding) and [Zarr V3](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/data-types/index.rst) specification documents for details. + """ + + dtype_cls = np.dtypes.UInt64DType + _zarr_v3_name: ClassVar[Literal["uint64"]] = "uint64" + _zarr_v2_names: ClassVar[tuple[Literal[">u8"], Literal["u8", " np.dtypes.UInt64DType: + """ + Convert the data type to a native NumPy dtype. + + Returns + ------- + np.dtypes.UInt64DType + The native NumPy dtype.eeeeeeeeeeeeeeeee + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls().newbyteorder(byte_order) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class unsigned 64-bit + integer. + """ + if cls._check_json_v2(data): + # Going via NumPy ensures that we get the endianness correct without + # annoying string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected one of the strings {cls._zarr_v2_names}." + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from Zarr V3-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class unsigned 64-bit + integer. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal[">u8", " Literal["uint64"]: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[Literal[">u8", "u8", " Self: + """ + Create an instance of this data type from a native NumPy dtype. + + Parameters + ---------- + dtype : TBaseDType + The native NumPy dtype. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input dtype is not a valid representation of this class unsigned 64-bit + integer. + """ + if cls._check_native_dtype(dtype): + return cls(endianness=get_endianness_from_numpy_dtype(dtype)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 8 diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/string.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/string.py new file mode 100644 index 0000000000000000000000000000000000000000..904280a33027d2b73d458165a4b3ca7f746ffe13 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/string.py @@ -0,0 +1,825 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + ClassVar, + Literal, + Protocol, + Self, + TypedDict, + TypeGuard, + overload, + runtime_checkable, +) + +import numpy as np + +from zarr.core.common import NamedConfig +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasEndianness, + HasItemSize, + HasLength, + HasObjectCodec, + check_dtype_spec_v2, + v3_unstable_dtype_warning, +) +from zarr.core.dtype.npy.common import ( + check_json_str, + endianness_to_numpy_str, + get_endianness_from_numpy_dtype, +) +from zarr.core.dtype.wrapper import TDType_co, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + from zarr.core.dtype.wrapper import TBaseDType + +_NUMPY_SUPPORTS_VLEN_STRING = hasattr(np.dtypes, "StringDType") + + +@runtime_checkable +class SupportsStr(Protocol): + def __str__(self) -> str: ... + + +class LengthBytesConfig(TypedDict): + """ + Configuration for a fixed-length string data type in Zarr V3. + + Attributes + ---------- + length_bytes : int + The length in bytes of the data associated with this configuration. + """ + + length_bytes: int + + +class FixedLengthUTF32JSON_V2(DTypeConfig_V2[str, None]): + """ + A wrapper around the JSON representation of the ``FixedLengthUTF32`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + Examples + -------- + + ```python + { + "name": " None: + """ + We don't allow instances of this class with length less than 1 because there is no way such + a data type can contain actual data. + """ + if self.length < 1: + raise ValueError(f"length must be >= 1, got {self.length}.") + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create a FixedLengthUTF32 from a NumPy data type. + + Parameters + ---------- + dtype : TBaseDType + The NumPy data type. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_native_dtype(dtype): + endianness = get_endianness_from_numpy_dtype(dtype) + return cls( + length=dtype.itemsize // (cls.code_point_bytes), + endianness=endianness, + ) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.StrDType[int]: + """ + Convert the FixedLengthUTF32 instance to a NumPy data type. + + Returns + ------- + np.dtypes.StrDType[int] + The NumPy data type. + """ + byte_order = endianness_to_numpy_str(self.endianness) + return self.dtype_cls(self.length).newbyteorder(byte_order) + + @classmethod + def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[FixedLengthUTF32JSON_V2]: + """ + Check that the input is a valid JSON representation of a NumPy U dtype. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + TypeGuard[FixedLengthUTF32JSON_V2] + Whether the input is a valid JSON representation of a NumPy U dtype. + """ + return ( + check_dtype_spec_v2(data) + and isinstance(data["name"], str) + and re.match(r"^[><]U\d+$", data["name"]) is not None + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[FixedLengthUTF32JSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + TypeGuard[FixedLengthUTF32JSONV3] + Whether the input is a valid JSON representation of a NumPy U dtype. + """ + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and "configuration" in data + and isinstance(data["configuration"], dict) + and set(data["configuration"].keys()) == {"length_bytes"} + and isinstance(data["configuration"]["length_bytes"], int) + ) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[str, None]: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> FixedLengthUTF32JSON_V3: ... + + def to_json( + self, zarr_format: ZarrFormat + ) -> DTypeConfig_V2[str, None] | FixedLengthUTF32JSON_V3: + """ + Convert the FixedLengthUTF32 instance to a JSON representation. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format to use. + + Returns + ------- + DTypeConfig_V2[str, None] | FixedLengthUTF32JSON_V3 + The JSON representation of the data type. + """ + if zarr_format == 2: + return {"name": self.to_native_dtype().str, "object_codec_id": None} + elif zarr_format == 3: + v3_unstable_dtype_warning(self) + return { + "name": self._zarr_v3_name, + "configuration": {"length_bytes": self.length * self.code_point_bytes}, + } + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create a FixedLengthUTF32 from a JSON representation of a NumPy U dtype. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v2(data): + # Construct the NumPy dtype instead of string parsing. + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + raise DataTypeValidationError( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a string representation of a NumPy U dtype." + ) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create a FixedLengthUTF32 from a JSON representation of a NumPy U dtype. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v3(data): + return cls(length=data["configuration"]["length_bytes"] // cls.code_point_bytes) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected {cls._zarr_v3_name}." + raise DataTypeValidationError(msg) + + def default_scalar(self) -> np.str_: + """ + Return the default scalar value for this data type. + + Returns + ------- + ``np.str_`` + The default scalar value. + """ + return np.str_("") + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert the scalar value to a JSON representation. + + Parameters + ---------- + data : object + The scalar value. + zarr_format : ZarrFormat + The Zarr format to use. + + Returns + ------- + str + The JSON representation of the scalar value. + """ + return str(data) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.str_: + """ + Convert the JSON representation of a scalar value to the native scalar value. + + Parameters + ---------- + data : JSON + The JSON data. + zarr_format : ZarrFormat + The Zarr format to use. + + Returns + ------- + ``np.str_`` + The native scalar value. + """ + if check_json_str(data): + return self.to_native_dtype().type(data) + raise TypeError(f"Invalid type: {data}. Expected a string.") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[SupportsStr]: + """ + Check that the input is a valid scalar value for this data type. + + Parameters + ---------- + data : object + The scalar value. + + Returns + ------- + TypeGuard[SupportsStr] + Whether the input is a valid scalar value for this data type. + """ + # this is generous for backwards compatibility + return isinstance(data, SupportsStr) + + def cast_scalar(self, data: object) -> np.str_: + """ + Cast the scalar value to the native scalar value. + + Parameters + ---------- + data : object + The scalar value. + + Returns + ------- + ``np.str_`` + The native scalar value. + """ + if self._check_scalar(data): + # We explicitly truncate before casting because of the following NumPy behavior: + # >>> x = np.dtype('U3').type('hello world') + # >>> x + # np.str_('hello world') + # >>> x.dtype + # dtype('U11') + + return self.to_native_dtype().type(str(data)[: self.length]) + + msg = ( # pragma: no cover + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) # pragma: no-cover + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return self.length * self.code_point_bytes + + +def check_vlen_string_json_scalar(data: object) -> TypeGuard[int | str | float]: + """ + Check if the input is a valid JSON scalar for a variable-length string. + + This function is generous for backwards compatibility, as Zarr Python v2 would use ints for + variable-length string fill values. + + Parameters + ---------- + data : object + The JSON value to check. + + Returns + ------- + TypeGuard[int | str | float] + True if the input is a valid scalar for a variable-length string. + """ + return isinstance(data, int | str | float) + + +class VariableLengthUTF8JSON_V2(DTypeConfig_V2[Literal["|O"], Literal["vlen-utf8"]]): + """ + A wrapper around the JSON representation of the ``VariableLengthUTF8`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. The ``object_codec_id`` field is always ``"vlen-utf8"``. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + + Examples + -------- + ```python + { + "name": "|O", + "object_codec_id": "vlen-utf8" + } + ``` + """ + + +# VariableLengthUTF8 is defined in two places, conditioned on the version of NumPy. +# If NumPy 2 is installed, then VariableLengthUTF8 is defined with the NumPy variable length +# string dtype as the native dtype. Otherwise, VariableLengthUTF8 is defined with the NumPy object +# dtype as the native dtype. +class UTF8Base(ZDType[TDType_co, str], HasObjectCodec): + """ + A base class for variable-length UTF-8 string data types. + + Not intended for direct use, but as a base for concrete implementations. + + Attributes + ---------- + object_codec_id : ClassVar[Literal["vlen-utf8"]] + The object codec ID for this data type. + + References + ---------- + This data type does not have a Zarr V3 specification. + + The Zarr V2 data type specification can be found [here](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + """ + + _zarr_v3_name: ClassVar[Literal["string"]] = "string" + object_codec_id: ClassVar[Literal["vlen-utf8"]] = "vlen-utf8" + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of this data type from a compatible NumPy data type. + + + Parameters + ---------- + dtype : TBaseDType + The native data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input is not compatible with this data type. + """ + if cls._check_native_dtype(dtype): + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + @classmethod + def _check_json_v2( + cls, + data: DTypeJSON, + ) -> TypeGuard[VariableLengthUTF8JSON_V2]: + """ + "Check if the input is a valid JSON representation of a variable-length UTF-8 string dtype + for Zarr v2." + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + ``TypeGuard[VariableLengthUTF8JSON_V2]`` + Whether the input is a valid JSON representation of a NumPy "object" data type, and that the + object codec id is appropriate for variable-length UTF-8 strings. + """ + return ( + check_dtype_spec_v2(data) + and data["name"] == "|O" + and data["object_codec_id"] == cls.object_codec_id + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["variable_length_utf8"]]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[Literal["variable_length_utf8"]] + Whether the input is a valid JSON representation of a variable length UTF-8 string + data type. + """ + return data == cls._zarr_v3_name + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from a JSON representation of a NumPy "object" dtype. + + Parameters + ---------- + data : DTypeJSON + The JSON data to create an instance from. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v2(data): + return cls() + msg = ( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string '|O'" + ) + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this class from a JSON representation of a variable length UTF-8 + string data type. + + Parameters + ---------- + data : DTypeJSON + The JSON data to create an instance from. + + Returns + ------- + Self + An instance of this data type. + """ + if cls._check_json_v3(data): + return cls() + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected {cls._zarr_v3_name}." + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> VariableLengthUTF8JSON_V2: ... + @overload + def to_json(self, zarr_format: Literal[3]) -> Literal["string"]: ... + + def to_json(self, zarr_format: ZarrFormat) -> VariableLengthUTF8JSON_V2 | Literal["string"]: + """ + Convert this data type to a JSON representation. + + Parameters + ---------- + zarr_format : int + The zarr format to use for the JSON representation. + + Returns + ------- + ``VariableLengthUTF8JSON_V2 | Literal["string"]`` + The JSON representation of this data type. + """ + if zarr_format == 2: + return {"name": "|O", "object_codec_id": self.object_codec_id} + elif zarr_format == 3: + return self._zarr_v3_name + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def default_scalar(self) -> str: + """ + Return the default scalar value for this data type. + + Returns + ------- + str + The default scalar value. + """ + return "" + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert a scalar value to a JSON representation. + + Parameters + ---------- + data : object + The scalar value to convert. + zarr_format : int + The zarr format to use for the JSON representation. + + Returns + ------- + str + The JSON representation of the scalar value. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + raise TypeError(f"Invalid type: {data}. Expected a string.") + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> str: + """ + Convert a JSON representation of a scalar value to the native scalar type. + + Parameters + ---------- + data : JSON + The JSON representation of the scalar value. + zarr_format : int + The zarr format to use for the JSON representation. + + Returns + ------- + str + The native scalar type of the scalar value. + """ + if not check_vlen_string_json_scalar(data): + raise TypeError(f"Invalid type: {data}. Expected a string or number.") + return str(data) + + def _check_scalar(self, data: object) -> TypeGuard[SupportsStr]: + """ + Check that the input is a valid scalar value for this data type. + + Parameters + ---------- + data : object + The scalar value to check. + + Returns + ------- + TypeGuard[SupportsStr] + Whether the input is a valid scalar value for this data type. + """ + return isinstance(data, SupportsStr) + + def _cast_scalar_unchecked(self, data: SupportsStr) -> str: + """ + Cast a scalar value to a string. + + Parameters + ---------- + data : object + The scalar value to cast. + + Returns + ------- + str + The string representation of the scalar value. + """ + return str(data) + + def cast_scalar(self, data: object) -> str: + """ + Cast an object to a string. + + Parameters + ---------- + data : object + The value to cast. + + Returns + ------- + str + The input cast to str. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( # pragma: no cover + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) # pragma: no cover + + +if _NUMPY_SUPPORTS_VLEN_STRING: + + @dataclass(frozen=True, kw_only=True) + class VariableLengthUTF8(UTF8Base[np.dtypes.StringDType]): # type: ignore[type-var] + """ + A Zarr data type for arrays containing variable-length UTF-8 strings. + + Wraps the ``np.dtypes.StringDType`` data type. Scalars for this data type are instances + of ``str``. + + + Attributes + ---------- + dtype_cls : Type[np.dtypes.StringDType] + The NumPy dtype class for this data type. + _zarr_v3_name : ClassVar[Literal["variable_length_utf8"]] = "variable_length_utf8" + The name of this data type in Zarr V3. + object_codec_id : ClassVar[Literal["vlen-utf8"]] = "vlen-utf8" + The object codec ID for this data type. + """ + + dtype_cls = np.dtypes.StringDType + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of this data type from a compatible NumPy data type. + We reject NumPy StringDType instances that have the `na_object` field set, + because this is not representable by the Zarr `string` data type. + + Parameters + ---------- + dtype : TBaseDType + The native data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input is not compatible with this data type. + ValueError + If the input is `numpy.dtypes.StringDType` and has `na_object` set. + """ + if cls._check_native_dtype(dtype): + if hasattr(dtype, "na_object"): + msg = ( + f"Zarr data type resolution from {dtype} failed. " + "Attempted to resolve a zarr data type from a `numpy.dtypes.StringDType` " + "with `na_object` set, which is not supported." + ) + raise ValueError(msg) + return cls() + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.StringDType: + """ + Create a NumPy string dtype from this VariableLengthUTF8 ZDType. + + Returns + ------- + np.dtypes.StringDType + The NumPy string dtype. + """ + return self.dtype_cls() + +else: + # Numpy pre-2 does not have a variable length string dtype, so we use the Object dtype instead. + @dataclass(frozen=True, kw_only=True) + class VariableLengthUTF8(UTF8Base[np.dtypes.ObjectDType]): # type: ignore[no-redef] + """ + A Zarr data type for arrays containing variable-length UTF-8 strings. + + Wraps the ``np.dtypes.ObjectDType`` data type. Scalars for this data type are instances + of ``str``. + + + Attributes + ---------- + dtype_cls : Type[np.dtypes.ObjectDType] + The NumPy dtype class for this data type. + _zarr_v3_name : ClassVar[Literal["variable_length_utf8"]] = "variable_length_utf8" + The name of this data type in Zarr V3. + object_codec_id : ClassVar[Literal["vlen-utf8"]] = "vlen-utf8" + The object codec ID for this data type. + """ + + dtype_cls = np.dtypes.ObjectDType + + def to_native_dtype(self) -> np.dtypes.ObjectDType: + """ + Create a NumPy object dtype from this VariableLengthUTF8 ZDType. + + Returns + ------- + np.dtypes.ObjectDType + The NumPy object dtype. + """ + return self.dtype_cls() diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/structured.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/structured.py new file mode 100644 index 0000000000000000000000000000000000000000..8bedee07ef6e55ed59ccb1d6df90db61c9c2c555 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/structured.py @@ -0,0 +1,477 @@ +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar, Literal, Self, TypeGuard, cast, overload + +import numpy as np + +from zarr.core.common import NamedConfig +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasItemSize, + StructuredName_V2, + check_dtype_spec_v2, + check_structured_dtype_name_v2, + v3_unstable_dtype_warning, +) +from zarr.core.dtype.npy.common import ( + bytes_from_json, + bytes_to_json, + check_json_str, +) +from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + +StructuredScalarLike = list[object] | tuple[object, ...] | bytes | int + + +class StructuredJSON_V2(DTypeConfig_V2[StructuredName_V2, None]): + """ + A wrapper around the JSON representation of the ``Structured`` data type in Zarr V2. + + The ``name`` field is a sequence of sequences, where each inner sequence has two values: + the field name and the data type name for that field (which could be another sequence). + The data type names are strings, and the object codec ID is always None. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + Examples + -------- + ```python + { + "name": [ + ["f0", " None: + if len(self.fields) < 1: + raise ValueError(f"must have at least one field. Got {self.fields!r}") + + @classmethod + def _check_native_dtype(cls, dtype: TBaseDType) -> TypeGuard[np.dtypes.VoidDType[int]]: + """ + Check that this dtype is a numpy structured dtype + + Parameters + ---------- + dtype : np.dtypes.DTypeLike + The dtype to check. + + Returns + ------- + TypeGuard[np.dtypes.VoidDType] + True if the dtype matches, False otherwise. + """ + return isinstance(dtype, cls.dtype_cls) and dtype.fields is not None + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create a Structured ZDType from a native NumPy data type. + + Parameters + ---------- + dtype : TBaseDType + The native data type. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input data type is not an instance of np.dtypes.VoidDType with a non-null + ``fields`` attribute. + + Notes + ----- + This method attempts to resolve the fields of the structured dtype using the data type + registry. + """ + from zarr.core.dtype import get_data_type_from_native_dtype + + fields: list[tuple[str, ZDType[TBaseDType, TBaseScalar]]] = [] + if cls._check_native_dtype(dtype): + # fields of a structured numpy dtype are either 2-tuples or 3-tuples. we only + # care about the first element in either case. + for key, (dtype_instance, *_) in dtype.fields.items(): # type: ignore[union-attr] + dtype_wrapped = get_data_type_from_native_dtype(dtype_instance) + fields.append((key, dtype_wrapped)) + + return cls(fields=tuple(fields)) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> np.dtypes.VoidDType[int]: + """ + Convert the structured Zarr data type to a native NumPy void dtype. + + This method constructs a NumPy dtype with fields corresponding to the + fields of the structured Zarr data type, by converting each field's + data type to its native dtype representation. + + Returns + ------- + np.dtypes.VoidDType[int] + The native NumPy void dtype representing the structured data type. + """ + + return cast( + "np.dtypes.VoidDType[int]", + np.dtype([(key, dtype.to_native_dtype()) for (key, dtype) in self.fields]), + ) + + @classmethod + def _check_json_v2( + cls, + data: DTypeJSON, + ) -> TypeGuard[StructuredJSON_V2]: + """ + Check if the input is a valid JSON representation of a Structured data type + for Zarr V2. + + The input data must be a mapping that contains a "name" key that is not a str, + and an "object_codec_id" key that is None. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[StructuredJSON_V2] + True if the input is a valid JSON representation of a Structured data type + for Zarr V2, False otherwise. + """ + return ( + check_dtype_spec_v2(data) + and not isinstance(data["name"], str) + and check_structured_dtype_name_v2(data["name"]) + and data["object_codec_id"] is None + ) + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[StructuredJSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[StructuredJSON_V3] + True if the input is a valid JSON representation of a structured data type for Zarr V3, + False otherwise. + """ + + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and isinstance(data["configuration"], dict) + and set(data["configuration"].keys()) == {"fields"} + ) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + # avoid circular import + from zarr.core.dtype import get_data_type_from_json + + if cls._check_json_v2(data): + # structured dtypes are constructed directly from a list of lists + # note that we do not handle the object codec here! this will prevent structured + # dtypes from containing object dtypes. + return cls( + fields=tuple( # type: ignore[misc] + ( # type: ignore[misc] + f_name, + get_data_type_from_json( + {"name": f_dtype, "object_codec_id": None}, zarr_format=2 + ), + ) + for f_name, f_dtype in data["name"] + ) + ) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a JSON array of arrays" + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + # avoid circular import + from zarr.core.dtype import get_data_type_from_json + + if cls._check_json_v3(data): + config = data["configuration"] + meta_fields = config["fields"] + return cls( + fields=tuple( + (f_name, get_data_type_from_json(f_dtype, zarr_format=3)) # type: ignore[misc] + for f_name, f_dtype in meta_fields + ) + ) + msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a JSON object with the key {cls._zarr_v3_name!r}" + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> StructuredJSON_V2: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> StructuredJSON_V3: ... + + def to_json(self, zarr_format: ZarrFormat) -> StructuredJSON_V2 | StructuredJSON_V3: + """ + Convert the structured data type to a JSON-serializable form. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version. Accepted values are 2 and 3. + + Returns + ------- + StructuredJSON_V2 | StructuredJSON_V3 + The JSON representation of the structured data type. + + Raises + ------ + ValueError + If the zarr_format is not 2 or 3. + """ + if zarr_format == 2: + fields = [ + [f_name, f_dtype.to_json(zarr_format=zarr_format)["name"]] + for f_name, f_dtype in self.fields + ] + return {"name": fields, "object_codec_id": None} + elif zarr_format == 3: + v3_unstable_dtype_warning(self) + fields = [ + [f_name, f_dtype.to_json(zarr_format=zarr_format)] # type: ignore[list-item] + for f_name, f_dtype in self.fields + ] + base_dict = { + "name": self._zarr_v3_name, + "configuration": {"fields": fields}, + } + return cast("StructuredJSON_V3", base_dict) + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[StructuredScalarLike]: + # TODO: implement something more precise here! + """ + Check that the input is a valid scalar value for this structured data type. + + Parameters + ---------- + data : object + The scalar value to check. + + Returns + ------- + TypeGuard[StructuredScalarLike] + Whether the input is a valid scalar value for this structured data type. + """ + return isinstance(data, (bytes, list, tuple, int, np.void)) + + def _cast_scalar_unchecked(self, data: StructuredScalarLike) -> np.void: + """ + Cast a python object to a numpy structured scalar without type checking. + + Parameters + ---------- + data : StructuredScalarLike + The data to cast. + + Returns + ------- + np.void + The casted data as a numpy structured scalar. + + Notes + ----- + This method does not perform any type checking. + The input data must be castable to a numpy structured scalar. + + """ + na_dtype = self.to_native_dtype() + if isinstance(data, bytes): + res = np.frombuffer(data, dtype=na_dtype)[0] + elif isinstance(data, list | tuple): + res = np.array([tuple(data)], dtype=na_dtype)[0] + else: + res = np.array([data], dtype=na_dtype)[0] + return cast("np.void", res) + + def cast_scalar(self, data: object) -> np.void: + """ + Cast a Python object to a NumPy structured scalar. + + This function attempts to cast the provided data to a NumPy structured scalar. + If the data is compatible with the structured scalar type, it is cast without + type checking. Otherwise, a TypeError is raised. + + Parameters + ---------- + data : object + The data to be cast to a NumPy structured scalar. + + Returns + ------- + np.void + The data cast as a NumPy structured scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a NumPy structured scalar. + """ + + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> np.void: + """ + Get the default scalar value for this structured data type. + + Returns + ------- + np.void + The default scalar value, which is the scalar representation of 0 + cast to this structured data type. + """ + + return self._cast_scalar_unchecked(0) + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.void: + """ + Read a JSON-serializable value as a NumPy structured scalar. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + np.void + The NumPy structured scalar. + + Raises + ------ + TypeError + If the input is not a base64-encoded string. + """ + if check_json_str(data): + as_bytes = bytes_from_json(data, zarr_format=zarr_format) + dtype = self.to_native_dtype() + return cast("np.void", np.array([as_bytes]).view(dtype)[0]) + raise TypeError(f"Invalid type: {data}. Expected a string.") + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str: + """ + Convert a scalar to a JSON-serializable string representation. + + Parameters + ---------- + data : object + The scalar to convert. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + str + A string representation of the scalar, which is a base64-encoded + string of the bytes that make up the scalar. + """ + return bytes_to_json(self.cast_scalar(data).tobytes(), zarr_format) + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return self.to_native_dtype().itemsize diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/time.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/time.py new file mode 100644 index 0000000000000000000000000000000000000000..402a1403214c282b472dc2956f41ec0ce5522133 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/npy/time.py @@ -0,0 +1,885 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import ( + TYPE_CHECKING, + ClassVar, + Literal, + Self, + TypedDict, + TypeGuard, + TypeVar, + cast, + get_args, + overload, +) + +import numpy as np +from typing_extensions import ReadOnly + +from zarr.core.common import NamedConfig +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeConfig_V2, + DTypeJSON, + HasEndianness, + HasItemSize, + check_dtype_spec_v2, +) +from zarr.core.dtype.npy.common import ( + DATETIME_UNIT, + DateTimeUnit, + check_json_int, + endianness_to_numpy_str, + get_endianness_from_numpy_dtype, +) +from zarr.core.dtype.wrapper import TBaseDType, ZDType + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + +TimeDeltaLike = str | int | bytes | np.timedelta64 | timedelta | None +DateTimeLike = str | int | bytes | np.datetime64 | datetime | None + + +def datetime_from_int(data: int, *, unit: DateTimeUnit, scale_factor: int) -> np.datetime64: + """ + Convert an integer to a datetime64. + + Parameters + ---------- + data : int + The integer to convert. + unit : DateTimeUnit + The unit of the datetime64. + scale_factor : int + The scale factor of the datetime64. + + Returns + ------- + numpy.datetime64 + The datetime64 value. + """ + dtype_name = f"datetime64[{scale_factor}{unit}]" + return cast("np.datetime64", np.int64(data).view(dtype_name)) + + +def datetimelike_to_int(data: np.datetime64 | np.timedelta64) -> int: + """ + Convert a datetime64 or a timedelta64 to an integer. + + Parameters + ---------- + data : np.datetime64 | numpy.timedelta64 + The value to convert. + + Returns + ------- + int + An integer representation of the scalar. + """ + return data.view(np.int64).item() + + +def check_json_time(data: JSON) -> TypeGuard[Literal["NaT"] | int]: + """ + Type guard to check if the input JSON data is the literal string "NaT" + or an integer. + """ + return check_json_int(data) or data == "NaT" + + +BaseTimeDType_co = TypeVar( + "BaseTimeDType_co", + bound=np.dtypes.TimeDelta64DType | np.dtypes.DateTime64DType, + covariant=True, +) +BaseTimeScalar_co = TypeVar( + "BaseTimeScalar_co", bound=np.timedelta64 | np.datetime64, covariant=True +) + + +class TimeConfig(TypedDict): + """ + The configuration for the numpy.timedelta64 or numpy.datetime64 data type in Zarr V3. + + Attributes + ---------- + unit : ReadOnly[DateTimeUnit] + A string encoding a unit of time. + scale_factor : ReadOnly[int] + A scale factor. + + Examples + -------- + ```python + {"unit": "ms", "scale_factor": 1} + ``` + """ + + unit: ReadOnly[DateTimeUnit] + scale_factor: ReadOnly[int] + + +class DateTime64JSON_V3(NamedConfig[Literal["numpy.datetime64"], TimeConfig]): + """ + The JSON representation of the ``numpy.datetime64`` data type in Zarr V3. + + References + ---------- + This representation is defined in the ``numpy.datetime64`` + [specification document](https://zarr-specs.readthedocs.io/en/latest/spec/v3/datatypes.html#numpy-datetime64). + + Examples + -------- + ```python + { + "name": "numpy.datetime64", + "configuration": { + "unit": "ms", + "scale_factor": 1 + } + } + ``` + """ + + +class TimeDelta64JSON_V3(NamedConfig[Literal["numpy.timedelta64"], TimeConfig]): + """ + The JSON representation of the ``TimeDelta64`` data type in Zarr V3. + + References + ---------- + This representation is defined in the numpy.timedelta64 + [specification document](https://zarr-specs.readthedocs.io/en/latest/spec/v3/datatypes.html#numpy-timedelta64). + + Examples + -------- + ```python + { + "name": "numpy.timedelta64", + "configuration": { + "unit": "ms", + "scale_factor": 1 + } + } + ``` + """ + + +class TimeDelta64JSON_V2(DTypeConfig_V2[str, None]): + """ + A wrapper around the JSON representation of the ``TimeDelta64`` data type in Zarr V2. + + The ``name`` field of this class contains the value that would appear under the + ``dtype`` field in Zarr V2 array metadata. + + References + ---------- + The structure of the ``name`` field is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + + Examples + -------- + ```python + { + "name": " None: + if self.scale_factor < 1: + raise ValueError(f"scale_factor must be > 0, got {self.scale_factor}.") + if self.scale_factor >= 2**31: + raise ValueError(f"scale_factor must be < 2147483648, got {self.scale_factor}.") + if self.unit not in get_args(DateTimeUnit): + raise ValueError(f"unit must be one of {get_args(DateTimeUnit)}, got {self.unit!r}.") + + @classmethod + def from_native_dtype(cls, dtype: TBaseDType) -> Self: + """ + Create an instance of this class from a native NumPy data type. + + Parameters + ---------- + dtype : TBaseDType + The native NumPy dtype to convert. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the dtype is not a valid representation of this class. + """ + + if cls._check_native_dtype(dtype): + unit, scale_factor = np.datetime_data(dtype.name) + unit = cast("DateTimeUnit", unit) + return cls( + unit=unit, + scale_factor=scale_factor, + endianness=get_endianness_from_numpy_dtype(dtype), + ) + raise DataTypeValidationError( + f"Invalid data type: {dtype}. Expected an instance of {cls.dtype_cls}" + ) + + def to_native_dtype(self) -> BaseTimeDType_co: + # Numpy does not allow creating datetime64 or timedelta64 via + # np.dtypes.{dtype_name}() + # so we use np.dtype with a formatted string. + """ + Convert this data type to a NumPy temporal data type with the appropriate + unit and scale factor. + + Returns + ------- + BaseTimeDType_co + A NumPy data type object representing the time data type with + the specified unit, scale factor, and byte order. + """ + + dtype_string = f"{self._numpy_name}[{self.scale_factor}{self.unit}]" + return np.dtype(dtype_string).newbyteorder(endianness_to_numpy_str(self.endianness)) # type: ignore[return-value] + + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> int: + """ + Convert a python object to a JSON representation of a datetime64 or timedelta64 scalar. + + Parameters + ---------- + data : object + The python object to convert. + zarr_format : ZarrFormat + The Zarr format version (2 or 3). + + Returns + ------- + int + The JSON representation of the scalar. + """ + return datetimelike_to_int(data) # type: ignore[arg-type] + + @property + def item_size(self) -> int: + """ + The size of a single scalar in bytes. + + Returns + ------- + int + The size of a single scalar in bytes. + """ + return 8 + + +@dataclass(frozen=True, kw_only=True, slots=True) +class TimeDelta64(TimeDTypeBase[np.dtypes.TimeDelta64DType, np.timedelta64], HasEndianness): + """ + A Zarr data type for arrays containing NumPy TimeDelta64 data. + + Wraps the ``np.dtypesTimeDelta64DType`` data type. Scalars for this data type + are instances of `np.timedelta64`. + + Attributes + ---------- + dtype_cls : Type[np.dtypesTimeDelta64DType] + The NumPy dtype class for this data type. + scale_factor : int + The scale factor for this data type. + unit : DateTimeUnit + The unit for this data type. + + References + ---------- + The Zarr V2 representation of this data type is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + The Zarr V3 representation of this data type is defined in the ``numpy.timedelta64`` + [specification document](https://github.com/zarr-developers/zarr-extensions/tree/main/data-types/numpy.timedelta64) + """ + + # mypy infers the type of np.dtypes.TimeDelta64DType to be + # "Callable[[Literal['Y', 'M', 'W', 'D'] | Literal['h', 'm', 's', 'ms', 'us', 'ns', 'ps', 'fs', 'as']], Never]" + dtype_cls = np.dtypes.TimeDelta64DType # type: ignore[assignment] + unit: DateTimeUnit = "generic" + scale_factor: int = 1 + _zarr_v3_name: ClassVar[Literal["numpy.timedelta64"]] = "numpy.timedelta64" + _zarr_v2_names: ClassVar[tuple[Literal[">m8"], Literal["m8", " TypeGuard[TimeDelta64JSON_V2]: + """ + Validate that the provided JSON input accurately represents a NumPy timedelta64 data type, + which could be in the form of strings like "m8[10s]". This method serves as a type + guard, helping to refine the type of unknown JSON input by confirming its adherence to the + expected format for NumPy timedelta64 data types. + + The JSON input should contain a "name" key with a value that matches the expected string + pattern for NumPy timedelta64 data types. The pattern includes an optional unit enclosed + within square brackets, following the base type identifier. + + Returns + ------- + bool + True if the JSON input is a valid representation of this class, + otherwise False. + """ + if not check_dtype_spec_v2(data): + return False + name = data["name"] + # match m[M], etc + # consider making this a standalone function + if not isinstance(name, str): + return False + if not name.startswith(cls._zarr_v2_names): + return False + if len(name) == 3: + # no unit, and + # we already checked that this string is either m8 + return True + else: + return name[4:-1].endswith(DATETIME_UNIT) and name[-1] == "]" + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[DateTime64JSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Returns + ------- + TypeGuard[DateTime64JSON_V3] + True if the JSON input is a valid representation of this class, + otherwise False. + """ + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and isinstance(data["configuration"], dict) + and set(data["configuration"].keys()) == {"unit", "scale_factor"} + ) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create a TimeDelta64 from a Zarr V2-flavored JSON. + + Parameters + ---------- + data : DTypeJSON + The JSON data. + + Returns + ------- + TimeDelta64 + An instance of TimeDelta64. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v2(data): + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = ( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a string " + f"representation of an instance of {cls.dtype_cls}" + ) + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create a TimeDelta64 from a Zarr V3-flavored JSON. + + The JSON representation of a TimeDelta64 in Zarr V3 is a dict with a 'name' key + with the value 'numpy.timedelta64', and a 'configuration' key with a value of a dict + with a 'unit' key and a 'scale_factor' key. + + For example: + + ```json + { + "name": "numpy.timedelta64", + "configuration": { + "unit": "generic", + "scale_factor": 1 + } + } + ``` + + """ + if cls._check_json_v3(data): + unit = data["configuration"]["unit"] + scale_factor = data["configuration"]["scale_factor"] + return cls(unit=unit, scale_factor=scale_factor) + msg = ( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a dict " + f"with a 'name' key with the value 'numpy.timedelta64', " + "and a 'configuration' key with a value of a dict with a 'unit' key and a " + "'scale_factor' key" + ) + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> TimeDelta64JSON_V2: ... + @overload + def to_json(self, zarr_format: Literal[3]) -> TimeDelta64JSON_V3: ... + + def to_json(self, zarr_format: ZarrFormat) -> TimeDelta64JSON_V2 | TimeDelta64JSON_V3: + """ + Serialize this data type to JSON. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version (2 or 3). + + Returns + ------- + TimeDelta64JSON_V2 | TimeDelta64JSON_V3 + The JSON representation of the data type. + + Raises + ------ + ValueError + If the zarr_format is not 2 or 3. + """ + if zarr_format == 2: + name = self.to_native_dtype().str + return {"name": name, "object_codec_id": None} + elif zarr_format == 3: + return { + "name": self._zarr_v3_name, + "configuration": {"unit": self.unit, "scale_factor": self.scale_factor}, + } + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[TimeDeltaLike]: + """ + Check if the input is a scalar of this data type. + + Parameters + ---------- + data : object + The object to check. + + Returns + ------- + TypeGuard[TimeDeltaLike] + True if the input is a scalar of this data type, False otherwise. + """ + if data is None: + return True + return isinstance(data, str | int | bytes | np.timedelta64 | timedelta) + + def _cast_scalar_unchecked(self, data: TimeDeltaLike) -> np.timedelta64: + """ + Cast the provided scalar input to a numpy timedelta64 without any type checking. + + This method assumes that the input data is already a valid scalar of this data type, + and does not perform any validation or type checks. It directly casts the input + to a numpy timedelta64 scalar using the unit and scale factor defined in the class. + + Parameters + ---------- + data : TimeDeltaLike + The scalar input data to cast. + + Returns + ------- + numpy.timedelta64 + The input data cast as a numpy timedelta64 scalar. + """ + return self.to_native_dtype().type(data, f"{self.scale_factor}{self.unit}") + + def cast_scalar(self, data: object) -> np.timedelta64: + """ + Cast the input to a numpy timedelta64 scalar. If the input is not a scalar of this data type, + raise a TypeError. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> np.timedelta64: + """ + Return a default scalar of this data type. + + This method provides a default value for the timedelta64 scalar, which is + a 'Not-a-Time' (NaT) value. + """ + return np.timedelta64("NaT") + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.timedelta64: + """ + Create a scalar of this data type from JSON input. + + Parameters + ---------- + data : JSON + The JSON representation of the scalar value. + zarr_format : int + The zarr format to use for the JSON representation. + + Returns + ------- + numpy.timedelta64 + The scalar value of this data type. + + Raises + ------ + TypeError + If the input JSON is not a valid representation of a scalar for this data type. + """ + if check_json_time(data): + return self.to_native_dtype().type(data, f"{self.scale_factor}{self.unit}") + raise TypeError(f"Invalid type: {data}. Expected an integer.") # pragma: no cover + + +@dataclass(frozen=True, kw_only=True, slots=True) +class DateTime64(TimeDTypeBase[np.dtypes.DateTime64DType, np.datetime64], HasEndianness): + """ + A Zarr data type for arrays containing NumPy Datetime64 data. + + Wraps the ``np.dtypes.TimeDelta64DType`` data type. Scalars for this data type + are instances of ``np.datetime64``. + + Attributes + ---------- + dtype_cls : Type[np.dtypesTimeDelta64DType] + The numpy dtype class for this data type. + unit : DateTimeUnit + The unit of time for this data type. + scale_factor : int + The scale factor for the time unit. + + References + ---------- + The Zarr V2 representation of this data type is defined in the Zarr V2 + [specification document](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v2/v2.0.rst#data-type-encoding). + + The Zarr V3 representation of this data type is defined in the ``numpy.datetime64`` + [specification document](https://github.com/zarr-developers/zarr-extensions/tree/main/data-types/numpy.datetime64) + """ + + dtype_cls = np.dtypes.DateTime64DType # type: ignore[assignment] + _zarr_v3_name: ClassVar[Literal["numpy.datetime64"]] = "numpy.datetime64" + _zarr_v2_names: ClassVar[tuple[Literal[">M8"], Literal["M8", " TypeGuard[DateTime64JSON_V2]: + """ + Check that the input is a valid JSON representation of this data type. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[DateTime64JSON_V2] + True if the input is a valid JSON representation of a NumPy datetime64 data type, + otherwise False. + """ + if not check_dtype_spec_v2(data): + return False + name = data["name"] + if not isinstance(name, str): + return False + if not name.startswith(cls._zarr_v2_names): + return False + if len(name) == 3: + # no unit, and + # we already checked that this string is either M8 + return True + else: + return name[4:-1].endswith(DATETIME_UNIT) and name[-1] == "]" + + @classmethod + def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[DateTime64JSON_V3]: + """ + Check that the input is a valid JSON representation of this class in Zarr V3. + + Parameters + ---------- + data : DTypeJSON + The JSON data to check. + + Returns + ------- + TypeGuard[DateTime64JSON_V3] + True if the input is a valid JSON representation of a numpy datetime64 data type in Zarr V3, False otherwise. + """ + + return ( + isinstance(data, dict) + and set(data.keys()) == {"name", "configuration"} + and data["name"] == cls._zarr_v3_name + and isinstance(data["configuration"], dict) + and set(data["configuration"].keys()) == {"unit", "scale_factor"} + ) + + @classmethod + def _from_json_v2(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from a Zarr V2-flavored JSON representation. + + This method checks if the provided JSON data is a valid representation of this class. + If valid, it creates an instance using the native NumPy dtype. Otherwise, it raises a + DataTypeValidationError. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + + if cls._check_json_v2(data): + name = data["name"] + return cls.from_native_dtype(np.dtype(name)) + msg = ( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a string " + f"representation of an instance of {cls.dtype_cls}" + ) + raise DataTypeValidationError(msg) + + @classmethod + def _from_json_v3(cls, data: DTypeJSON) -> Self: + """ + Create an instance of this data type from a Zarr V3-flavored JSON representation. + + This method checks if the provided JSON data is a valid representation of this class. + If valid, it creates an instance using the native NumPy dtype. Otherwise, it raises a + DataTypeValidationError. + + Parameters + ---------- + data : DTypeJSON + The JSON data to parse. + + Returns + ------- + Self + An instance of this data type. + + Raises + ------ + DataTypeValidationError + If the input JSON is not a valid representation of this class. + """ + if cls._check_json_v3(data): + unit = data["configuration"]["unit"] + scale_factor = data["configuration"]["scale_factor"] + return cls(unit=unit, scale_factor=scale_factor) + msg = ( + f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected a dict " + f"with a 'name' key with the value 'numpy.datetime64', " + "and a 'configuration' key with a value of a dict with a 'unit' key and a " + "'scale_factor' key" + ) + raise DataTypeValidationError(msg) + + @overload + def to_json(self, zarr_format: Literal[2]) -> DateTime64JSON_V2: ... + @overload + def to_json(self, zarr_format: Literal[3]) -> DateTime64JSON_V3: ... + + def to_json(self, zarr_format: ZarrFormat) -> DateTime64JSON_V2 | DateTime64JSON_V3: + """ + Serialize this data type to JSON. + + Parameters + ---------- + zarr_format : ZarrFormat + The Zarr format version (2 or 3). + + Returns + ------- + DateTime64JSON_V2 | DateTime64JSON_V3 + The JSON representation of the data type. + + Raises + ------ + ValueError + If the zarr_format is not 2 or 3. + """ + if zarr_format == 2: + name = self.to_native_dtype().str + return {"name": name, "object_codec_id": None} + elif zarr_format == 3: + return { + "name": self._zarr_v3_name, + "configuration": {"unit": self.unit, "scale_factor": self.scale_factor}, + } + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + def _check_scalar(self, data: object) -> TypeGuard[DateTimeLike]: + """ + Check if the input is convertible to a scalar of this data type. + + Parameters + ---------- + data : object + The object to check. + + Returns + ------- + TypeGuard[DateTimeLike] + True if the input is a scalar of this data type, False otherwise. + """ + if data is None: + return True + return isinstance(data, str | int | bytes | np.datetime64 | datetime) + + def _cast_scalar_unchecked(self, data: DateTimeLike) -> np.datetime64: + """ + Cast the input to a scalar of this data type without any type checking. + + Parameters + ---------- + data : DateTimeLike + The scalar data to cast. + + Returns + ------- + numpy.datetime64 + The input cast to a NumPy datetime scalar. + """ + return self.to_native_dtype().type(data, f"{self.scale_factor}{self.unit}") + + def cast_scalar(self, data: object) -> np.datetime64: + """ + Cast the input to a scalar of this data type after a type check. + + Parameters + ---------- + data : object + The scalar value to cast. + + Returns + ------- + numpy.datetime64 + The input cast to a NumPy datetime scalar. + + Raises + ------ + TypeError + If the data cannot be converted to a numpy datetime scalar. + """ + if self._check_scalar(data): + return self._cast_scalar_unchecked(data) + msg = ( + f"Cannot convert object {data!r} with type {type(data)} to a scalar compatible with the " + f"data type {self}." + ) + raise TypeError(msg) + + def default_scalar(self) -> np.datetime64: + """ + Return the default scalar value for this data type. + + Returns + ------- + numpy.datetime64 + The default scalar value, which is a 'Not-a-Time' (NaT) value + """ + + return np.datetime64("NaT") + + def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> np.datetime64: + """ + Read a JSON-serializable value as a scalar. + + Parameters + ---------- + data : JSON + The JSON-serializable value. + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + numpy.datetime64 + The numpy datetime scalar. + + Raises + ------ + TypeError + If the input is not a valid integer type. + """ + if check_json_time(data): + return self._cast_scalar_unchecked(data) + raise TypeError(f"Invalid type: {data}. Expected an integer.") # pragma: no cover diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/registry.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..315945cf4ec19668c1251d0c6a3863cc62f3109b --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/registry.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import contextlib +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Self + +import numpy as np + +from zarr.core.dtype.common import ( + DataTypeValidationError, + DTypeJSON, +) + +if TYPE_CHECKING: + from importlib.metadata import EntryPoint + + from zarr.core.common import ZarrFormat + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + + +# This class is different from the other registry classes, which inherit from +# dict. IMO it's simpler to just do a dataclass. But long-term we should +# have just 1 registry class in use. +@dataclass(frozen=True, kw_only=True) +class DataTypeRegistry: + """ + A registry for ZDType classes. + + This registry is a mapping from Zarr data type names to their + corresponding ZDType classes. + + Attributes + ---------- + contents : dict[str, type[ZDType[TBaseDType, TBaseScalar]]] + The mapping from Zarr data type names to their corresponding + ZDType classes. + """ + + contents: dict[str, type[ZDType[TBaseDType, TBaseScalar]]] = field( + default_factory=dict, init=False + ) + + _lazy_load_list: list[EntryPoint] = field(default_factory=list, init=False) + + def _lazy_load(self) -> None: + """ + Load all data types from the lazy load list and register them with + the registry. After loading, clear the lazy load list. + """ + for e in self._lazy_load_list: + self.register(e.load()._zarr_v3_name, e.load()) + + self._lazy_load_list.clear() + + def register(self: Self, key: str, cls: type[ZDType[TBaseDType, TBaseScalar]]) -> None: + """ + Register a data type with the registry. + + Parameters + ---------- + key : str + The Zarr V3 name of the data type. + cls : type[ZDType[TBaseDType, TBaseScalar]] + The class of the data type to register. + + Notes + ----- + This method is idempotent. If the data type is already registered, this + method does nothing. + """ + if key not in self.contents or self.contents[key] != cls: + self.contents[key] = cls + + def unregister(self, key: str) -> None: + """ + Unregister a data type from the registry. + + Parameters + ---------- + key : str + The key associated with the ZDType class to be unregistered. + + Returns + ------- + None + + Raises + ------ + KeyError + If the data type is not found in the registry. + """ + if key in self.contents: + del self.contents[key] + else: + raise KeyError(f"Data type '{key}' not found in registry.") + + def get(self, key: str) -> type[ZDType[TBaseDType, TBaseScalar]]: + """ + Retrieve a registered ZDType class by its key. + + Parameters + ---------- + key : str + The key associated with the desired ZDType class. + + Returns + ------- + type[ZDType[TBaseDType, TBaseScalar]] + The ZDType class registered under the given key. + + Raises + ------ + KeyError + If the key is not found in the registry. + """ + + return self.contents[key] + + def match_dtype(self, dtype: TBaseDType) -> ZDType[TBaseDType, TBaseScalar]: + """ + Match a native data type, e.g. a NumPy data type, to a registered ZDType. + + Parameters + ---------- + dtype : TBaseDType + The native data type to match. + + Returns + ------- + ZDType[TBaseDType, TBaseScalar] + The matched ZDType corresponding to the provided NumPy data type. + + Raises + ------ + ValueError + If the data type is a NumPy "Object" type, which is ambiguous, or if multiple + or no Zarr data types are found that match the provided dtype. + + Notes + ----- + This function attempts to resolve a Zarr data type from a given native data type. + If the dtype is a NumPy "Object" data type, it raises a ValueError, as this type + can represent multiple Zarr data types. In such cases, a specific Zarr data type + should be explicitly constructed instead of relying on dynamic resolution. + + If multiple matches are found, it will also raise a ValueError. In this case + conflicting data types must be unregistered, or the Zarr data type should be explicitly + constructed. + """ + + if dtype == np.dtype("O"): + msg = ( + f"Zarr data type resolution from {dtype} failed. " + 'Attempted to resolve a zarr data type from a numpy "Object" data type, which is ' + 'ambiguous, as multiple zarr data types can be represented by the numpy "Object" ' + "data type. " + "In this case you should construct your array by providing a specific Zarr data " + 'type. For a list of Zarr data types that are compatible with the numpy "Object"' + "data type, see https://github.com/zarr-developers/zarr-python/issues/3117" + ) + raise ValueError(msg) + matched: list[ZDType[TBaseDType, TBaseScalar]] = [] + for val in self.contents.values(): + # DataTypeValidationError means "this dtype doesn't match me", which is + # expected and suppressed. Other exceptions (e.g. ValueError for a dtype + # that matches the type but has an invalid configuration) are propagated + # to the caller. + with contextlib.suppress(DataTypeValidationError): + matched.append(val.from_native_dtype(dtype)) + if len(matched) == 1: + return matched[0] + elif len(matched) > 1: + msg = ( + f"Zarr data type resolution from {dtype} failed. " + f"Multiple data type wrappers found that match dtype '{dtype}': {matched}. " + "You should unregister one of these data types, or avoid Zarr data type inference " + "entirely by providing a specific Zarr data type when creating your array." + "For more information, see https://github.com/zarr-developers/zarr-python/issues/3117" + ) + raise ValueError(msg) + raise ValueError(f"No Zarr data type found that matches dtype '{dtype!r}'") + + def match_json( + self, data: DTypeJSON, *, zarr_format: ZarrFormat + ) -> ZDType[TBaseDType, TBaseScalar]: + """ + Match a JSON representation of a data type to a registered ZDType. + + Parameters + ---------- + data : DTypeJSON + The JSON representation of a data type to match. + zarr_format : ZarrFormat + The Zarr format version to consider when matching data types. + + Returns + ------- + ZDType[TBaseDType, TBaseScalar] + The matched ZDType corresponding to the JSON representation. + + Raises + ------ + ValueError + If no matching Zarr data type is found for the given JSON data. + """ + + for val in self.contents.values(): + try: + return val.from_json(data, zarr_format=zarr_format) + except DataTypeValidationError: + pass + raise ValueError(f"No Zarr data type found that matches {data!r}") diff --git a/.venv/lib/python3.11/site-packages/zarr/core/dtype/wrapper.py b/.venv/lib/python3.11/site-packages/zarr/core/dtype/wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..fdc5f747f0fdbd5d1f0500cc181364358b3a356e --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/dtype/wrapper.py @@ -0,0 +1,300 @@ +""" +Wrapper for native array data types. + +The ``ZDType`` class is an abstract base class for wrapping native array data types, e.g. NumPy dtypes. +``ZDType`` provides a common interface for working with data types in a way that is independent of the +underlying data type system. + +The wrapper class encapsulates a native data type. Instances of the class can be created from a +native data type instance, and a native data type instance can be created from an instance of the +wrapper class. + +The wrapper class is responsible for: +- Serializing and deserializing a native data type to Zarr V2 or Zarr V3 metadata. + This ensures that the data type can be properly stored and retrieved from array metadata. +- Serializing and deserializing scalar values to Zarr V2 or Zarr V3 metadata. This is important for + storing a fill value for an array in a manner that is valid for the data type. + +You can add support for a new data type in Zarr by subclassing ``ZDType`` wrapper class and adapt its methods +to support your native data type. The wrapper class must be added to a data type registry +(defined elsewhere) before array creation routines or array reading routines can use your new data +type. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + ClassVar, + Generic, + Literal, + Self, + TypeGuard, + TypeVar, + overload, +) + +import numpy as np + +if TYPE_CHECKING: + from zarr.core.common import JSON, ZarrFormat + from zarr.core.dtype.common import DTypeJSON, DTypeSpec_V2, DTypeSpec_V3 + +# This the upper bound for the scalar types we support. It's numpy scalars + str, +# because the new variable-length string dtype in numpy does not have a corresponding scalar type +TBaseScalar = np.generic | str | bytes +# This is the bound for the dtypes that we support. If we support non-numpy dtypes, +# then this bound will need to be widened. +TBaseDType = np.dtype[np.generic] + +# These two type parameters are covariant because we want +# x : ZDType[BaseDType, BaseScalar] = ZDType[SubDType, SubScalar] +# to type check +TScalar_co = TypeVar("TScalar_co", bound=TBaseScalar, covariant=True) +TDType_co = TypeVar("TDType_co", bound=TBaseDType, covariant=True) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class ZDType(ABC, Generic[TDType_co, TScalar_co]): + """ + Abstract base class for wrapping native array data types, e.g. numpy dtypes + + Attributes + ---------- + dtype_cls : ClassVar[type[TDType]] + The wrapped dtype class. This is a class variable. + _zarr_v3_name : ClassVar[str] + The name given to the data type by a Zarr v3 data type specification. This is a + class variable, and it should generally be unique across different data types. + """ + + # this class will create a native data type + dtype_cls: ClassVar[type[TDType_co]] + _zarr_v3_name: ClassVar[str] + + @classmethod + def _check_native_dtype(cls: type[Self], dtype: TBaseDType) -> TypeGuard[TDType_co]: + """ + Check that a native data type matches the dtype_cls class attribute. + + Used as a type guard. + + Parameters + ---------- + dtype : TDType + The dtype to check. + + Returns + ------- + Bool + True if the dtype matches, False otherwise. + """ + return type(dtype) is cls.dtype_cls + + @classmethod + @abstractmethod + def from_native_dtype(cls: type[Self], dtype: TBaseDType) -> Self: + """ + Create a ZDType instance from a native data type. + + This method is used when taking a user-provided native data type, like a NumPy data type, + and creating the corresponding ZDType instance from them. + + Parameters + ---------- + dtype : TDType + The native data type object to wrap. + + Returns + ------- + Self + The ZDType that wraps the native data type. + + Raises + ------ + TypeError + If the native data type is not consistent with the wrapped data type. + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def to_native_dtype(self: Self) -> TDType_co: + """ + Return an instance of the wrapped data type. This operation inverts ``from_native_dtype``. + + Returns + ------- + TDType + The native data type wrapped by this ZDType. + """ + raise NotImplementedError # pragma: no cover + + @classmethod + @abstractmethod + def _from_json_v2(cls: type[Self], data: DTypeJSON) -> Self: + raise NotImplementedError # pragma: no cover + + @classmethod + @abstractmethod + def _from_json_v3(cls: type[Self], data: DTypeJSON) -> Self: + raise NotImplementedError # pragma: no cover + + @classmethod + def from_json(cls: type[Self], data: DTypeJSON, *, zarr_format: ZarrFormat) -> Self: + """ + Create an instance of this ZDType from JSON data. + + Parameters + ---------- + data : DTypeJSON + The JSON representation of the data type. + + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + Self + An instance of this data type. + """ + if zarr_format == 2: + return cls._from_json_v2(data) + if zarr_format == 3: + return cls._from_json_v3(data) + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + @overload + def to_json(self, zarr_format: Literal[2]) -> DTypeSpec_V2: ... + + @overload + def to_json(self, zarr_format: Literal[3]) -> DTypeSpec_V3: ... + + @abstractmethod + def to_json(self, zarr_format: ZarrFormat) -> DTypeSpec_V2 | DTypeSpec_V3: + """ + Serialize this ZDType to JSON. + + Parameters + ---------- + zarr_format : ZarrFormat + The zarr format version. + + Returns + ------- + DTypeJSON_V2 | DTypeJSON_V3 + The JSON-serializable representation of the wrapped data type + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def _check_scalar(self, data: object) -> bool: + """ + Check that a python object is a valid scalar value for the wrapped data type. + + Parameters + ---------- + data : object + A value to check. + + Returns + ------- + Bool + True if the object is valid, False otherwise. + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def cast_scalar(self, data: object) -> TScalar_co: + """ + Cast a python object to the wrapped scalar type. + + The type of the provided scalar is first checked for compatibility. + If it's incompatible with the associated scalar type, a ``TypeError`` will be raised. + + Parameters + ---------- + data : object + The python object to cast. + + Returns + ------- + TScalar + The cast value. + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def default_scalar(self) -> TScalar_co: + """ + Get the default scalar value for the wrapped data type. + + This is a method, rather than an attribute, because the default value for some data types depends on parameters that are + not known until a concrete data type is wrapped. For example, data types parametrized by a + length like fixed-length strings or bytes will generate scalars consistent with that length. + + Returns + ------- + TScalar + The default value for this data type. + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def from_json_scalar(self: Self, data: JSON, *, zarr_format: ZarrFormat) -> TScalar_co: + """ + Read a JSON-serializable value as a scalar. + + Parameters + ---------- + data : JSON + A JSON representation of a scalar value. + zarr_format : ZarrFormat + The zarr format version. This is specified because the JSON serialization of scalars + differs between Zarr V2 and Zarr V3. + + Returns + ------- + TScalar + The deserialized scalar value. + """ + raise NotImplementedError # pragma: no cover + + @abstractmethod + def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> JSON: + """ + Serialize a python object to the JSON representation of a scalar. + + The value will first be cast to the scalar type associated with this ZDType, then serialized + to JSON. + + Parameters + ---------- + data : object + The value to convert. + zarr_format : ZarrFormat + The zarr format version. This is specified because the JSON serialization of scalars + differs between Zarr V2 and Zarr V3. + + Returns + ------- + JSON + The JSON-serialized scalar. + """ + raise NotImplementedError # pragma: no cover + + +def scalar_failed_type_check_msg( + cls_instance: ZDType[TBaseDType, TBaseScalar], bad_scalar: object +) -> str: + """ + Generate an error message reporting that a particular value failed a type check when attempting + to cast that value to a scalar. + """ + return ( + f"The value {bad_scalar!r} failed a type check. " + f"It cannot be safely cast to a scalar compatible with {cls_instance}. " + f"Consult the documentation for {cls_instance} to determine the possible values that can " + "be cast to scalars of the wrapped data type." + ) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/group.py b/.venv/lib/python3.11/site-packages/zarr/core/group.py new file mode 100644 index 0000000000000000000000000000000000000000..080e90ff0f19c4d537050427b957190193569dca --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/group.py @@ -0,0 +1,3910 @@ +from __future__ import annotations + +import asyncio +import itertools +import json +import logging +import unicodedata +import warnings +from collections import defaultdict +from dataclasses import asdict, dataclass, field, fields, replace +from itertools import accumulate +from typing import TYPE_CHECKING, Literal, TypeVar, assert_never, cast, overload + +import numpy as np +import numpy.typing as npt +from typing_extensions import deprecated + +import zarr.api.asynchronous as async_api +from zarr.abc.metadata import Metadata +from zarr.abc.store import Store, set_or_delete +from zarr.core._info import GroupInfo +from zarr.core.array import ( + DEFAULT_FILL_VALUE, + Array, + AsyncArray, + CompressorLike, + CompressorsLike, + FiltersLike, + SerializerLike, + ShardsLike, + _parse_deprecated_compressor, + create_array, +) +from zarr.core.attributes import Attributes +from zarr.core.buffer import default_buffer_prototype +from zarr.core.common import ( + JSON, + ZARR_JSON, + ZARRAY_JSON, + ZATTRS_JSON, + ZGROUP_JSON, + ZMETADATA_V2_JSON, + DimensionNames, + NodeType, + ShapeLike, + ZarrFormat, + parse_shapelike, +) +from zarr.core.config import config +from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.io import save_metadata +from zarr.core.sync import SyncMixin, sync +from zarr.errors import ( + ContainsArrayError, + ContainsGroupError, + GroupNotFoundError, + MetadataValidationError, + ZarrDeprecationWarning, + ZarrUserWarning, +) +from zarr.storage import StoreLike, StorePath +from zarr.storage._common import ensure_no_existing_node, make_store_path +from zarr.storage._utils import _join_paths, _normalize_path_keys, normalize_path + +if TYPE_CHECKING: + from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Coroutine, + Generator, + Iterable, + Iterator, + Mapping, + ) + from typing import Any + + from zarr.core.array_spec import ArrayConfigLike + from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike + from zarr.core.common import MemoryOrder + from zarr.core.dtype import ZDTypeLike + from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 + +logger = logging.getLogger("zarr.group") + +DefaultT = TypeVar("DefaultT") + + +def parse_zarr_format(data: Any) -> ZarrFormat: + """Parse the zarr_format field from metadata.""" + if data in (2, 3): + return cast("ZarrFormat", data) + msg = f"Invalid zarr_format. Expected one of 2 or 3. Got {data}." + raise ValueError(msg) + + +def parse_node_type(data: Any) -> NodeType: + """Parse the node_type field from metadata.""" + if data in ("array", "group"): + return cast("Literal['array', 'group']", data) + msg = f"Invalid value for 'node_type'. Expected 'array' or 'group'. Got '{data}'." + raise MetadataValidationError(msg) + + +# todo: convert None to empty dict +def parse_attributes(data: Any) -> dict[str, Any]: + """Parse the attributes field from metadata.""" + if data is None: + return {} + elif isinstance(data, dict) and all(isinstance(k, str) for k in data): + return data + msg = f"Expected dict with string keys. Got {type(data)} instead." + raise TypeError(msg) + + +@overload +def _parse_async_node(node: AsyncArrayV3) -> ArrayV3: ... + + +@overload +def _parse_async_node(node: AsyncArrayV2) -> ArrayV2: ... + + +@overload +def _parse_async_node(node: AsyncGroup) -> Group: ... + + +def _parse_async_node( + node: AnyAsyncArray | AsyncGroup, +) -> AnyArray | Group: + """Wrap an AsyncArray in an Array, or an AsyncGroup in a Group.""" + if isinstance(node, AsyncArray): + return Array(node) + elif isinstance(node, AsyncGroup): + return Group(node) + else: + raise TypeError(f"Unknown node type, got {type(node)}") + + +@dataclass(frozen=True) +class ConsolidatedMetadata: + """ + Consolidated Metadata for this Group. + + This stores the metadata of child nodes below this group. Any child groups + will have their consolidated metadata set appropriately. + """ + + metadata: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] + kind: Literal["inline"] = "inline" + must_understand: Literal[False] = False + + def to_dict(self) -> dict[str, JSON]: + return { + "kind": self.kind, + "must_understand": self.must_understand, + "metadata": { + k: v.to_dict() + for k, v in sorted( + self.flattened_metadata.items(), + key=lambda item: ( + item[0].count("/"), + unicodedata.normalize("NFKC", item[0]).casefold(), + ), + ) + }, + } + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> ConsolidatedMetadata: + data = dict(data) + + kind = data.get("kind") + if kind != "inline": + raise ValueError(f"Consolidated metadata kind='{kind}' is not supported.") + + raw_metadata = data.get("metadata") + if not isinstance(raw_metadata, dict): + raise TypeError(f"Unexpected type for 'metadata': {type(raw_metadata)}") + + metadata: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] = {} + if raw_metadata: + for k, v in raw_metadata.items(): + if not isinstance(v, dict): + raise TypeError( + f"Invalid value for metadata items. key='{k}', type='{type(v).__name__}'" + ) + + # zarr_format is present in v2 and v3. + zarr_format = parse_zarr_format(v["zarr_format"]) + + if zarr_format == 3: + node_type = parse_node_type(v.get("node_type", None)) + if node_type == "group": + metadata[k] = GroupMetadata.from_dict(v) + elif node_type == "array": + metadata[k] = ArrayV3Metadata.from_dict(v) + else: + assert_never(node_type) + elif zarr_format == 2: + if "shape" in v: + metadata[k] = ArrayV2Metadata.from_dict(v) + else: + metadata[k] = GroupMetadata.from_dict(v) + else: + assert_never(zarr_format) + + cls._flat_to_nested(metadata) + + return cls(metadata=metadata) + + @staticmethod + def _flat_to_nested( + metadata: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], + ) -> None: + """ + Convert a flat metadata representation to a nested one. + + Notes + ----- + Flat metadata is used when persisting the consolidated metadata. The keys + include the full path, not just the node name. The key prefixes can be + used to determine which nodes are children of which other nodes. + + Nested metadata is used in-memory. The outermost level will only have the + *immediate* children of the Group. All nested child groups will be stored + under the consolidated metadata of their immediate parent. + """ + # We have a flat mapping from {k: v} where the keys include the *full* + # path segment: + # { + # "/a/b": { group_metadata }, + # "/a/b/array-0": { array_metadata }, + # "/a/b/array-1": { array_metadata }, + # } + # + # We want to reorganize the metadata such that each Group contains the + # array metadata of its immediate children. + # In the example, the group at `/a/b` will have consolidated metadata + # for its children `array-0` and `array-1`. + # + # metadata = dict(metadata) + + keys = sorted(metadata, key=lambda k: k.count("/")) + grouped = { + k: list(v) for k, v in itertools.groupby(keys, key=lambda k: k.rsplit("/", 1)[0]) + } + + # we go top down and directly manipulate metadata. + for key, children_keys in grouped.items(): + # key is a key like "a", "a/b", "a/b/c" + # The basic idea is to find the immediate parent (so "", "a", or "a/b") + # and update that node's consolidated metadata to include the metadata + # in children_keys + *prefixes, name = key.split("/") + parent = metadata + + while prefixes: + # e.g. a/b/c has a parent "a/b". Walk through to get + # metadata["a"]["b"] + part = prefixes.pop(0) + # we can assume that parent[part] here is a group + # otherwise we wouldn't have a node with this `part` prefix. + # We can also assume that the parent node will have consolidated metadata, + # because we're walking top to bottom. + parent = parent[part].consolidated_metadata.metadata # type: ignore[union-attr] + + node = parent[name] + children_keys = list(children_keys) + + if isinstance(node, ArrayV2Metadata | ArrayV3Metadata): + # These are already present, either thanks to being an array in the + # root, or by being collected as a child in the else clause + continue + children_keys = list(children_keys) + # We pop from metadata, since we're *moving* this under group + children = { + child_key.split("/")[-1]: metadata.pop(child_key) + for child_key in children_keys + if child_key != key + } + parent[name] = replace( + node, consolidated_metadata=ConsolidatedMetadata(metadata=children) + ) + + @property + def flattened_metadata(self) -> dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata]: + """ + Return the flattened representation of Consolidated Metadata. + + The returned dictionary will have a key for each child node in the hierarchy + under this group. Under the default (nested) representation available through + ``self.metadata``, the dictionary only contains keys for immediate children. + + The keys of the dictionary will include the full path to a child node from + the current group, where segments are joined by ``/``. + + Examples + -------- + ```python + from zarr.core.group import ConsolidatedMetadata, GroupMetadata + cm = ConsolidatedMetadata( + metadata={ + "group-0": GroupMetadata( + consolidated_metadata=ConsolidatedMetadata( + { + "group-0-0": GroupMetadata(), + } + ) + ), + "group-1": GroupMetadata(), + } + ) + # {'group-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'), + # 'group-0/group-0-0': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'), + # 'group-1': GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group')} + ``` + """ + metadata = {} + + def flatten( + key: str, group: GroupMetadata | ArrayV2Metadata | ArrayV3Metadata + ) -> dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata]: + children: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] = {} + if isinstance(group, ArrayV2Metadata | ArrayV3Metadata): + children[key] = group + else: + if group.consolidated_metadata and group.consolidated_metadata.metadata is not None: + children[key] = replace( + group, consolidated_metadata=ConsolidatedMetadata(metadata={}) + ) + for name, val in group.consolidated_metadata.metadata.items(): + full_key = f"{key}/{name}" + if isinstance(val, GroupMetadata): + children.update(flatten(full_key, val)) + else: + children[full_key] = val + else: + children[key] = replace(group, consolidated_metadata=None) + return children + + for k, v in self.metadata.items(): + metadata.update(flatten(k, v)) + + return metadata + + +@dataclass(frozen=True) +class GroupMetadata(Metadata): + """ + Metadata for a Group. + """ + + attributes: dict[str, Any] = field(default_factory=dict) + zarr_format: ZarrFormat = 3 + consolidated_metadata: ConsolidatedMetadata | None = None + node_type: Literal["group"] = field(default="group", init=False) + + def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + json_indent = config.get("json_indent") + if self.zarr_format == 3: + return { + ZARR_JSON: prototype.buffer.from_bytes( + json.dumps(self.to_dict(), indent=json_indent, allow_nan=True).encode() + ) + } + else: + items = { + ZGROUP_JSON: prototype.buffer.from_bytes( + json.dumps({"zarr_format": self.zarr_format}, indent=json_indent).encode() + ), + ZATTRS_JSON: prototype.buffer.from_bytes( + json.dumps(self.attributes, indent=json_indent, allow_nan=True).encode() + ), + } + if self.consolidated_metadata: + d = { + ZGROUP_JSON: {"zarr_format": self.zarr_format}, + ZATTRS_JSON: self.attributes, + } + consolidated_metadata = self.consolidated_metadata.to_dict()["metadata"] + assert isinstance(consolidated_metadata, dict) + for k, v in consolidated_metadata.items(): + attrs = v.pop("attributes", {}) + d[f"{k}/{ZATTRS_JSON}"] = attrs + if "shape" in v: + # it's an array + d[f"{k}/{ZARRAY_JSON}"] = v + else: + d[f"{k}/{ZGROUP_JSON}"] = { + "zarr_format": self.zarr_format, + "consolidated_metadata": { + "metadata": {}, + "must_understand": False, + "kind": "inline", + }, + } + + items[ZMETADATA_V2_JSON] = prototype.buffer.from_bytes( + json.dumps( + {"metadata": d, "zarr_consolidated_format": 1}, allow_nan=True + ).encode() + ) + + return items + + def __init__( + self, + attributes: dict[str, Any] | None = None, + zarr_format: ZarrFormat = 3, + consolidated_metadata: ConsolidatedMetadata | None = None, + ) -> None: + attributes_parsed = parse_attributes(attributes) + zarr_format_parsed = parse_zarr_format(zarr_format) + + object.__setattr__(self, "attributes", attributes_parsed) + object.__setattr__(self, "zarr_format", zarr_format_parsed) + object.__setattr__(self, "consolidated_metadata", consolidated_metadata) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> GroupMetadata: + data = dict(data) + assert data.pop("node_type", None) in ("group", None) + consolidated_metadata = data.pop("consolidated_metadata", None) + if consolidated_metadata: + data["consolidated_metadata"] = ConsolidatedMetadata.from_dict(consolidated_metadata) + + zarr_format = data.get("zarr_format") + if zarr_format == 2 or zarr_format is None: + # zarr v2 allowed arbitrary keys here. + # We don't want the GroupMetadata constructor to fail just because someone put an + # extra key in the metadata. + expected = {x.name for x in fields(cls)} + data = {k: v for k, v in data.items() if k in expected} + + return cls(**data) + + def to_dict(self) -> dict[str, Any]: + result = asdict(replace(self, consolidated_metadata=None)) + if self.consolidated_metadata is not None: + result["consolidated_metadata"] = self.consolidated_metadata.to_dict() + else: + # Leave consolidated metadata unset if it's None + result.pop("consolidated_metadata") + return result + + +@dataclass(frozen=True) +class ImplicitGroupMarker(GroupMetadata): + """ + Marker for an implicit group. Instances of this class are only used in the context of group + creation as a placeholder to represent groups that should only be created if they do not + already exist in storage + """ + + +@dataclass(frozen=True) +class AsyncGroup: + """ + Asynchronous Group object. + """ + + metadata: GroupMetadata + store_path: StorePath + + # TODO: make this correct and work + # TODO: ensure that this can be bound properly to subclass of AsyncGroup + + @classmethod + async def from_store( + cls, + store: StoreLike, + *, + attributes: dict[str, Any] | None = None, + overwrite: bool = False, + zarr_format: ZarrFormat = 3, + ) -> AsyncGroup: + store_path = await make_store_path(store) + + if overwrite: + if store_path.store.supports_deletes: + await store_path.delete_dir() + else: + await ensure_no_existing_node(store_path, zarr_format=zarr_format) + else: + await ensure_no_existing_node(store_path, zarr_format=zarr_format) + attributes = attributes or {} + group = cls( + metadata=GroupMetadata(attributes=attributes, zarr_format=zarr_format), + store_path=store_path, + ) + await group._save_metadata(ensure_parents=True) + return group + + @classmethod + async def open( + cls, + store: StoreLike, + zarr_format: ZarrFormat | None = 3, + use_consolidated: bool | str | None = None, + ) -> AsyncGroup: + """Open a new AsyncGroup + + Parameters + ---------- + store : StoreLike + zarr_format : {2, 3}, optional + use_consolidated : bool or str, default None + Whether to use consolidated metadata. + + By default, consolidated metadata is used if it's present in the + store (in the ``zarr.json`` for Zarr format 3 and in the ``.zmetadata`` file + for Zarr format 2) and the Store supports it. + + To explicitly require consolidated metadata, set ``use_consolidated=True``. + In this case, if the Store doesn't support consolidation or consolidated metadata is + not found, a ``ValueError`` exception is raised. + + To explicitly *not* use consolidated metadata, set ``use_consolidated=False``, + which will fall back to using the regular, non consolidated metadata. + + Zarr format 2 allowed configuring the key storing the consolidated metadata + (``.zmetadata`` by default). Specify the custom key as ``use_consolidated`` + to load consolidated metadata from a non-default key. + """ + store_path = await make_store_path(store) + if not store_path.store.supports_consolidated_metadata: + # Fail if consolidated metadata was requested but the Store doesn't support it + if use_consolidated: + store_name = type(store_path.store).__name__ + raise ValueError( + f"The Zarr store in use ({store_name}) doesn't support consolidated metadata." + ) + + # if use_consolidated was None (optional), the Store dictates it doesn't want consolidation + use_consolidated = False + + consolidated_key = ZMETADATA_V2_JSON + + if (zarr_format == 2 or zarr_format is None) and isinstance(use_consolidated, str): + consolidated_key = use_consolidated + + if zarr_format == 2: + paths = [store_path / ZGROUP_JSON, store_path / ZATTRS_JSON] + if use_consolidated or use_consolidated is None: + paths.append(store_path / consolidated_key) + + zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather( + *[path.get() for path in paths] + ) + if zgroup_bytes is None: + raise FileNotFoundError(store_path) + + if use_consolidated or use_consolidated is None: + maybe_consolidated_metadata_bytes = rest[0] + + else: + maybe_consolidated_metadata_bytes = None + + elif zarr_format == 3: + zarr_json_bytes = await (store_path / ZARR_JSON).get() + if zarr_json_bytes is None: + raise FileNotFoundError(store_path) + elif zarr_format is None: + ( + zarr_json_bytes, + zgroup_bytes, + zattrs_bytes, + maybe_consolidated_metadata_bytes, + ) = await asyncio.gather( + (store_path / ZARR_JSON).get(), + (store_path / ZGROUP_JSON).get(), + (store_path / ZATTRS_JSON).get(), + (store_path / str(consolidated_key)).get(), + ) + if zarr_json_bytes is not None and zgroup_bytes is not None: + # warn and favor v3 + msg = f"Both zarr.json (Zarr format 3) and .zgroup (Zarr format 2) metadata objects exist at {store_path}. Zarr format 3 will be used." + warnings.warn(msg, category=ZarrUserWarning, stacklevel=1) + if zarr_json_bytes is None and zgroup_bytes is None: + raise FileNotFoundError( + f"could not find zarr.json or .zgroup objects in {store_path}" + ) + # set zarr_format based on which keys were found + if zarr_json_bytes is not None: + zarr_format = 3 + else: + zarr_format = 2 + else: + msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable] + raise MetadataValidationError(msg) + + if zarr_format == 2: + # this is checked above, asserting here for mypy + assert zgroup_bytes is not None + + if use_consolidated and maybe_consolidated_metadata_bytes is None: + # the user requested consolidated metadata, but it was missing + raise ValueError(consolidated_key) + + elif use_consolidated is False: + # the user explicitly opted out of consolidated_metadata. + # Discard anything we might have read. + maybe_consolidated_metadata_bytes = None + + return cls._from_bytes_v2( + store_path, zgroup_bytes, zattrs_bytes, maybe_consolidated_metadata_bytes + ) + else: + # V3 groups are comprised of a zarr.json object + assert zarr_json_bytes is not None + if not isinstance(use_consolidated, bool | None): + raise TypeError("use_consolidated must be a bool or None for Zarr format 3.") + + return cls._from_bytes_v3( + store_path, + zarr_json_bytes, + use_consolidated=use_consolidated, + ) + + @classmethod + def _from_bytes_v2( + cls, + store_path: StorePath, + zgroup_bytes: Buffer, + zattrs_bytes: Buffer | None, + consolidated_metadata_bytes: Buffer | None, + ) -> AsyncGroup: + # V2 groups are comprised of a .zgroup and .zattrs objects + zgroup = json.loads(zgroup_bytes.to_bytes()) + zattrs = json.loads(zattrs_bytes.to_bytes()) if zattrs_bytes is not None else {} + group_metadata = {**zgroup, "attributes": zattrs} + + if consolidated_metadata_bytes is not None: + v2_consolidated_metadata = json.loads(consolidated_metadata_bytes.to_bytes()) + v2_consolidated_metadata = v2_consolidated_metadata["metadata"] + # We already read zattrs and zgroup. Should we ignore these? + v2_consolidated_metadata.pop(".zattrs", None) + v2_consolidated_metadata.pop(".zgroup", None) + + consolidated_metadata: defaultdict[str, dict[str, Any]] = defaultdict(dict) + + # keys like air/.zarray, air/.zattrs + for k, v in v2_consolidated_metadata.items(): + path, kind = k.rsplit("/.", 1) + + if kind == "zarray": + consolidated_metadata[path].update(v) + elif kind == "zattrs": + consolidated_metadata[path]["attributes"] = v + elif kind == "zgroup": + consolidated_metadata[path].update(v) + else: + raise ValueError(f"Invalid file type '{kind}' at path '{path}") + + group_metadata["consolidated_metadata"] = { + "metadata": dict(consolidated_metadata), + "kind": "inline", + "must_understand": False, + } + + return cls.from_dict(store_path, group_metadata) + + @classmethod + def _from_bytes_v3( + cls, + store_path: StorePath, + zarr_json_bytes: Buffer, + use_consolidated: bool | None, + ) -> AsyncGroup: + group_metadata = json.loads(zarr_json_bytes.to_bytes()) + if use_consolidated and group_metadata.get("consolidated_metadata") is None: + msg = f"Consolidated metadata requested with 'use_consolidated=True' but not found in '{store_path.path}'." + raise ValueError(msg) + + elif use_consolidated is False: + # Drop consolidated metadata if it's there. + group_metadata.pop("consolidated_metadata", None) + + return cls.from_dict(store_path, group_metadata) + + @classmethod + def from_dict( + cls, + store_path: StorePath, + data: dict[str, Any], + ) -> AsyncGroup: + node_type = data.pop("node_type", None) + if node_type == "array": + msg = f"An array already exists in store {store_path.store} at path {store_path.path}." + raise ContainsArrayError(msg) + elif node_type not in ("group", None): + msg = f"Node type in metadata ({node_type}) is not 'group'" + raise GroupNotFoundError(msg) + return cls( + metadata=GroupMetadata.from_dict(data), + store_path=store_path, + ) + + async def setitem(self, key: str, value: Any) -> None: + """ + Fastpath for creating a new array + New arrays will be created with default array settings for the array type. + + Parameters + ---------- + key : str + Array name + value : array-like + Array data + """ + path = self.store_path / key + await async_api.save_array( + store=path, arr=value, zarr_format=self.metadata.zarr_format, overwrite=True + ) + + async def getitem( + self, + key: str, + ) -> AnyAsyncArray | AsyncGroup: + """ + Get a subarray or subgroup from the group. + + Parameters + ---------- + key : str + Array or group name + + Returns + ------- + AsyncArray or AsyncGroup + """ + store_path = self.store_path / key + logger.debug("key=%s, store_path=%s", key, store_path) + + # Consolidated metadata lets us avoid some I/O operations so try that first. + if self.metadata.consolidated_metadata is not None: + return self._getitem_consolidated(store_path, key, prefix=self.name) + try: + return await get_node( + store=store_path.store, path=store_path.path, zarr_format=self.metadata.zarr_format + ) + except FileNotFoundError as e: + raise KeyError(key) from e + + def _getitem_consolidated( + self, store_path: StorePath, key: str, prefix: str + ) -> AnyAsyncArray | AsyncGroup: + # getitem, in the special case where we have consolidated metadata. + # Note that this is a regular def (non async) function. + # This shouldn't do any additional I/O. + + # the caller needs to verify this! + assert self.metadata.consolidated_metadata is not None + + # we support nested getitems like group/subgroup/array + indexers = normalize_path(key).split("/") + indexers.reverse() + metadata: ArrayV2Metadata | ArrayV3Metadata | GroupMetadata = self.metadata + + while indexers: + indexer = indexers.pop() + if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): + # we've indexed into an array with group["array/subarray"]. Invalid. + raise KeyError(key) + if metadata.consolidated_metadata is None: + # we've indexed into a group without consolidated metadata. + # This isn't normal; typically, consolidated metadata + # will include explicit markers for when there are no child + # nodes as metadata={}. + # We have some freedom in exactly how we interpret this case. + # For now, we treat None as the same as {}, i.e. we don't + # have any children. + raise KeyError(key) + try: + metadata = metadata.consolidated_metadata.metadata[indexer] + except KeyError as e: + # The Group Metadata has consolidated metadata, but the key + # isn't present. We trust this to mean that the key isn't in + # the hierarchy, and *don't* fall back to checking the store. + msg = f"'{key}' not found in consolidated metadata." + raise KeyError(msg) from e + + # update store_path to ensure that AsyncArray/Group.name is correct + if prefix != "/": + key = "/".join([prefix.lstrip("/"), key]) + store_path = StorePath(store=store_path.store, path=key) + + if isinstance(metadata, GroupMetadata): + return AsyncGroup(metadata=metadata, store_path=store_path) + else: + return AsyncArray(metadata=metadata, store_path=store_path) + + async def delitem(self, key: str) -> None: + """Delete a group member. + + Parameters + ---------- + key : str + Array or group name + """ + store_path = self.store_path / key + + await store_path.delete_dir() + if self.metadata.consolidated_metadata: + self.metadata.consolidated_metadata.metadata.pop(key, None) + await self._save_metadata() + + async def get( + self, key: str, default: DefaultT | None = None + ) -> AnyAsyncArray | AsyncGroup | DefaultT | None: + """Obtain a group member, returning default if not found. + + Parameters + ---------- + key : str + Group member name. + default : object + Default value to return if key is not found (default: None). + + Returns + ------- + object + Group member (AsyncArray or AsyncGroup) or default if not found. + """ + try: + return await self.getitem(key) + except KeyError: + return default + + async def _save_metadata(self, ensure_parents: bool = False) -> None: + await save_metadata(self.store_path, self.metadata, ensure_parents=ensure_parents) + + @property + def path(self) -> str: + """Storage path.""" + return self.store_path.path + + @property + def name(self) -> str: + """Group name following h5py convention.""" + if self.path: + # follow h5py convention: add leading slash + name = self.path + if name[0] != "/": + name = "/" + name + return name + return "/" + + @property + def basename(self) -> str: + """Final component of name.""" + return self.name.split("/")[-1] + + @property + def attrs(self) -> dict[str, Any]: + return self.metadata.attributes + + @property + def info(self) -> Any: + """ + Return a visual representation of the statically known information about a group. + + Note that this doesn't include dynamic information, like the number of child + Groups or Arrays. + + Returns + ------- + GroupInfo + + Related + ------- + [zarr.AsyncGroup.info_complete][] + All information about a group, including dynamic information + """ + + if self.metadata.consolidated_metadata: + members = list(self.metadata.consolidated_metadata.flattened_metadata.values()) + else: + members = None + return self._info(members=members) + + async def info_complete(self) -> Any: + """ + Return all the information for a group. + + This includes dynamic information like the number + of child Groups or Arrays. If this group doesn't contain consolidated + metadata then this will need to read from the backing Store. + + Returns + ------- + GroupInfo + + Related + ------- + [zarr.AsyncGroup.info][] + """ + members = [x[1].metadata async for x in self.members(max_depth=None)] + return self._info(members=members) + + def _info( + self, members: list[ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] | None = None + ) -> Any: + kwargs = {} + if members is not None: + kwargs["_count_members"] = len(members) + count_arrays = 0 + count_groups = 0 + for member in members: + if isinstance(member, GroupMetadata): + count_groups += 1 + else: + count_arrays += 1 + kwargs["_count_arrays"] = count_arrays + kwargs["_count_groups"] = count_groups + + return GroupInfo( + _name=self.store_path.path, + _read_only=self.read_only, + _store_type=type(self.store_path.store).__name__, + _zarr_format=self.metadata.zarr_format, + # maybe do a typeddict + **kwargs, # type: ignore[arg-type] + ) + + @property + def store(self) -> Store: + return self.store_path.store + + @property + def read_only(self) -> bool: + # Backwards compatibility for 2.x + return self.store_path.read_only + + @property + def synchronizer(self) -> None: + # Backwards compatibility for 2.x + # Not implemented in 3.x yet. + return None + + async def create_group( + self, + name: str, + *, + overwrite: bool = False, + attributes: dict[str, Any] | None = None, + ) -> AsyncGroup: + """Create a sub-group. + + Parameters + ---------- + name : str + Group name. + overwrite : bool, optional + If True, do not raise an error if the group already exists. + attributes : dict, optional + Group attributes. + + Returns + ------- + g : AsyncGroup + """ + attributes = attributes or {} + return await type(self).from_store( + self.store_path / name, + attributes=attributes, + overwrite=overwrite, + zarr_format=self.metadata.zarr_format, + ) + + async def require_group(self, name: str, overwrite: bool = False) -> AsyncGroup: + """Obtain a sub-group, creating one if it doesn't exist. + + Parameters + ---------- + name : str + Group name. + overwrite : bool, optional + Overwrite any existing group with given `name` if present. + + Returns + ------- + g : AsyncGroup + """ + if overwrite: + # TODO: check that overwrite=True errors if an array exists where the group is being created + grp = await self.create_group(name, overwrite=True) + else: + try: + item: AsyncGroup | AnyAsyncArray = await self.getitem(name) + if not isinstance(item, AsyncGroup): + raise TypeError( + f"Incompatible object ({item.__class__.__name__}) already exists" + ) + assert isinstance(item, AsyncGroup) # make mypy happy + grp = item + except KeyError: + grp = await self.create_group(name) + return grp + + async def require_groups(self, *names: str) -> tuple[AsyncGroup, ...]: + """Convenience method to require multiple groups in a single call. + + Parameters + ---------- + *names : str + Group names. + + Returns + ------- + Tuple[AsyncGroup, ...] + """ + if not names: + return () + return tuple(await asyncio.gather(*(self.require_group(name) for name in names))) + + async def create_array( + self, + name: str, + *, + shape: ShapeLike | None = None, + dtype: ZDTypeLike | None = None, + data: np.ndarray[Any, np.dtype[Any]] | None = None, + chunks: tuple[int, ...] | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + compressor: CompressorLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNames = None, + storage_options: dict[str, Any] | None = None, + overwrite: bool = False, + config: ArrayConfigLike | None = None, + write_data: bool = True, + ) -> AnyAsyncArray: + """Create an array within this group. + + This method lightly wraps [zarr.core.array.create_array][]. + + Parameters + ---------- + name : str + The name of the array relative to the group. If ``path`` is ``None``, the array will be located + at the root of the store. + shape : tuple[int, ...] + Shape of the array. + dtype : npt.DTypeLike + Data type of the array. + chunks : tuple[int, ...], optional + Chunk shape of the array. + If not specified, default are guessed based on the shape and dtype. + shards : tuple[int, ...], optional + Shard shape of the array. The default value of ``None`` results in no sharding at all. + filters : Iterable[Codec] | Literal["auto"], optional + Iterable of filters to apply to each chunk of the array, in order, before serializing that + chunk to bytes. + + For Zarr format 3, a "filter" is a codec that takes an array and returns an array, + and these values must be instances of [`zarr.abc.codec.ArrayArrayCodec`][], or a + dict representations of [`zarr.abc.codec.ArrayArrayCodec`][]. + + For Zarr format 2, a "filter" can be any numcodecs codec; you should ensure that the + the order if your filters is consistent with the behavior of each filter. + + The default value of ``"auto"`` instructs Zarr to use a default used based on the data + type of the array and the Zarr format specified. For all data types in Zarr V3, and most + data types in Zarr V2, the default filters are empty. The only cases where default filters + are not empty is when the Zarr format is 2, and the data type is a variable-length data type like + [`zarr.dtype.VariableLengthUTF8`][] or [`zarr.dtype.VariableLengthUTF8`][]. In these cases, + the default filters contains a single element which is a codec specific to that particular data type. + + To create an array with no filters, provide an empty iterable or the value ``None``. + compressors : Iterable[Codec], optional + List of compressors to apply to the array. Compressors are applied in order, and after any + filters are applied (if any are specified) and the data is serialized into bytes. + + For Zarr format 3, a "compressor" is a codec that takes a bytestream, and + returns another bytestream. Multiple compressors my be provided for Zarr format 3. + If no ``compressors`` are provided, a default set of compressors will be used. + These defaults can be changed by modifying the value of ``array.v3_default_compressors`` + in [`zarr.config`][zarr.config]. + Use ``None`` to omit default compressors. + + For Zarr format 2, a "compressor" can be any numcodecs codec. Only a single compressor may + be provided for Zarr format 2. + If no ``compressor`` is provided, a default compressor will be used. + in [`zarr.config`][zarr.config]. + Use ``None`` to omit the default compressor. + compressor : Codec, optional + Deprecated in favor of ``compressors``. + serializer : dict[str, JSON] | ArrayBytesCodec, optional + Array-to-bytes codec to use for encoding the array data. + Zarr format 3 only. Zarr format 2 arrays use implicit array-to-bytes conversion. + If no ``serializer`` is provided, a default serializer will be used. + These defaults can be changed by modifying the value of ``array.v3_default_serializer`` + in [`zarr.config`][zarr.config]. + fill_value : Any, optional + Fill value for the array. + order : {"C", "F"}, optional + The memory of the array (default is "C"). + For Zarr format 2, this parameter sets the memory order of the array. + For Zarr format 3, this parameter is deprecated, because memory order + is a runtime parameter for Zarr format 3 arrays. The recommended way to specify the memory + order for Zarr format 3 arrays is via the ``config`` parameter, e.g. ``{'config': 'C'}``. + If no ``order`` is provided, a default order will be used. + This default can be changed by modifying the value of ``array.order`` in [`zarr.config`][zarr.config]. + attributes : dict, optional + Attributes for the array. + chunk_key_encoding : ChunkKeyEncoding, optional + A specification of how the chunk keys are represented in storage. + For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``. + For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``. + dimension_names : Iterable[str], optional + The names of the dimensions (default is None). + Zarr format 3 only. Zarr format 2 arrays should not use this parameter. + storage_options : dict, optional + If using an fsspec URL to create the store, these will be passed to the backend implementation. + Ignored otherwise. + overwrite : bool, default False + Whether to overwrite an array with the same name in the store, if one exists. + config : ArrayConfig or ArrayConfigLike, optional + Runtime configuration for the array. + write_data : bool + If a pre-existing array-like object was provided to this function via the ``data`` parameter + then ``write_data`` determines whether the values in that array-like object should be + written to the Zarr array created by this function. If ``write_data`` is ``False``, then the + array will be left empty. + + Returns + ------- + AsyncArray + + """ + compressors = _parse_deprecated_compressor( + compressor, compressors, zarr_format=self.metadata.zarr_format + ) + return await create_array( + store=self.store_path, + name=name, + shape=shape, + dtype=dtype, + data=data, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=self.metadata.zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + storage_options=storage_options, + overwrite=overwrite, + config=config, + write_data=write_data, + ) + + @deprecated("Use AsyncGroup.create_array instead.", category=ZarrDeprecationWarning) + async def create_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyAsyncArray: + """Create an array. + + !!! warning "Deprecated" + `AsyncGroup.create_dataset()` is deprecated since v3.0.0 and will be removed in v3.1.0. + Use `AsyncGroup.create_array` instead. + + Arrays are known as "datasets" in HDF5 terminology. For compatibility + with h5py, Zarr groups also implement the [zarr.AsyncGroup.require_dataset][] method. + + Parameters + ---------- + name : str + Array name. + **kwargs : dict + Additional arguments passed to [zarr.AsyncGroup.create_array][]. + + Returns + ------- + a : AsyncArray + """ + data = kwargs.pop("data", None) + # create_dataset in zarr 2.x requires shape but not dtype if data is + # provided. Allow this configuration by inferring dtype from data if + # necessary and passing it to create_array + if "dtype" not in kwargs and data is not None: + kwargs["dtype"] = data.dtype + array = await self.create_array(name, shape=shape, **kwargs) + if data is not None: + await array.setitem(slice(None), data) + return array + + @deprecated("Use AsyncGroup.require_array instead.", category=ZarrDeprecationWarning) + async def require_dataset( + self, + name: str, + *, + shape: tuple[int, ...], + dtype: npt.DTypeLike = None, + exact: bool = False, + **kwargs: Any, + ) -> AnyAsyncArray: + """Obtain an array, creating if it doesn't exist. + + !!! warning "Deprecated" + `AsyncGroup.require_dataset()` is deprecated since v3.0.0 and will be removed in v3.1.0. + Use `AsyncGroup.require_dataset` instead. + + Arrays are known as "datasets" in HDF5 terminology. For compatibility + with h5py, Zarr groups also implement the [zarr.AsyncGroup.create_dataset][] method. + + Other `kwargs` are as per [zarr.AsyncGroup.create_dataset][]. + + Parameters + ---------- + name : str + Array name. + shape : int or tuple of ints + Array shape. + dtype : str or dtype, optional + NumPy dtype. + exact : bool, optional + If True, require `dtype` to match exactly. If false, require + `dtype` can be cast from array dtype. + + Returns + ------- + a : AsyncArray + """ + return await self.require_array(name, shape=shape, dtype=dtype, exact=exact, **kwargs) + + async def require_array( + self, + name: str, + *, + shape: ShapeLike, + dtype: npt.DTypeLike = None, + exact: bool = False, + **kwargs: Any, + ) -> AnyAsyncArray: + """Obtain an array, creating if it doesn't exist. + + Other `kwargs` are as per [zarr.AsyncGroup.create_dataset][]. + + Parameters + ---------- + name : str + Array name. + shape : int or tuple of ints + Array shape. + dtype : str or dtype, optional + NumPy dtype. + exact : bool, optional + If True, require `dtype` to match exactly. If false, require + `dtype` can be cast from array dtype. + + Returns + ------- + a : AsyncArray + """ + try: + ds = await self.getitem(name) + if not isinstance(ds, AsyncArray): + raise TypeError(f"Incompatible object ({ds.__class__.__name__}) already exists") + + shape = parse_shapelike(shape) + if shape != ds.shape: + raise TypeError(f"Incompatible shape ({ds.shape} vs {shape})") + + dtype = np.dtype(dtype) + if exact: + if ds.dtype != dtype: + raise TypeError(f"Incompatible dtype ({ds.dtype} vs {dtype})") + else: + if not np.can_cast(ds.dtype, dtype): + raise TypeError(f"Incompatible dtype ({ds.dtype} vs {dtype})") + except KeyError: + ds = await self.create_array(name, shape=shape, dtype=dtype, **kwargs) + + return ds + + async def update_attributes(self, new_attributes: dict[str, Any]) -> AsyncGroup: + """Update group attributes. + + Parameters + ---------- + new_attributes : dict + New attributes to set on the group. + + Returns + ------- + self : AsyncGroup + """ + self.metadata.attributes.update(new_attributes) + + # Write new metadata + await self._save_metadata() + + return self + + def __repr__(self) -> str: + return f"" + + async def nmembers( + self, + max_depth: int | None = 0, + ) -> int: + """Count the number of members in this group. + + Parameters + ---------- + max_depth : int, default 0 + The maximum number of levels of the hierarchy to include. By + default, (``max_depth=0``) only immediate children are included. Set + ``max_depth=None`` to include all nodes, and some positive integer + to consider children within that many levels of the root Group. + + Returns + ------- + count : int + """ + # check if we can use consolidated metadata, which requires that we have non-None + # consolidated metadata at all points in the hierarchy. + if self.metadata.consolidated_metadata is not None: + if max_depth is not None and max_depth < 0: + raise ValueError(f"max_depth must be None or >= 0. Got '{max_depth}' instead") + if max_depth is None: + return len(self.metadata.consolidated_metadata.flattened_metadata) + else: + return len( + [ + x + for x in self.metadata.consolidated_metadata.flattened_metadata + if x.count("/") <= max_depth + ] + ) + # TODO: consider using aioitertools.builtins.sum for this + # return await aioitertools.builtins.sum((1 async for _ in self.members()), start=0) + n = 0 + async for _ in self.members(max_depth=max_depth): + n += 1 + return n + + async def members( + self, + max_depth: int | None = 0, + *, + use_consolidated_for_children: bool = True, + ) -> AsyncGenerator[ + tuple[str, AnyAsyncArray | AsyncGroup], + None, + ]: + """ + Returns an AsyncGenerator over the arrays and groups contained in this group. + This method requires that `store_path.store` supports directory listing. + + The results are not guaranteed to be ordered. + + Parameters + ---------- + max_depth : int, default 0 + The maximum number of levels of the hierarchy to include. By + default, (``max_depth=0``) only immediate children are included. Set + ``max_depth=None`` to include all nodes, and some positive integer + to consider children within that many levels of the root Group. + use_consolidated_for_children : bool, default True + Whether to use the consolidated metadata of child groups loaded + from the store. Note that this only affects groups loaded from the + store. If the current Group already has consolidated metadata, it + will always be used. + + Returns + ------- + path: + A string giving the path to the target, relative to the Group ``self``. + value: AsyncArray or AsyncGroup + The AsyncArray or AsyncGroup that is a child of ``self``. + """ + if max_depth is not None and max_depth < 0: + raise ValueError(f"max_depth must be None or >= 0. Got '{max_depth}' instead") + async for item in self._members( + max_depth=max_depth, use_consolidated_for_children=use_consolidated_for_children + ): + yield item + + def _members_consolidated( + self, max_depth: int | None, prefix: str = "" + ) -> Generator[ + tuple[str, AnyAsyncArray | AsyncGroup], + None, + ]: + consolidated_metadata = self.metadata.consolidated_metadata + + do_recursion = max_depth is None or max_depth > 0 + + # we kind of just want the top-level keys. + if consolidated_metadata is not None: + for key in consolidated_metadata.metadata: + obj = self._getitem_consolidated( + self.store_path, key, prefix=self.name + ) # Metadata -> Group/Array + key = f"{prefix}/{key}".lstrip("/") + yield key, obj + + if do_recursion and isinstance(obj, AsyncGroup): + if max_depth is None: + new_depth = None + else: + new_depth = max_depth - 1 + yield from obj._members_consolidated(new_depth, prefix=key) + + async def _members( + self, max_depth: int | None, *, use_consolidated_for_children: bool = True + ) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: + skip_keys: tuple[str, ...] + if self.metadata.zarr_format == 2: + skip_keys = (".zattrs", ".zgroup", ".zarray", ".zmetadata") + elif self.metadata.zarr_format == 3: + skip_keys = ("zarr.json",) + else: + raise ValueError(f"Unknown Zarr format: {self.metadata.zarr_format}") + + if self.metadata.consolidated_metadata is not None: + members = self._members_consolidated(max_depth=max_depth) + for member in members: + yield member + return + + if not self.store_path.store.supports_listing: + msg = ( + f"The store associated with this group ({type(self.store_path.store)}) " + "does not support listing, " + "specifically via the `list_dir` method. " + "This function requires a store that supports listing." + ) + + raise ValueError(msg) + # enforce a concurrency limit by passing a semaphore to all the recursive functions + semaphore = asyncio.Semaphore(config.get("async.concurrency")) + async for member in _iter_members_deep( + self, + max_depth=max_depth, + skip_keys=skip_keys, + semaphore=semaphore, + use_consolidated_for_children=use_consolidated_for_children, + ): + yield member + + async def create_hierarchy( + self, + nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], + *, + overwrite: bool = False, + ) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: + """ + Create a hierarchy of arrays or groups rooted at this group. + + This function will parse its input to ensure that the hierarchy is complete. Any implicit groups + will be inserted as needed. For example, an input like + ```{'a/b': GroupMetadata}``` will be parsed to + ```{'': GroupMetadata, 'a': GroupMetadata, 'b': Groupmetadata}```. + + Explicitly specifying a root group, e.g. with ``nodes = {'': GroupMetadata()}`` is an error + because this group instance is the root group. + + After input parsing, this function then creates all the nodes in the hierarchy concurrently. + + Arrays and Groups are yielded in the order they are created. This order is not stable and + should not be relied on. + + Parameters + ---------- + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes in the hierarchy, + relative to the path of the group. The values are instances of ``GroupMetadata`` or ``ArrayMetadata``. Note that + all values must have the same ``zarr_format`` as the parent group -- it is an error to mix zarr versions in the + same hierarchy. + + Leading "/" characters from keys will be removed. + overwrite : bool + Whether to overwrite existing nodes. Defaults to ``False``, in which case an error is + raised instead of overwriting an existing array or group. + + This function will not erase an existing group unless that group is explicitly named in + ``nodes``. If ``nodes`` defines implicit groups, e.g. ``{`'a/b/c': GroupMetadata}``, and a + group already exists at path ``a``, then this function will leave the group at ``a`` as-is. + + Yields + ------ + tuple[str, AsyncArray | AsyncGroup]. + """ + # check that all the nodes have the same zarr_format as Self + prefix = self.path + nodes_parsed = {} + for key, value in nodes.items(): + if value.zarr_format != self.metadata.zarr_format: + msg = ( + "The zarr_format of the nodes must be the same as the parent group. " + f"The node at {key} has zarr_format {value.zarr_format}, but the parent group" + f" has zarr_format {self.metadata.zarr_format}." + ) + raise ValueError(msg) + if normalize_path(key) == "": + msg = ( + "The input defines a root node, but a root node already exists, namely this Group instance." + "It is an error to use this method to create a root node. " + "Remove the root node from the input dict, or use a function like " + "create_rooted_hierarchy to create a rooted hierarchy." + ) + raise ValueError(msg) + else: + nodes_parsed[_join_paths([prefix, key])] = value + + async for key, node in create_hierarchy( + store=self.store, + nodes=nodes_parsed, + overwrite=overwrite, + ): + if prefix == "": + out_key = key + else: + out_key = key.removeprefix(prefix + "/") + yield out_key, node + + async def keys(self) -> AsyncGenerator[str, None]: + """Iterate over member names.""" + async for key, _ in self.members(): + yield key + + async def contains(self, member: str) -> bool: + """Check if a member exists in the group. + + Parameters + ---------- + member : str + Member name. + + Returns + ------- + bool + """ + # TODO: this can be made more efficient. + try: + await self.getitem(member) + except KeyError: + return False + else: + return True + + async def groups(self) -> AsyncGenerator[tuple[str, AsyncGroup], None]: + """Iterate over subgroups.""" + async for name, value in self.members(): + if isinstance(value, AsyncGroup): + yield name, value + + async def group_keys(self) -> AsyncGenerator[str, None]: + """Iterate over group names.""" + async for key, _ in self.groups(): + yield key + + async def group_values(self) -> AsyncGenerator[AsyncGroup, None]: + """Iterate over group values.""" + async for _, group in self.groups(): + yield group + + async def arrays( + self, + ) -> AsyncGenerator[tuple[str, AnyAsyncArray], None]: + """Iterate over arrays.""" + async for key, value in self.members(): + if isinstance(value, AsyncArray): + yield key, value + + async def array_keys(self) -> AsyncGenerator[str, None]: + """Iterate over array names.""" + async for key, _ in self.arrays(): + yield key + + async def array_values( + self, + ) -> AsyncGenerator[AnyAsyncArray, None]: + """Iterate over array values.""" + async for _, array in self.arrays(): + yield array + + async def tree( + self, + expand: bool | None = None, + level: int | None = None, + *, + max_nodes: int = 500, + plain: bool = False, + ) -> Any: + """ + Return a tree-like representation of a hierarchy. + + Parameters + ---------- + expand : bool, optional + This keyword is not yet supported. A NotImplementedError is raised if + it's used. + level : int, optional + The maximum depth below this Group to display in the tree. + max_nodes : int + Maximum number of nodes to display before truncating. Default is 500. + plain : bool, optional + If True, return a plain-text tree without ANSI styling. This is + useful when the output will be consumed by an LLM or written to a + file. Default is False. + + Returns + ------- + TreeRepr + A pretty-printable object displaying the hierarchy. + """ + from zarr.core._tree import group_tree_async + + if expand is not None: + raise NotImplementedError("'expand' is not yet implemented.") + return await group_tree_async(self, max_depth=level, max_nodes=max_nodes, plain=plain) + + async def empty(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray: + """Create an empty array with the specified shape in this Group. The contents will + be filled with the array's fill value or zeros if no fill value is provided. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Notes + ----- + The contents of an empty Zarr array are not defined. On attempting to + retrieve data from an empty Zarr array, any values may be returned, + and these are not guaranteed to be stable from one access to the next. + """ + return await async_api.empty(shape=shape, store=self.store_path, path=name, **kwargs) + + async def zeros(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray: + """Create an array, with zero being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.zeros(shape=shape, store=self.store_path, path=name, **kwargs) + + async def ones(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray: + """Create an array, with one being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.ones(shape=shape, store=self.store_path, path=name, **kwargs) + + async def full( + self, *, name: str, shape: tuple[int, ...], fill_value: Any | None, **kwargs: Any + ) -> AnyAsyncArray: + """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + fill_value : scalar + Value to fill the array with. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.full( + shape=shape, + fill_value=fill_value, + store=self.store_path, + path=name, + **kwargs, + ) + + async def empty_like( + self, *, name: str, data: async_api.ArrayLike, **kwargs: Any + ) -> AnyAsyncArray: + """Create an empty sub-array like `data`. The contents will be filled with + the array's fill value or zeros if no fill value is provided. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create an empty array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.empty_like(a=data, store=self.store_path, path=name, **kwargs) + + async def zeros_like( + self, *, name: str, data: async_api.ArrayLike, **kwargs: Any + ) -> AnyAsyncArray: + """Create a sub-array of zeros like `data`. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.zeros_like(a=data, store=self.store_path, path=name, **kwargs) + + async def ones_like( + self, *, name: str, data: async_api.ArrayLike, **kwargs: Any + ) -> AnyAsyncArray: + """Create a sub-array of ones like `data`. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.ones_like(a=data, store=self.store_path, path=name, **kwargs) + + async def full_like( + self, *, name: str, data: async_api.ArrayLike, **kwargs: Any + ) -> AnyAsyncArray: + """Create a sub-array like `data` filled with the `fill_value` of `data` . + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + AsyncArray + The new array. + """ + return await async_api.full_like(a=data, store=self.store_path, path=name, **kwargs) + + async def move(self, source: str, dest: str) -> None: + """Move a sub-group or sub-array from one path to another. + + Notes + ----- + Not implemented + """ + raise NotImplementedError + + +@dataclass(frozen=True) +class Group(SyncMixin): + """ + A Zarr group. + """ + + _async_group: AsyncGroup + + @classmethod + def from_store( + cls, + store: StoreLike, + *, + attributes: dict[str, Any] | None = None, + zarr_format: ZarrFormat = 3, + overwrite: bool = False, + ) -> Group: + """Instantiate a group from an initialized store. + + Parameters + ---------- + store : StoreLike + StoreLike containing the Group. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + attributes : dict, optional + A dictionary of JSON-serializable values with user-defined attributes. + zarr_format : {2, 3}, optional + Zarr storage format version. + overwrite : bool, optional + If True, do not raise an error if the group already exists. + + Returns + ------- + Group + Group instantiated from the store. + + Raises + ------ + ContainsArrayError, ContainsGroupError, ContainsArrayAndGroupError + """ + attributes = attributes or {} + obj = sync( + AsyncGroup.from_store( + store, + attributes=attributes, + overwrite=overwrite, + zarr_format=zarr_format, + ), + ) + + return cls(obj) + + @classmethod + def open( + cls, + store: StoreLike, + zarr_format: ZarrFormat | None = 3, + ) -> Group: + """Open a group from an initialized store. + + Parameters + ---------- + store : StoreLike + Store containing the Group. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + zarr_format : {2, 3, None}, optional + Zarr storage format version. + + Returns + ------- + Group + Group instantiated from the store. + """ + obj = sync(AsyncGroup.open(store, zarr_format=zarr_format)) + return cls(obj) + + def __getitem__(self, path: str) -> AnyArray | Group: + """Obtain a group member. + + Parameters + ---------- + path : str + Group member name. + + Returns + ------- + Array | Group + Group member (Array or Group) at the specified key + + Examples + -------- + ```python + import zarr + from zarr.core.group import Group + group = Group.from_store(zarr.storage.MemoryStore()) + group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="float64") + group.create_group(name="subgroup").create_array(name="subarray", shape=(10,), chunks=(10,), dtype="float64") + group["subarray"] + # + group["subgroup"] + # + group["subgroup"]["subarray"] + # + ``` + + """ + obj = self._sync(self._async_group.getitem(path)) + if isinstance(obj, AsyncArray): + return Array(obj) + else: + return Group(obj) + + def get(self, path: str, default: DefaultT | None = None) -> AnyArray | Group | DefaultT | None: + """Obtain a group member, returning default if not found. + + Parameters + ---------- + path : str + Group member name. + default : object + Default value to return if key is not found (default: None). + + Returns + ------- + object + Group member (Array or Group) or default if not found. + + Examples + -------- + ```python + import zarr + from zarr.core.group import Group + group = Group.from_store(zarr.storage.MemoryStore()) + group.create_array(name="subarray", shape=(10,), chunks=(10,), dtype="float64") + group.create_group(name="subgroup") + group.get("subarray") + # + group.get("subgroup") + # + group.get("nonexistent", None) + # None + ``` + + """ + try: + return self[path] + except KeyError: + return default + + def __delitem__(self, key: str) -> None: + """Delete a group member. + + Parameters + ---------- + key : str + Group member name. + + Examples + -------- + >>> import zarr + >>> group = Group.from_store(zarr.storage.MemoryStore() + >>> group.create_array(name="subarray", shape=(10,), chunks=(10,)) + >>> del group["subarray"] + >>> "subarray" in group + False + """ + self._sync(self._async_group.delitem(key)) + + def __iter__(self) -> Iterator[str]: + """Return an iterator over group member names. + Examples + -------- + >>> import zarr + >>> g1 = zarr.group() + >>> g2 = g1.create_group('foo') + >>> g3 = g1.create_group('bar') + >>> d1 = g1.create_array('baz', shape=(10,), chunks=(10,)) + >>> d2 = g1.create_array('quux', shape=(10,), chunks=(10,)) + >>> for name in g1: + ... print(name) + baz + bar + foo + quux + """ + yield from self.keys() + + def __len__(self) -> int: + """Number of members.""" + return self.nmembers() + + def __setitem__(self, key: str, value: Any) -> None: + """Fastpath for creating a new array. + + New arrays will be created using default settings for the array type. + If you need to create an array with custom settings, use the `create_array` method. + + Parameters + ---------- + key : str + Array name. + value : Any + Array data. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group["foo"] = zarr.zeros((10,)) + >>> group["foo"] + + """ + self._sync(self._async_group.setitem(key, value)) + + def __repr__(self) -> str: + return f"" + + async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group: + """Update the attributes of this group. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> await group.update_attributes_async({"foo": "bar"}) + >>> group.attrs.asdict() + {'foo': 'bar'} + """ + new_metadata = replace(self.metadata, attributes=new_attributes) + + # Write new metadata + to_save = new_metadata.to_buffer_dict(default_buffer_prototype()) + awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()] + await asyncio.gather(*awaitables) + + async_group = replace(self._async_group, metadata=new_metadata) + return replace(self, _async_group=async_group) + + @property + def store_path(self) -> StorePath: + """Path-like interface for the Store.""" + return self._async_group.store_path + + @property + def metadata(self) -> GroupMetadata: + """Group metadata.""" + return self._async_group.metadata + + @property + def path(self) -> str: + """Storage path.""" + return self._async_group.path + + @property + def name(self) -> str: + """Group name following h5py convention.""" + return self._async_group.name + + @property + def basename(self) -> str: + """Final component of name.""" + return self._async_group.basename + + @property + def attrs(self) -> Attributes: + """Attributes of this Group""" + return Attributes(self) + + @property + def info(self) -> Any: + """ + Return the statically known information for a group. + + Returns + ------- + GroupInfo + + Related + ------- + [zarr.Group.info_complete][] + All information about a group, including dynamic information + like the children members. + """ + return self._async_group.info + + def info_complete(self) -> Any: + """ + Return information for a group. + + If this group doesn't contain consolidated metadata then + this will need to read from the backing Store. + + Returns + ------- + GroupInfo + + Related + ------- + [zarr.Group.info][] + """ + return self._sync(self._async_group.info_complete()) + + @property + def store(self) -> Store: + # Backwards compatibility for 2.x + return self._async_group.store + + @property + def read_only(self) -> bool: + # Backwards compatibility for 2.x + return self._async_group.read_only + + @property + def synchronizer(self) -> None: + # Backwards compatibility for 2.x + # Not implemented in 3.x yet. + return self._async_group.synchronizer + + def update_attributes(self, new_attributes: dict[str, Any]) -> Group: + """Update the attributes of this group. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.update_attributes({"foo": "bar"}) + >>> group.attrs.asdict() + {'foo': 'bar'} + """ + self._sync(self._async_group.update_attributes(new_attributes)) + return self + + def nmembers(self, max_depth: int | None = 0) -> int: + """Count the number of members in this group. + + Parameters + ---------- + max_depth : int, default 0 + The maximum number of levels of the hierarchy to include. By + default, (``max_depth=0``) only immediate children are included. Set + ``max_depth=None`` to include all nodes, and some positive integer + to consider children within that many levels of the root Group. + + Returns + ------- + count : int + """ + + return self._sync(self._async_group.nmembers(max_depth=max_depth)) + + def members( + self, max_depth: int | None = 0, *, use_consolidated_for_children: bool = True + ) -> tuple[tuple[str, AnyArray | Group], ...]: + """ + Returns an AsyncGenerator over the arrays and groups contained in this group. + This method requires that `store_path.store` supports directory listing. + + The results are not guaranteed to be ordered. + + Parameters + ---------- + max_depth : int, default 0 + The maximum number of levels of the hierarchy to include. By + default, (``max_depth=0``) only immediate children are included. Set + ``max_depth=None`` to include all nodes, and some positive integer + to consider children within that many levels of the root Group. + use_consolidated_for_children : bool, default True + Whether to use the consolidated metadata of child groups loaded + from the store. Note that this only affects groups loaded from the + store. If the current Group already has consolidated metadata, it + will always be used. + + Returns + ------- + path: + A string giving the path to the target, relative to the Group ``self``. + value: AsyncArray or AsyncGroup + The AsyncArray or AsyncGroup that is a child of ``self``. + """ + _members = self._sync_iter(self._async_group.members(max_depth=max_depth)) + + return tuple((kv[0], _parse_async_node(kv[1])) for kv in _members) + + def create_hierarchy( + self, + nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], + *, + overwrite: bool = False, + ) -> Iterator[tuple[str, Group | AnyArray]]: + """ + Create a hierarchy of arrays or groups rooted at this group. + + This function will parse its input to ensure that the hierarchy is complete. Any implicit groups + will be inserted as needed. For example, an input like + ```{'a/b': GroupMetadata}``` will be parsed to + ```{'': GroupMetadata, 'a': GroupMetadata, 'b': Groupmetadata}```. + + Explicitly specifying a root group, e.g. with ``nodes = {'': GroupMetadata()}`` is an error + because this group instance is the root group. + + After input parsing, this function then creates all the nodes in the hierarchy concurrently. + + Arrays and Groups are yielded in the order they are created. This order is not stable and + should not be relied on. + + Parameters + ---------- + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes in the hierarchy, + relative to the path of the group. The values are instances of ``GroupMetadata`` or ``ArrayMetadata``. Note that + all values must have the same ``zarr_format`` as the parent group -- it is an error to mix zarr versions in the + same hierarchy. + + Leading "/" characters from keys will be removed. + overwrite : bool + Whether to overwrite existing nodes. Defaults to ``False``, in which case an error is + raised instead of overwriting an existing array or group. + + This function will not erase an existing group unless that group is explicitly named in + ``nodes``. If ``nodes`` defines implicit groups, e.g. ``{`'a/b/c': GroupMetadata}``, and a + group already exists at path ``a``, then this function will leave the group at ``a`` as-is. + + Yields + ------ + tuple[str, Array | Group]. + + Examples + -------- + >>> import zarr + >>> from zarr.core.group import GroupMetadata + >>> root = zarr.create_group(store={}) + >>> for key, val in root.create_hierarchy({'a/b/c': GroupMetadata()}): + ... print(key, val) + ... + + + + """ + for key, node in self._sync_iter( + self._async_group.create_hierarchy(nodes, overwrite=overwrite) + ): + yield (key, _parse_async_node(node)) + + def keys(self) -> Generator[str, None]: + """Return an iterator over group member names. + + Examples + -------- + >>> import zarr + >>> g1 = zarr.group() + >>> g2 = g1.create_group('foo') + >>> g3 = g1.create_group('bar') + >>> d1 = g1.create_array('baz', shape=(10,), chunks=(10,)) + >>> d2 = g1.create_array('quux', shape=(10,), chunks=(10,)) + >>> for name in g1.keys(): + ... print(name) + baz + bar + foo + quux + """ + yield from self._sync_iter(self._async_group.keys()) + + def __contains__(self, member: str) -> bool: + """Test for group membership. + + Examples + -------- + >>> import zarr + >>> g1 = zarr.group() + >>> g2 = g1.create_group('foo') + >>> d1 = g1.create_array('bar', shape=(10,), chunks=(10,)) + >>> 'foo' in g1 + True + >>> 'bar' in g1 + True + >>> 'baz' in g1 + False + + """ + return self._sync(self._async_group.contains(member)) + + def groups(self) -> Generator[tuple[str, Group], None]: + """Return the sub-groups of this group as a generator of (name, group) pairs. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_group("subgroup") + >>> for name, subgroup in group.groups(): + ... print(name, subgroup) + subgroup + """ + for name, async_group in self._sync_iter(self._async_group.groups()): + yield name, Group(async_group) + + def group_keys(self) -> Generator[str, None]: + """Return an iterator over group member names. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_group("subgroup") + >>> for name in group.group_keys(): + ... print(name) + subgroup + """ + for name, _ in self.groups(): + yield name + + def group_values(self) -> Generator[Group, None]: + """Return an iterator over group members. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_group("subgroup") + >>> for subgroup in group.group_values(): + ... print(subgroup) + + """ + for _, group in self.groups(): + yield group + + def arrays(self) -> Generator[tuple[str, AnyArray], None]: + """Return the sub-arrays of this group as a generator of (name, array) pairs + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_array("subarray", shape=(10,), chunks=(10,)) + >>> for name, subarray in group.arrays(): + ... print(name, subarray) + subarray + """ + for name, async_array in self._sync_iter(self._async_group.arrays()): + yield name, Array(async_array) + + def array_keys(self) -> Generator[str, None]: + """Return an iterator over group member names. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_array("subarray", shape=(10,), chunks=(10,)) + >>> for name in group.array_keys(): + ... print(name) + subarray + """ + + for name, _ in self.arrays(): + yield name + + def array_values(self) -> Generator[AnyArray, None]: + """Return an iterator over group members. + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> group.create_array("subarray", shape=(10,), chunks=(10,)) + >>> for subarray in group.array_values(): + ... print(subarray) + + """ + for _, array in self.arrays(): + yield array + + def tree( + self, + expand: bool | None = None, + level: int | None = None, + *, + max_nodes: int = 500, + plain: bool = False, + ) -> Any: + """ + Return a tree-like representation of a hierarchy. + + Parameters + ---------- + expand : bool, optional + This keyword is not yet supported. A NotImplementedError is raised if + it's used. + level : int, optional + The maximum depth below this Group to display in the tree. + max_nodes : int + Maximum number of nodes to display before truncating. Default is 500. + plain : bool, optional + If True, return a plain-text tree without ANSI styling. This is + useful when the output will be consumed by an LLM or written to a + file. Default is False. + + Returns + ------- + TreeRepr + A pretty-printable object displaying the hierarchy. + """ + return self._sync( + self._async_group.tree(expand=expand, level=level, max_nodes=max_nodes, plain=plain) + ) + + def create_group(self, name: str, **kwargs: Any) -> Group: + """Create a sub-group. + + Parameters + ---------- + name : str + Name of the new subgroup. + + Returns + ------- + Group + + Examples + -------- + >>> import zarr + >>> group = zarr.group() + >>> subgroup = group.create_group("subgroup") + >>> subgroup + + """ + return Group(self._sync(self._async_group.create_group(name, **kwargs))) + + def require_group(self, name: str, **kwargs: Any) -> Group: + """Obtain a sub-group, creating one if it doesn't exist. + + Parameters + ---------- + name : str + Group name. + + Returns + ------- + g : Group + """ + return Group(self._sync(self._async_group.require_group(name, **kwargs))) + + def require_groups(self, *names: str) -> tuple[Group, ...]: + """Convenience method to require multiple groups in a single call. + + Parameters + ---------- + *names : str + Group names. + + Returns + ------- + groups : tuple of Groups + """ + return tuple(map(Group, self._sync(self._async_group.require_groups(*names)))) + + def create( + self, + name: str, + *, + shape: ShapeLike | None = None, + dtype: ZDTypeLike | None = None, + data: np.ndarray[Any, np.dtype[Any]] | None = None, + chunks: tuple[int, ...] | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + compressor: CompressorLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNames = None, + storage_options: dict[str, Any] | None = None, + overwrite: bool = False, + config: ArrayConfigLike | None = None, + write_data: bool = True, + ) -> AnyArray: + """Create an array within this group. + + This method lightly wraps [`zarr.core.array.create_array`][]. + + Parameters + ---------- + name : str + The name of the array relative to the group. If ``path`` is ``None``, the array will be located + at the root of the store. + shape : ShapeLike, optional + Shape of the array. Must be ``None`` if ``data`` is provided. + dtype : npt.DTypeLike | None + Data type of the array. Must be ``None`` if ``data`` is provided. + data : Array-like data to use for initializing the array. If this parameter is provided, the + ``shape`` and ``dtype`` parameters must be ``None``. + chunks : tuple[int, ...], optional + Chunk shape of the array. + If not specified, default are guessed based on the shape and dtype. + shards : tuple[int, ...], optional + Shard shape of the array. The default value of ``None`` results in no sharding at all. + filters : Iterable[Codec] | Literal["auto"], optional + Iterable of filters to apply to each chunk of the array, in order, before serializing that + chunk to bytes. + + For Zarr format 3, a "filter" is a codec that takes an array and returns an array, + and these values must be instances of [`zarr.abc.codec.ArrayArrayCodec`][], or a + dict representations of [`zarr.abc.codec.ArrayArrayCodec`][]. + + For Zarr format 2, a "filter" can be any numcodecs codec; you should ensure that the + the order if your filters is consistent with the behavior of each filter. + + The default value of ``"auto"`` instructs Zarr to use a default used based on the data + type of the array and the Zarr format specified. For all data types in Zarr V3, and most + data types in Zarr V2, the default filters are empty. The only cases where default filters + are not empty is when the Zarr format is 2, and the data type is a variable-length data type like + [`zarr.dtype.VariableLengthUTF8`][] or [`zarr.dtype.VariableLengthUTF8`][]. In these cases, + the default filters contains a single element which is a codec specific to that particular data type. + + To create an array with no filters, provide an empty iterable or the value ``None``. + compressors : Iterable[Codec], optional + List of compressors to apply to the array. Compressors are applied in order, and after any + filters are applied (if any are specified) and the data is serialized into bytes. + + For Zarr format 3, a "compressor" is a codec that takes a bytestream, and + returns another bytestream. Multiple compressors my be provided for Zarr format 3. + If no ``compressors`` are provided, a default set of compressors will be used. + These defaults can be changed by modifying the value of ``array.v3_default_compressors`` + in [`zarr.config`][]. + Use ``None`` to omit default compressors. + + For Zarr format 2, a "compressor" can be any numcodecs codec. Only a single compressor may + be provided for Zarr format 2. + If no ``compressor`` is provided, a default compressor will be used. + in [`zarr.config`][]. + Use ``None`` to omit the default compressor. + compressor : Codec, optional + Deprecated in favor of ``compressors``. + serializer : dict[str, JSON] | ArrayBytesCodec, optional + Array-to-bytes codec to use for encoding the array data. + Zarr format 3 only. Zarr format 2 arrays use implicit array-to-bytes conversion. + If no ``serializer`` is provided, a default serializer will be used. + These defaults can be changed by modifying the value of ``array.v3_default_serializer`` + in [`zarr.config`][]. + fill_value : Any, optional + Fill value for the array. + order : {"C", "F"}, optional + The memory of the array (default is "C"). + For Zarr format 2, this parameter sets the memory order of the array. + For Zarr format 3, this parameter is deprecated, because memory order + is a runtime parameter for Zarr format 3 arrays. The recommended way to specify the memory + order for Zarr format 3 arrays is via the ``config`` parameter, e.g. ``{'config': 'C'}``. + If no ``order`` is provided, a default order will be used. + This default can be changed by modifying the value of ``array.order`` in [`zarr.config`][]. + attributes : dict, optional + Attributes for the array. + chunk_key_encoding : ChunkKeyEncoding, optional + A specification of how the chunk keys are represented in storage. + For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``. + For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``. + dimension_names : Iterable[str], optional + The names of the dimensions (default is None). + Zarr format 3 only. Zarr format 2 arrays should not use this parameter. + storage_options : dict, optional + If using an fsspec URL to create the store, these will be passed to the backend implementation. + Ignored otherwise. + overwrite : bool, default False + Whether to overwrite an array with the same name in the store, if one exists. + config : ArrayConfig or ArrayConfigLike, optional + Runtime configuration for the array. + write_data : bool + If a pre-existing array-like object was provided to this function via the ``data`` parameter + then ``write_data`` determines whether the values in that array-like object should be + written to the Zarr array created by this function. If ``write_data`` is ``False``, then the + array will be left empty. + + Returns + ------- + AsyncArray + """ + return self.create_array( + name, + shape=shape, + dtype=dtype, + data=data, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + compressor=compressor, + serializer=serializer, + fill_value=fill_value, + order=order, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + storage_options=storage_options, + overwrite=overwrite, + config=config, + write_data=write_data, + ) + + def create_array( + self, + name: str, + *, + shape: ShapeLike | None = None, + dtype: ZDTypeLike | None = None, + data: np.ndarray[Any, np.dtype[Any]] | None = None, + chunks: tuple[int, ...] | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + compressor: CompressorLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNames = None, + storage_options: dict[str, Any] | None = None, + overwrite: bool = False, + config: ArrayConfigLike | None = None, + write_data: bool = True, + ) -> AnyArray: + """Create an array within this group. + + This method lightly wraps [zarr.core.array.create_array][]. + + Parameters + ---------- + name : str + The name of the array relative to the group. If ``path`` is ``None``, the array will be located + at the root of the store. + shape : ShapeLike, optional + Shape of the array. Must be ``None`` if ``data`` is provided. + dtype : npt.DTypeLike | None + Data type of the array. Must be ``None`` if ``data`` is provided. + data : Array-like data to use for initializing the array. If this parameter is provided, the + ``shape`` and ``dtype`` parameters must be ``None``. + chunks : tuple[int, ...], optional + Chunk shape of the array. + If not specified, default are guessed based on the shape and dtype. + shards : tuple[int, ...], optional + Shard shape of the array. The default value of ``None`` results in no sharding at all. + filters : Iterable[Codec] | Literal["auto"], optional + Iterable of filters to apply to each chunk of the array, in order, before serializing that + chunk to bytes. + + For Zarr format 3, a "filter" is a codec that takes an array and returns an array, + and these values must be instances of [`zarr.abc.codec.ArrayArrayCodec`][], or a + dict representations of [`zarr.abc.codec.ArrayArrayCodec`][]. + + For Zarr format 2, a "filter" can be any numcodecs codec; you should ensure that the + the order if your filters is consistent with the behavior of each filter. + + The default value of ``"auto"`` instructs Zarr to use a default used based on the data + type of the array and the Zarr format specified. For all data types in Zarr V3, and most + data types in Zarr V2, the default filters are empty. The only cases where default filters + are not empty is when the Zarr format is 2, and the data type is a variable-length data type like + [`zarr.dtype.VariableLengthUTF8`][] or [`zarr.dtype.VariableLengthUTF8`][]. In these cases, + the default filters contains a single element which is a codec specific to that particular data type. + + To create an array with no filters, provide an empty iterable or the value ``None``. + compressors : Iterable[Codec], optional + List of compressors to apply to the array. Compressors are applied in order, and after any + filters are applied (if any are specified) and the data is serialized into bytes. + + For Zarr format 3, a "compressor" is a codec that takes a bytestream, and + returns another bytestream. Multiple compressors my be provided for Zarr format 3. + If no ``compressors`` are provided, a default set of compressors will be used. + These defaults can be changed by modifying the value of ``array.v3_default_compressors`` + in [`zarr.config`][zarr.config]. + Use ``None`` to omit default compressors. + + For Zarr format 2, a "compressor" can be any numcodecs codec. Only a single compressor may + be provided for Zarr format 2. + If no ``compressor`` is provided, a default compressor will be used. + in [`zarr.config`][zarr.config]. + Use ``None`` to omit the default compressor. + compressor : Codec, optional + Deprecated in favor of ``compressors``. + serializer : dict[str, JSON] | ArrayBytesCodec, optional + Array-to-bytes codec to use for encoding the array data. + Zarr format 3 only. Zarr format 2 arrays use implicit array-to-bytes conversion. + If no ``serializer`` is provided, a default serializer will be used. + These defaults can be changed by modifying the value of ``array.v3_default_serializer`` + in [`zarr.config`][zarr.config]. + fill_value : Any, optional + Fill value for the array. + order : {"C", "F"}, optional + The memory of the array (default is "C"). + For Zarr format 2, this parameter sets the memory order of the array. + For Zarr format 3, this parameter is deprecated, because memory order + is a runtime parameter for Zarr format 3 arrays. The recommended way to specify the memory + order for Zarr format 3 arrays is via the ``config`` parameter, e.g. ``{'config': 'C'}``. + If no ``order`` is provided, a default order will be used. + This default can be changed by modifying the value of ``array.order`` in [`zarr.config`][zarr.config]. + attributes : dict, optional + Attributes for the array. + chunk_key_encoding : ChunkKeyEncoding, optional + A specification of how the chunk keys are represented in storage. + For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``. + For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``. + dimension_names : Iterable[str], optional + The names of the dimensions (default is None). + Zarr format 3 only. Zarr format 2 arrays should not use this parameter. + storage_options : dict, optional + If using an fsspec URL to create the store, these will be passed to the backend implementation. + Ignored otherwise. + overwrite : bool, default False + Whether to overwrite an array with the same name in the store, if one exists. + config : ArrayConfig or ArrayConfigLike, optional + Runtime configuration for the array. + write_data : bool + If a pre-existing array-like object was provided to this function via the ``data`` parameter + then ``write_data`` determines whether the values in that array-like object should be + written to the Zarr array created by this function. If ``write_data`` is ``False``, then the + array will be left empty. + + Returns + ------- + AsyncArray + """ + compressors = _parse_deprecated_compressor( + compressor, compressors, zarr_format=self.metadata.zarr_format + ) + return Array( + self._sync( + self._async_group.create_array( + name=name, + shape=shape, + dtype=dtype, + data=data, + chunks=chunks, + shards=shards, + fill_value=fill_value, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + compressors=compressors, + serializer=serializer, + dimension_names=dimension_names, + order=order, + filters=filters, + overwrite=overwrite, + storage_options=storage_options, + config=config, + write_data=write_data, + ) + ) + ) + + @deprecated("Use Group.create_array instead.", category=ZarrDeprecationWarning) + def create_dataset(self, name: str, **kwargs: Any) -> AnyArray: + """Create an array. + + !!! warning "Deprecated" + `Group.create_dataset()` is deprecated since v3.0.0 and will be removed in v3.1.0. + Use `Group.create_array` instead. + + + Arrays are known as "datasets" in HDF5 terminology. For compatibility + with h5py, Zarr groups also implement the [zarr.Group.require_dataset][] method. + + Parameters + ---------- + name : str + Array name. + **kwargs : dict + Additional arguments passed to [zarr.Group.create_array][] + + Returns + ------- + a : Array + """ + return Array(self._sync(self._async_group.create_dataset(name, **kwargs))) + + @deprecated("Use Group.require_array instead.", category=ZarrDeprecationWarning) + def require_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray: + """Obtain an array, creating if it doesn't exist. + + !!! warning "Deprecated" + `Group.require_dataset()` is deprecated since v3.0.0 and will be removed in v3.1.0. + Use `Group.require_array` instead. + + Arrays are known as "datasets" in HDF5 terminology. For compatibility + with h5py, Zarr groups also implement the [zarr.Group.create_dataset][] method. + + Other `kwargs` are as per [zarr.Group.create_dataset][]. + + Parameters + ---------- + name : str + Array name. + **kwargs : + See [zarr.Group.create_dataset][]. + + Returns + ------- + a : Array + """ + return Array(self._sync(self._async_group.require_array(name, shape=shape, **kwargs))) + + def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray: + """Obtain an array, creating if it doesn't exist. + + Other `kwargs` are as per [zarr.Group.create_array][]. + + Parameters + ---------- + name : str + Array name. + **kwargs : + See [zarr.Group.create_array][]. + + Returns + ------- + a : Array + """ + return Array(self._sync(self._async_group.require_array(name, shape=shape, **kwargs))) + + def empty(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyArray: + """Create an empty array with the specified shape in this Group. The contents will be filled with + the array's fill value or zeros if no fill value is provided. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Notes + ----- + The contents of an empty Zarr array are not defined. On attempting to + retrieve data from an empty Zarr array, any values may be returned, + and these are not guaranteed to be stable from one access to the next. + """ + return Array(self._sync(self._async_group.empty(name=name, shape=shape, **kwargs))) + + def zeros(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyArray: + """Create an array, with zero being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + return Array(self._sync(self._async_group.zeros(name=name, shape=shape, **kwargs))) + + def ones(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyArray: + """Create an array, with one being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + return Array(self._sync(self._async_group.ones(name=name, shape=shape, **kwargs))) + + def full( + self, *, name: str, shape: tuple[int, ...], fill_value: Any | None, **kwargs: Any + ) -> AnyArray: + """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. + + Parameters + ---------- + name : str + Name of the array. + shape : int or tuple of int + Shape of the empty array. + fill_value : scalar + Value to fill the array with. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + return Array( + self._sync( + self._async_group.full(name=name, shape=shape, fill_value=fill_value, **kwargs) + ) + ) + + def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: + """Create an empty sub-array like `data`. The contents will be filled + with the array's fill value or zeros if no fill value is provided. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create an empty array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + + Notes + ----- + The contents of an empty Zarr array are not defined. On attempting to + retrieve data from an empty Zarr array, any values may be returned, + and these are not guaranteed to be stable from one access to the next. + """ + return Array(self._sync(self._async_group.empty_like(name=name, data=data, **kwargs))) + + def zeros_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: + """Create a sub-array of zeros like `data`. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + + return Array(self._sync(self._async_group.zeros_like(name=name, data=data, **kwargs))) + + def ones_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: + """Create a sub-array of ones like `data`. + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + return Array(self._sync(self._async_group.ones_like(name=name, data=data, **kwargs))) + + def full_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: + """Create a sub-array like `data` filled with the `fill_value` of `data` . + + Parameters + ---------- + name : str + Name of the array. + data : array-like + The array to create the new array like. + **kwargs + Keyword arguments passed to [zarr.api.asynchronous.create][]. + + Returns + ------- + Array + The new array. + """ + return Array(self._sync(self._async_group.full_like(name=name, data=data, **kwargs))) + + def move(self, source: str, dest: str) -> None: + """Move a sub-group or sub-array from one path to another. + + Notes + ----- + Not implemented + """ + return self._sync(self._async_group.move(source, dest)) + + @deprecated("Use Group.create_array instead.", category=ZarrDeprecationWarning) + def array( + self, + name: str, + *, + shape: ShapeLike, + dtype: npt.DTypeLike, + chunks: tuple[int, ...] | Literal["auto"] = "auto", + shards: tuple[int, ...] | Literal["auto"] | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + compressor: CompressorLike = None, + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNames = None, + storage_options: dict[str, Any] | None = None, + overwrite: bool = False, + config: ArrayConfigLike | None = None, + data: npt.ArrayLike | None = None, + ) -> AnyArray: + """Create an array within this group. + + !!! warning "Deprecated" + `Group.array()` is deprecated since v3.0.0 and will be removed in a future release. + Use `Group.create_array` instead. + + This method lightly wraps [zarr.core.array.create_array][]. + + Parameters + ---------- + name : str + The name of the array relative to the group. If ``path`` is ``None``, the array will be located + at the root of the store. + shape : tuple[int, ...] + Shape of the array. + dtype : npt.DTypeLike + Data type of the array. + chunks : tuple[int, ...], optional + Chunk shape of the array. + If not specified, default are guessed based on the shape and dtype. + shards : tuple[int, ...], optional + Shard shape of the array. The default value of ``None`` results in no sharding at all. + filters : Iterable[Codec] | Literal["auto"], optional + Iterable of filters to apply to each chunk of the array, in order, before serializing that + chunk to bytes. + + For Zarr format 3, a "filter" is a codec that takes an array and returns an array, + and these values must be instances of [`zarr.abc.codec.ArrayArrayCodec`][], or a + dict representations of [`zarr.abc.codec.ArrayArrayCodec`][]. + + For Zarr format 2, a "filter" can be any numcodecs codec; you should ensure that the + the order if your filters is consistent with the behavior of each filter. + + The default value of ``"auto"`` instructs Zarr to use a default used based on the data + type of the array and the Zarr format specified. For all data types in Zarr V3, and most + data types in Zarr V2, the default filters are empty. The only cases where default filters + are not empty is when the Zarr format is 2, and the data type is a variable-length data type like + [`zarr.dtype.VariableLengthUTF8`][] or [`zarr.dtype.VariableLengthUTF8`][]. In these cases, + the default filters contains a single element which is a codec specific to that particular data type. + + To create an array with no filters, provide an empty iterable or the value ``None``. + compressors : Iterable[Codec], optional + List of compressors to apply to the array. Compressors are applied in order, and after any + filters are applied (if any are specified) and the data is serialized into bytes. + + For Zarr format 3, a "compressor" is a codec that takes a bytestream, and + returns another bytestream. Multiple compressors my be provided for Zarr format 3. + If no ``compressors`` are provided, a default set of compressors will be used. + These defaults can be changed by modifying the value of ``array.v3_default_compressors`` + in [`zarr.config`][zarr.config]. + Use ``None`` to omit default compressors. + + For Zarr format 2, a "compressor" can be any numcodecs codec. Only a single compressor may + be provided for Zarr format 2. + If no ``compressor`` is provided, a default compressor will be used. + in [`zarr.config`][zarr.config]. + Use ``None`` to omit the default compressor. + compressor : Codec, optional + Deprecated in favor of ``compressors``. + serializer : dict[str, JSON] | ArrayBytesCodec, optional + Array-to-bytes codec to use for encoding the array data. + Zarr format 3 only. Zarr format 2 arrays use implicit array-to-bytes conversion. + If no ``serializer`` is provided, a default serializer will be used. + These defaults can be changed by modifying the value of ``array.v3_default_serializer`` + in [`zarr.config`][zarr.config]. + fill_value : Any, optional + Fill value for the array. + order : {"C", "F"}, optional + The memory of the array (default is "C"). + For Zarr format 2, this parameter sets the memory order of the array. + For Zarr format 3, this parameter is deprecated, because memory order + is a runtime parameter for Zarr format 3 arrays. The recommended way to specify the memory + order for Zarr format 3 arrays is via the ``config`` parameter, e.g. ``{'config': 'C'}``. + If no ``order`` is provided, a default order will be used. + This default can be changed by modifying the value of ``array.order`` in [`zarr.config`][zarr.config]. + attributes : dict, optional + Attributes for the array. + chunk_key_encoding : ChunkKeyEncoding, optional + A specification of how the chunk keys are represented in storage. + For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``. + For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``. + dimension_names : Iterable[str], optional + The names of the dimensions (default is None). + Zarr format 3 only. Zarr format 2 arrays should not use this parameter. + storage_options : dict, optional + If using an fsspec URL to create the store, these will be passed to the backend implementation. + Ignored otherwise. + overwrite : bool, default False + Whether to overwrite an array with the same name in the store, if one exists. + config : ArrayConfig or ArrayConfigLike, optional + Runtime configuration for the array. + data : array_like + The data to fill the array with. + + Returns + ------- + AsyncArray + """ + compressors = _parse_deprecated_compressor(compressor, compressors) + return Array( + self._sync( + self._async_group.create_dataset( + name=name, + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + fill_value=fill_value, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + compressors=compressors, + serializer=serializer, + dimension_names=dimension_names, + order=order, + filters=filters, + overwrite=overwrite, + storage_options=storage_options, + config=config, + data=data, + ) + ) + ) + + +async def create_hierarchy( + *, + store: Store, + nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], + overwrite: bool = False, +) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: + """ + Create a complete zarr hierarchy from a collection of metadata objects. + + This function will parse its input to ensure that the hierarchy is complete. Any implicit groups + will be inserted as needed. For example, an input like + ```{'a/b': GroupMetadata}``` will be parsed to + ```{'': GroupMetadata, 'a': GroupMetadata, 'b': Groupmetadata}``` + + After input parsing, this function then creates all the nodes in the hierarchy concurrently. + + Arrays and Groups are yielded in the order they are created. This order is not stable and + should not be relied on. + + Parameters + ---------- + store : Store + The storage backend to use. + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes in the hierarchy, + relative to the root of the ``Store``. The root of the store can be specified with the empty + string ``''``. The values are instances of ``GroupMetadata`` or ``ArrayMetadata``. Note that + all values must have the same ``zarr_format`` -- it is an error to mix zarr versions in the + same hierarchy. + + Leading "/" characters from keys will be removed. + overwrite : bool + Whether to overwrite existing nodes. Defaults to ``False``, in which case an error is + raised instead of overwriting an existing array or group. + + This function will not erase an existing group unless that group is explicitly named in + ``nodes``. If ``nodes`` defines implicit groups, e.g. ``{`'a/b/c': GroupMetadata}``, and a + group already exists at path ``a``, then this function will leave the group at ``a`` as-is. + + Yields + ------ + tuple[str, AsyncGroup | AsyncArray] + This function yields (path, node) pairs, in the order the nodes were created. + + Examples + -------- + >>> from zarr.api.asynchronous import create_hierarchy + >>> from zarr.storage import MemoryStore + >>> from zarr.core.group import GroupMetadata + >>> import asyncio + >>> store = MemoryStore() + >>> nodes = {'a': GroupMetadata(attributes={'name': 'leaf'})} + >>> async def run(): + ... print(dict([x async for x in create_hierarchy(store=store, nodes=nodes)])) + >>> asyncio.run(run()) + # {'a': , '': } + """ + # normalize the keys to be valid paths + nodes_normed_keys = _normalize_path_keys(nodes) + + # ensure that all nodes have the same zarr_format, and add implicit groups as needed + nodes_parsed = _parse_hierarchy_dict(data=nodes_normed_keys) + redundant_implicit_groups = [] + + # empty hierarchies should be a no-op + if len(nodes_parsed) > 0: + # figure out which zarr format we are using + zarr_format = next(iter(nodes_parsed.values())).zarr_format + + # check which implicit groups will require materialization + implicit_group_keys = tuple( + filter(lambda k: isinstance(nodes_parsed[k], ImplicitGroupMarker), nodes_parsed) + ) + # read potential group metadata for each implicit group + maybe_extant_group_coros = ( + _read_group_metadata(store, k, zarr_format=zarr_format) for k in implicit_group_keys + ) + maybe_extant_groups = await asyncio.gather( + *maybe_extant_group_coros, return_exceptions=True + ) + + for key, value in zip(implicit_group_keys, maybe_extant_groups, strict=True): + if isinstance(value, BaseException): + if isinstance(value, FileNotFoundError): + # this is fine -- there was no group there, so we will create one + pass + else: + raise value + else: + # a loop exists already at ``key``, so we can avoid creating anything there + redundant_implicit_groups.append(key) + + if overwrite: + # we will remove any nodes that collide with arrays and non-implicit groups defined in + # nodes + + # track the keys of nodes we need to delete + to_delete_keys = [] + to_delete_keys.extend( + [k for k, v in nodes_parsed.items() if k not in implicit_group_keys] + ) + await asyncio.gather(*(store.delete_dir(key) for key in to_delete_keys)) + else: + # This type is long. + coros: ( + Generator[Coroutine[Any, Any, ArrayV2Metadata | GroupMetadata], None, None] + | Generator[Coroutine[Any, Any, ArrayV3Metadata | GroupMetadata], None, None] + ) + if zarr_format == 2: + coros = (_read_metadata_v2(store=store, path=key) for key in nodes_parsed) + elif zarr_format == 3: + coros = (_read_metadata_v3(store=store, path=key) for key in nodes_parsed) + else: # pragma: no cover + raise ValueError(f"Invalid zarr_format: {zarr_format}") # pragma: no cover + + extant_node_query = dict( + zip( + nodes_parsed.keys(), + await asyncio.gather(*coros, return_exceptions=True), + strict=False, + ) + ) + # iterate over the existing arrays / groups and figure out which of them conflict + # with the arrays / groups we want to create + for key, extant_node in extant_node_query.items(): + proposed_node = nodes_parsed[key] + if isinstance(extant_node, BaseException): + if isinstance(extant_node, FileNotFoundError): + # ignore FileNotFoundError, because they represent nodes we can safely create + pass + else: + # Any other exception is a real error + raise extant_node + else: + # this is a node that already exists, but a node with the same key was specified + # in nodes_parsed. + if isinstance(extant_node, GroupMetadata): + # a group already exists where we want to create a group + if isinstance(proposed_node, ImplicitGroupMarker): + # we have proposed an implicit group, which is OK -- we will just skip + # creating this particular metadata document + redundant_implicit_groups.append(key) + else: + # we have proposed an explicit group, which is an error, given that a + # group already exists. + msg = f"A group exists in store {store!r} at path {key!r}." + raise ContainsGroupError(msg) + elif isinstance(extant_node, ArrayV2Metadata | ArrayV3Metadata): + # we are trying to overwrite an existing array. this is an error. + msg = f"An array exists in store {store!r} at path {key!r}." + raise ContainsArrayError(msg) + + nodes_explicit: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata] = {} + + for k, v in nodes_parsed.items(): + if k not in redundant_implicit_groups: + if isinstance(v, ImplicitGroupMarker): + nodes_explicit[k] = GroupMetadata(zarr_format=v.zarr_format) + else: + nodes_explicit[k] = v + + async for key, node in create_nodes(store=store, nodes=nodes_explicit): + yield key, node + + +async def create_nodes( + *, + store: Store, + nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], +) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: + """Create a collection of arrays and / or groups concurrently. + + Note: no attempt is made to validate that these arrays and / or groups collectively form a + valid Zarr hierarchy. It is the responsibility of the caller of this function to ensure that + the ``nodes`` parameter satisfies any correctness constraints. + + Parameters + ---------- + store : Store + The storage backend to use. + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes + in the hierarchy, and the values are the metadata of the nodes. The + metadata must be either an instance of GroupMetadata, ArrayV3Metadata + or ArrayV2Metadata. + + Yields + ------ + AsyncGroup | AsyncArray + The created nodes in the order they are created. + """ + + # Note: the only way to alter this value is via the config. If that's undesirable for some reason, + # then we should consider adding a keyword argument this this function + semaphore = asyncio.Semaphore(config.get("async.concurrency")) + create_tasks: list[Coroutine[None, None, str]] = [] + + for key, value in nodes.items(): + # make the key absolute + create_tasks.extend(_persist_metadata(store, key, value, semaphore=semaphore)) + + created_object_keys = [] + + for coro in asyncio.as_completed(create_tasks): + created_key = await coro + # we need this to track which metadata documents were written so that we can yield a + # complete v2 Array / Group class after both .zattrs and the metadata JSON was created. + created_object_keys.append(created_key) + + # get the node name from the object key + if len(created_key.split("/")) == 1: + # this is the root node + meta_out = nodes[""] + node_name = "" + else: + # turn "foo/" into "foo" + node_name = created_key[: created_key.rfind("/")] + meta_out = nodes[node_name] + if meta_out.zarr_format == 3: + yield node_name, _build_node(store=store, path=node_name, metadata=meta_out) + else: + # For zarr v2 + # we only want to yield when both the metadata and attributes are created + # so we track which keys have been created, and wait for both the meta key and + # the attrs key to be created before yielding back the AsyncArray / AsyncGroup + + attrs_done = _join_paths([node_name, ZATTRS_JSON]) in created_object_keys + + if isinstance(meta_out, GroupMetadata): + meta_done = _join_paths([node_name, ZGROUP_JSON]) in created_object_keys + else: + meta_done = _join_paths([node_name, ZARRAY_JSON]) in created_object_keys + + if meta_done and attrs_done: + yield node_name, _build_node(store=store, path=node_name, metadata=meta_out) + + continue + + +def _get_roots( + data: Iterable[str], +) -> tuple[str, ...]: + """ + Return the keys of the root(s) of the hierarchy. A root is a key with the fewest number of + path segments. + """ + if "" in data: + return ("",) + keys_split = sorted((key.split("/") for key in data), key=len) + groups: defaultdict[int, list[str]] = defaultdict(list) + for key_split in keys_split: + groups[len(key_split)].append("/".join(key_split)) + return tuple(groups[min(groups.keys())]) + + +def _parse_hierarchy_dict( + *, + data: Mapping[str, ImplicitGroupMarker | GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], +) -> dict[str, ImplicitGroupMarker | GroupMetadata | ArrayV2Metadata | ArrayV3Metadata]: + """ + Take an input with type Mapping[str, ArrayMetadata | GroupMetadata] and parse it into + a dict of str: node pairs that models a valid, complete Zarr hierarchy. + + If the input represents a complete Zarr hierarchy, i.e. one with no implicit groups, + then return a dict with the exact same data as the input. + + Otherwise, return a dict derived from the input with GroupMetadata inserted as needed to make + the hierarchy complete. + + For example, an input of {'a/b': ArrayMetadata} is incomplete, because it references two + groups (the root group '' and a group at 'a') that are not specified in the input. Applying this function + to that input will result in a return value of + {'': GroupMetadata, 'a': GroupMetadata, 'a/b': ArrayMetadata}, i.e. the implied groups + were added. + + The input is also checked for the following conditions; an error is raised if any are violated: + + - No arrays can contain group or arrays (i.e., all arrays must be leaf nodes). + - All arrays and groups must have the same ``zarr_format`` value. + + This function ensures that the input is transformed into a specification of a complete and valid + Zarr hierarchy. + """ + + # ensure that all nodes have the same zarr format + data_purified = _ensure_consistent_zarr_format(data) + + # ensure that keys are normalized to zarr paths + data_normed_keys = _normalize_path_keys(data_purified) + + # insert an implicit root group if a root was not specified + # but not if an empty dict was provided, because any empty hierarchy has no nodes + if len(data_normed_keys) > 0 and "" not in data_normed_keys: + z_format = next(iter(data_normed_keys.values())).zarr_format + data_normed_keys = data_normed_keys | {"": ImplicitGroupMarker(zarr_format=z_format)} + + out: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata] = {**data_normed_keys} + + for k, v in data_normed_keys.items(): + key_split = k.split("/") + + # get every parent path + *subpaths, _ = accumulate(key_split, lambda a, b: _join_paths([a, b])) + + for subpath in subpaths: + # If a component is not already in the output dict, add ImplicitGroupMetadata + if subpath not in out: + out[subpath] = ImplicitGroupMarker(zarr_format=v.zarr_format) + else: + if not isinstance(out[subpath], GroupMetadata | ImplicitGroupMarker): + msg = ( + f"The node at {subpath} contains other nodes, but it is not a Zarr group. " + "This is invalid. Only Zarr groups can contain other nodes." + ) + raise ValueError(msg) + return out + + +def _ensure_consistent_zarr_format( + data: Mapping[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], +) -> Mapping[str, GroupMetadata | ArrayV2Metadata] | Mapping[str, GroupMetadata | ArrayV3Metadata]: + """ + Ensure that all values of the input dict have the same zarr format. If any do not, + then a value error is raised. + """ + observed_zarr_formats: dict[ZarrFormat, list[str]] = {2: [], 3: []} + + for k, v in data.items(): + observed_zarr_formats[v.zarr_format].append(k) + + if len(observed_zarr_formats[2]) > 0 and len(observed_zarr_formats[3]) > 0: + msg = ( + "Got data with both Zarr v2 and Zarr v3 nodes, which is invalid. " + f"The following keys map to Zarr v2 nodes: {observed_zarr_formats.get(2)}. " + f"The following keys map to Zarr v3 nodes: {observed_zarr_formats.get(3)}." + "Ensure that all nodes have the same Zarr format." + ) + raise ValueError(msg) + + return cast( + "Mapping[str, GroupMetadata | ArrayV2Metadata] | Mapping[str, GroupMetadata | ArrayV3Metadata]", + data, + ) + + +async def _getitem_semaphore( + node: AsyncGroup, key: str, semaphore: asyncio.Semaphore | None +) -> AnyAsyncArray | AsyncGroup: + """ + Wrap Group.getitem with an optional semaphore. + + If the semaphore parameter is an + asyncio.Semaphore instance, then the getitem operation is performed inside an async context + manager provided by that semaphore. If the semaphore parameter is None, then getitem is invoked + without a context manager. + """ + if semaphore is not None: + async with semaphore: + return await node.getitem(key) + else: + return await node.getitem(key) + + +async def _iter_members( + node: AsyncGroup, + skip_keys: tuple[str, ...], + semaphore: asyncio.Semaphore | None, +) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: + """ + Iterate over the arrays and groups contained in a group. + + Parameters + ---------- + node : AsyncGroup + The group to traverse. + skip_keys : tuple[str, ...] + A tuple of keys to skip when iterating over the possible members of the group. + semaphore : asyncio.Semaphore | None + An optional semaphore to use for concurrency control. + + Yields + ------ + tuple[str, AnyAsyncArray | AsyncGroup] + """ + + # retrieve keys from storage + keys = [key async for key in node.store.list_dir(node.path)] + keys_filtered = tuple(filter(lambda v: v not in skip_keys, keys)) + + node_tasks = tuple( + asyncio.create_task(_getitem_semaphore(node, key, semaphore), name=key) + for key in keys_filtered + ) + + for fetched_node_coro in asyncio.as_completed(node_tasks): + try: + fetched_node = await fetched_node_coro + except KeyError as e: + # keyerror is raised when `key` names an object (in the object storage sense), + # as opposed to a prefix, in the store under the prefix associated with this group + # in which case `key` cannot be the name of a sub-array or sub-group. + warnings.warn( + f"Object at {e.args[0]} is not recognized as a component of a Zarr hierarchy.", + ZarrUserWarning, + stacklevel=1, + ) + continue + match fetched_node: + case AsyncArray() | AsyncGroup(): + yield fetched_node.basename, fetched_node + case _: + raise ValueError(f"Unexpected type: {type(fetched_node)}") + + +async def _iter_members_deep( + group: AsyncGroup, + *, + max_depth: int | None, + skip_keys: tuple[str, ...], + semaphore: asyncio.Semaphore | None = None, + use_consolidated_for_children: bool = True, +) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: + """ + Iterate over the arrays and groups contained in a group, and optionally the + arrays and groups contained in those groups. + + Parameters + ---------- + group : AsyncGroup + The group to traverse. + max_depth : int | None + The maximum depth of recursion. + skip_keys : tuple[str, ...] + A tuple of keys to skip when iterating over the possible members of the group. + semaphore : asyncio.Semaphore | None + An optional semaphore to use for concurrency control. + use_consolidated_for_children : bool, default True + Whether to use the consolidated metadata of child groups loaded + from the store. Note that this only affects groups loaded from the + store. If the current Group already has consolidated metadata, it + will always be used. + + Yields + ------ + tuple[str, AnyAsyncArray | AsyncGroup] + """ + + to_recurse = {} + do_recursion = max_depth is None or max_depth > 0 + + if max_depth is None: + new_depth = None + else: + new_depth = max_depth - 1 + async for name, node in _iter_members(group, skip_keys=skip_keys, semaphore=semaphore): + is_group = isinstance(node, AsyncGroup) + if ( + is_group + and not use_consolidated_for_children + and node.metadata.consolidated_metadata is not None + ): + node = cast("AsyncGroup", node) + # We've decided not to trust consolidated metadata at this point, because we're + # reconsolidating the metadata, for example. + node = replace(node, metadata=replace(node.metadata, consolidated_metadata=None)) + yield name, node + if is_group and do_recursion: + node = cast("AsyncGroup", node) + to_recurse[name] = _iter_members_deep( + node, max_depth=new_depth, skip_keys=skip_keys, semaphore=semaphore + ) + + for prefix, subgroup_iter in to_recurse.items(): + async for name, node in subgroup_iter: + key = f"{prefix}/{name}".lstrip("/") + yield key, node + + +async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata: + """ + Given a store_path, return ArrayV3Metadata or GroupMetadata defined by the metadata + document stored at store_path.path / zarr.json. If no such document is found, raise a + FileNotFoundError. + """ + zarr_json_bytes = await store.get( + _join_paths([path, ZARR_JSON]), prototype=default_buffer_prototype() + ) + if zarr_json_bytes is None: + raise FileNotFoundError(path) + else: + zarr_json = json.loads(zarr_json_bytes.to_bytes()) + return _build_metadata_v3(zarr_json) + + +async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupMetadata: + """ + Given a store_path, return ArrayV2Metadata or GroupMetadata defined by the metadata + document stored at store_path.path / (.zgroup | .zarray). If no such document is found, + raise a FileNotFoundError. + """ + # TODO: consider first fetching array metadata, and only fetching group metadata when we don't + # find an array + zarray_bytes, zgroup_bytes, zattrs_bytes = await asyncio.gather( + store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()), + store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()), + store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()), + ) + + if zattrs_bytes is None: + zattrs = {} + else: + zattrs = json.loads(zattrs_bytes.to_bytes()) + + # TODO: decide how to handle finding both array and group metadata. The spec does not seem to + # consider this situation. A practical approach would be to ignore that combination, and only + # return the array metadata. + if zarray_bytes is not None: + zmeta = json.loads(zarray_bytes.to_bytes()) + else: + if zgroup_bytes is None: + # neither .zarray or .zgroup were found results in KeyError + raise FileNotFoundError(path) + else: + zmeta = json.loads(zgroup_bytes.to_bytes()) + + return _build_metadata_v2(zmeta, zattrs) + + +async def _read_group_metadata_v2(store: Store, path: str) -> GroupMetadata: + """ + Read group metadata or error + """ + meta = await _read_metadata_v2(store=store, path=path) + if not isinstance(meta, GroupMetadata): + raise FileNotFoundError(f"Group metadata was not found in {store} at {path}") + return meta + + +async def _read_group_metadata_v3(store: Store, path: str) -> GroupMetadata: + """ + Read group metadata or error + """ + meta = await _read_metadata_v3(store=store, path=path) + if not isinstance(meta, GroupMetadata): + raise FileNotFoundError(f"Group metadata was not found in {store} at {path}") + return meta + + +async def _read_group_metadata( + store: Store, path: str, *, zarr_format: ZarrFormat +) -> GroupMetadata: + if zarr_format == 2: + return await _read_group_metadata_v2(store=store, path=path) + return await _read_group_metadata_v3(store=store, path=path) + + +def _build_metadata_v3(zarr_json: dict[str, JSON]) -> ArrayV3Metadata | GroupMetadata: + """ + Convert a dict representation of Zarr V3 metadata into the corresponding metadata class. + """ + if "node_type" not in zarr_json: + msg = "Required key 'node_type' is missing from the provided metadata document." + raise MetadataValidationError(msg) + match zarr_json: + case {"node_type": "array"}: + return ArrayV3Metadata.from_dict(zarr_json) + case {"node_type": "group"}: + return GroupMetadata.from_dict(zarr_json) + case _: # pragma: no cover + raise ValueError( + "invalid value for `node_type` key in metadata document" + ) # pragma: no cover + + +def _build_metadata_v2( + zarr_json: dict[str, JSON], attrs_json: dict[str, JSON] +) -> ArrayV2Metadata | GroupMetadata: + """ + Convert a dict representation of Zarr V2 metadata into the corresponding metadata class. + """ + match zarr_json: + case {"shape": _}: + return ArrayV2Metadata.from_dict(zarr_json | {"attributes": attrs_json}) + case _: # pragma: no cover + return GroupMetadata.from_dict(zarr_json | {"attributes": attrs_json}) + + +@overload +def _build_node(*, store: Store, path: str, metadata: ArrayV2Metadata) -> AsyncArrayV2: ... + + +@overload +def _build_node(*, store: Store, path: str, metadata: ArrayV3Metadata) -> AsyncArrayV3: ... + + +@overload +def _build_node(*, store: Store, path: str, metadata: GroupMetadata) -> AsyncGroup: ... + + +def _build_node( + *, store: Store, path: str, metadata: ArrayV3Metadata | ArrayV2Metadata | GroupMetadata +) -> AnyAsyncArray | AsyncGroup: + """ + Take a metadata object and return a node (AsyncArray or AsyncGroup). + """ + store_path = StorePath(store=store, path=path) + match metadata: + case ArrayV2Metadata() | ArrayV3Metadata(): + return AsyncArray(metadata, store_path=store_path) + case GroupMetadata(): + return AsyncGroup(metadata, store_path=store_path) + case _: # pragma: no cover + raise ValueError(f"Unexpected metadata type: {type(metadata)}") # pragma: no cover + + +async def _get_node_v2(store: Store, path: str) -> AsyncArrayV2 | AsyncGroup: + """ + Read a Zarr v2 AsyncArray or AsyncGroup from a path in a Store. + + Parameters + ---------- + store : Store + The store-like object to read from. + path : str + The path to the node to read. + + Returns + ------- + AsyncArray | AsyncGroup + """ + metadata = await _read_metadata_v2(store=store, path=path) + return _build_node(store=store, path=path, metadata=metadata) + + +async def _get_node_v3(store: Store, path: str) -> AsyncArrayV3 | AsyncGroup: + """ + Read a Zarr v3 AsyncArray or AsyncGroup from a path in a Store. + + Parameters + ---------- + store : Store + The store-like object to read from. + path : str + The path to the node to read. + + Returns + ------- + AsyncArray | AsyncGroup + """ + metadata = await _read_metadata_v3(store=store, path=path) + return _build_node(store=store, path=path, metadata=metadata) + + +async def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyAsyncArray | AsyncGroup: + """ + Get an AsyncArray or AsyncGroup from a path in a Store. + + Parameters + ---------- + store : Store + The store-like object to read from. + path : str + The path to the node to read. + zarr_format : {2, 3} + The zarr format of the node to read. + + Returns + ------- + AsyncArray | AsyncGroup + """ + + match zarr_format: + case 2: + return await _get_node_v2(store=store, path=path) + case 3: + return await _get_node_v3(store=store, path=path) + case _: # pragma: no cover + raise ValueError(f"Unexpected zarr format: {zarr_format}") # pragma: no cover + + +async def _set_return_key( + *, store: Store, key: str, value: Buffer, semaphore: asyncio.Semaphore | None = None +) -> str: + """ + Write a value to storage at the given key. The key is returned. + Useful when saving values via routines that return results in execution order, + like asyncio.as_completed, because in this case we need to know which key was saved in order + to yield the right object to the caller. + + Parameters + ---------- + store : Store + The store to save the value to. + key : str + The key to save the value to. + value : Buffer + The value to save. + semaphore : asyncio.Semaphore | None + An optional semaphore to use to limit the number of concurrent writes. + """ + + if semaphore is not None: + async with semaphore: + await store.set(key, value) + else: + await store.set(key, value) + return key + + +def _persist_metadata( + store: Store, + path: str, + metadata: ArrayV2Metadata | ArrayV3Metadata | GroupMetadata, + semaphore: asyncio.Semaphore | None = None, +) -> tuple[Coroutine[None, None, str], ...]: + """ + Prepare to save a metadata document to storage, returning a tuple of coroutines that must be awaited. + """ + + to_save = metadata.to_buffer_dict(default_buffer_prototype()) + return tuple( + _set_return_key(store=store, key=_join_paths([path, key]), value=value, semaphore=semaphore) + for key, value in to_save.items() + ) + + +async def create_rooted_hierarchy( + *, + store: Store, + nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], + overwrite: bool = False, +) -> AsyncGroup | AnyAsyncArray: + """ + Create an ``AsyncGroup`` or ``AsyncArray`` from a store and a dict of metadata documents. + This function ensures that its input contains a specification of a root node, + calls ``create_hierarchy`` to create nodes, and returns the root node of the hierarchy. + """ + roots = _get_roots(nodes.keys()) + if len(roots) != 1: + msg = ( + "The input does not specify a root node. " + "This function can only create hierarchies that contain a root node, which is " + "defined as a group that is ancestral to all the other arrays and " + "groups in the hierarchy, or a single array." + ) + raise ValueError(msg) + else: + root_key = roots[0] + + nodes_created = [ + x async for x in create_hierarchy(store=store, nodes=nodes, overwrite=overwrite) + ] + return dict(nodes_created)[root_key] diff --git a/.venv/lib/python3.11/site-packages/zarr/core/indexing.py b/.venv/lib/python3.11/site-packages/zarr/core/indexing.py new file mode 100644 index 0000000000000000000000000000000000000000..d226c036757ecd92104dc0ca7e7e6b64f522f0e3 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/indexing.py @@ -0,0 +1,1595 @@ +from __future__ import annotations + +import itertools +import math +import numbers +import operator +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache, reduce +from types import EllipsisType +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Literal, + NamedTuple, + Protocol, + TypeAlias, + TypeGuard, + TypeVar, + cast, + runtime_checkable, +) + +import numpy as np +import numpy.typing as npt + +from zarr.core.common import ceildiv, product +from zarr.core.metadata import T_ArrayMetadata +from zarr.errors import ( + ArrayIndexError, + BoundsCheckError, + NegativeStepError, + VindexInvalidSelectionError, +) + +if TYPE_CHECKING: + from zarr.core.array import AsyncArray + from zarr.core.buffer import NDArrayLikeOrScalar + from zarr.core.chunk_grids import ChunkGrid + from zarr.types import AnyArray + + +IntSequence = list[int] | npt.NDArray[np.intp] +ArrayOfIntOrBool = npt.NDArray[np.intp] | npt.NDArray[np.bool_] +BasicSelector = int | slice | EllipsisType +Selector = BasicSelector | ArrayOfIntOrBool +BasicSelection = BasicSelector | tuple[BasicSelector, ...] # also used for BlockIndex +CoordinateSelection = IntSequence | tuple[IntSequence, ...] +MaskSelection = npt.NDArray[np.bool_] +OrthogonalSelection = Selector | tuple[Selector, ...] +Selection = BasicSelection | CoordinateSelection | MaskSelection | OrthogonalSelection +CoordinateSelectionNormalized = tuple[npt.NDArray[np.intp], ...] +SelectionNormalized = tuple[Selector, ...] | ArrayOfIntOrBool +SelectionWithFields = Selection | str | Sequence[str] +SelectorTuple = tuple[Selector, ...] | npt.NDArray[np.intp] | slice +Fields = str | list[str] | tuple[str, ...] + + +def err_too_many_indices(selection: Any, shape: tuple[int, ...]) -> None: + raise IndexError(f"too many indices for array; expected {len(shape)}, got {len(selection)}") + + +def _zarr_array_to_int_or_bool_array(arr: AnyArray) -> npt.NDArray[np.intp] | npt.NDArray[np.bool_]: + if arr.dtype.kind in ("i", "b"): + return np.asarray(arr) + else: + raise IndexError( + f"Invalid array dtype: {arr.dtype}. Arrays used as indices must be of integer or boolean type" + ) + + +@runtime_checkable +class Indexer(Protocol): + shape: tuple[int, ...] + drop_axes: tuple[int, ...] + + def __iter__(self) -> Iterator[ChunkProjection]: ... + + +_ArrayIndexingOrder: TypeAlias = Literal["lexicographic"] + + +def _iter_grid( + grid_shape: Sequence[int], + *, + origin: Sequence[int] | None = None, + selection_shape: Sequence[int] | None = None, + order: _ArrayIndexingOrder = "lexicographic", +) -> Iterator[tuple[int, ...]]: + """ + Iterate over the elements of grid of integers, with the option to restrict the domain of + iteration to a contiguous subregion of that grid. + + Parameters + ---------- + grid_shape : Sequence[int] + The size of the domain to iterate over. + origin : Sequence[int] | None, default=None + The first coordinate of the domain to return. + selection_shape : Sequence[int] | None, default=None + The shape of the selection. + order : Literal["lexicographic"], default="lexicographic" + The linear indexing order to use. + + Returns + ------- + + Iterator[tuple[int, ...]] + An iterator over tuples of integers + + Examples + -------- + ```python + from zarr.core.indexing import _iter_grid + tuple(_iter_grid((1,))) + # ((0,),) + + tuple(_iter_grid((2,3))) + # ((0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)) + + tuple(_iter_grid((2,3), origin=(1,1))) + # ((1, 1), (1, 2)) + + tuple(_iter_grid((2,3), origin=(0,0), selection_shape=(2,2))) + # ((0, 0), (0, 1), (1, 0), (1, 1)) + ``` + """ + if origin is None: + origin_parsed = (0,) * len(grid_shape) + else: + if len(origin) != len(grid_shape): + msg = ( + "Shape and origin parameters must have the same length." + f"Got {len(grid_shape)} elements in shape, but {len(origin)} elements in origin." + ) + raise ValueError(msg) + origin_parsed = tuple(origin) + if selection_shape is None: + selection_shape_parsed = tuple( + g - o for o, g in zip(origin_parsed, grid_shape, strict=True) + ) + else: + selection_shape_parsed = tuple(selection_shape) + if order == "lexicographic": + dimensions: tuple[range, ...] = () + for idx, (o, gs, ss) in enumerate( + zip(origin_parsed, grid_shape, selection_shape_parsed, strict=True) + ): + if o + ss > gs: + raise IndexError( + f"Invalid selection shape ({ss}) for origin ({o}) and grid shape ({gs}) at axis {idx}." + ) + dimensions += (range(o, o + ss),) + return itertools.product(*(dimensions)) + + else: + msg = f"Indexing order {order} is not supported at this time." # type: ignore[unreachable] # pragma: no cover + raise NotImplementedError(msg) # pragma: no cover + + +def _iter_regions( + domain_shape: Sequence[int], + region_shape: Sequence[int], + *, + origin: Sequence[int] | None = None, + selection_shape: Sequence[int] | None = None, + order: _ArrayIndexingOrder = "lexicographic", + trim_excess: bool = True, +) -> Iterator[tuple[slice, ...]]: + """ + Iterate over contiguous regions on a grid of integers, with the option to restrict the + domain of iteration to a contiguous subregion of that grid. + + Parameters + ---------- + domain_shape : Sequence[int] + The size of the domain to iterate over. + region_shape : Sequence[int] + The shape of the region to iterate over. + origin : Sequence[int] | None, default=None + The location, in grid coordinates, of the first region to return. + selection_shape : Sequence[int] | None, default=None + The shape of the selection, in grid coordinates. + order : Literal["lexicographic"], default="lexicographic" + The linear indexing order to use. + + Yields + ------ + + Iterator[tuple[slice, ...]] + An iterator over tuples of slices, where each slice spans a separate contiguous region + + Examples + -------- + ```python + from zarr.core.indexing import _iter_regions + tuple(_iter_regions((1,), (1,))) + # ((slice(0, 1, 1),),) + + tuple(_iter_regions((2, 3), (1, 2))) + # ((slice(0, 1, 1), slice(0, 2, 1)), (slice(1, 2, 1), slice(0, 2, 1))) + + tuple(_iter_regions((2,3), (1,2), origin=(1,1))) + # ((slice(1, 2, 1), slice(1, 3, 1)), (slice(2, 3, 1), slice(1, 3, 1))) + + tuple(_iter_regions((2,3), (1,2), origin=(0,0), selection_shape=(2,2))) + # ((slice(0, 1, 1), slice(0, 2, 1)), (slice(1, 2, 1), slice(0, 2, 1))) + ``` + """ + grid_shape = tuple(itertools.starmap(ceildiv, zip(domain_shape, region_shape, strict=True))) + for grid_position in _iter_grid( + grid_shape=grid_shape, origin=origin, selection_shape=selection_shape, order=order + ): + out: list[slice] = [] + for g_pos, r_shape, d_shape in zip(grid_position, region_shape, domain_shape, strict=True): + start = g_pos * r_shape + stop = start + r_shape + if trim_excess: + stop = min(stop, d_shape) + out.append(slice(start, stop, 1)) + yield tuple(out) + + +def is_integer(x: Any) -> TypeGuard[int]: + """True if x is an integer (both pure Python or NumPy).""" + return isinstance(x, numbers.Integral) and not is_bool(x) + + +def is_bool(x: Any) -> TypeGuard[bool | np.bool_]: + """True if x is a boolean (both pure Python or NumPy).""" + return type(x) in [bool, np.bool_] + + +def is_integer_list(x: Any) -> TypeGuard[list[int]]: + """True if x is a list of integers.""" + return isinstance(x, list) and len(x) > 0 and all(is_integer(i) for i in x) + + +def is_bool_list(x: Any) -> TypeGuard[list[bool | np.bool_]]: + """True if x is a list of boolean.""" + return isinstance(x, list) and len(x) > 0 and all(is_bool(i) for i in x) + + +def is_integer_array(x: Any, ndim: int | None = None) -> TypeGuard[npt.NDArray[np.intp]]: + t = not np.isscalar(x) and hasattr(x, "shape") and hasattr(x, "dtype") and x.dtype.kind in "ui" + if ndim is not None: + t = t and hasattr(x, "shape") and len(x.shape) == ndim + return t + + +def is_bool_array(x: Any, ndim: int | None = None) -> TypeGuard[npt.NDArray[np.bool_]]: + t = hasattr(x, "shape") and hasattr(x, "dtype") and x.dtype == bool + if ndim is not None: + t = t and hasattr(x, "shape") and len(x.shape) == ndim + return t + + +def is_int_or_bool_iterable(x: Any) -> bool: + return is_integer_list(x) or is_integer_array(x) or is_bool_array(x) or is_bool_list(x) + + +def is_scalar(value: Any, dtype: np.dtype[Any]) -> bool: + if np.isscalar(value): + return True + if hasattr(value, "shape") and value.shape == (): + return True + return isinstance(value, tuple) and dtype.names is not None and len(value) == len(dtype.names) + + +def is_pure_fancy_indexing(selection: Any, ndim: int) -> bool: + """Check whether a selection contains only scalars or integer/bool array-likes. + + Parameters + ---------- + selection : tuple, slice, or scalar + A valid selection value for indexing into arrays. + + Returns + ------- + is_pure : bool + True if the selection is a pure fancy indexing expression (ie not mixed + with boolean or slices). + """ + if is_bool_array(selection): + # is mask selection + return True + + if ndim == 1 and ( + is_integer_list(selection) or is_integer_array(selection) or is_bool_list(selection) + ): + return True + + # if not, we go through the normal path below, because a 1-tuple + # of integers is also allowed. + no_slicing = ( + isinstance(selection, tuple) + and len(selection) == ndim + and not (any(isinstance(elem, slice) or elem is Ellipsis for elem in selection)) + ) + return ( + no_slicing + and all( + is_integer(elem) or is_integer_list(elem) or is_integer_array(elem) + for elem in selection + ) + and any(is_integer_list(elem) or is_integer_array(elem) for elem in selection) + ) + + +def is_pure_orthogonal_indexing(selection: Selection, ndim: int) -> TypeGuard[OrthogonalSelection]: + if not ndim: + return False + + selection_normalized = (selection,) if not isinstance(selection, tuple) else selection + + # Case 1: Selection contains of iterable of integers or boolean + if len(selection_normalized) == ndim and all( + is_int_or_bool_iterable(s) for s in selection_normalized + ): + return True + + # Case 2: selection contains either zero or one integer iterables. + # All other selection elements are slices or integers + return ( + len(selection_normalized) <= ndim + and sum(is_int_or_bool_iterable(s) for s in selection_normalized) <= 1 + and all( + is_int_or_bool_iterable(s) or isinstance(s, int | slice) for s in selection_normalized + ) + ) + + +def get_chunk_shape(chunk_grid: ChunkGrid) -> tuple[int, ...]: + from zarr.core.chunk_grids import RegularChunkGrid + + assert isinstance(chunk_grid, RegularChunkGrid), ( + "Only regular chunk grid is supported, currently." + ) + return chunk_grid.chunk_shape + + +def normalize_integer_selection(dim_sel: int, dim_len: int) -> int: + # normalize type to int + dim_sel = int(dim_sel) + + # handle wraparound + if dim_sel < 0: + dim_sel = dim_len + dim_sel + + # handle out of bounds + if dim_sel >= dim_len or dim_sel < 0: + msg = f"index out of bounds for dimension with length {dim_len}" + raise BoundsCheckError(msg) + + return dim_sel + + +class ChunkDimProjection(NamedTuple): + """A mapping from chunk to output array for a single dimension. + + Attributes + ---------- + dim_chunk_ix + Index of chunk. + dim_chunk_sel + Selection of items from chunk array. + dim_out_sel + Selection of items in target (output) array. + """ + + dim_chunk_ix: int + dim_chunk_sel: Selector + dim_out_sel: Selector | None + is_complete_chunk: bool + + +@dataclass(frozen=True) +class IntDimIndexer: + dim_sel: int + dim_len: int + dim_chunk_len: int + nitems: int = 1 + + def __init__(self, dim_sel: int, dim_len: int, dim_chunk_len: int) -> None: + object.__setattr__(self, "dim_sel", normalize_integer_selection(dim_sel, dim_len)) + object.__setattr__(self, "dim_len", dim_len) + object.__setattr__(self, "dim_chunk_len", dim_chunk_len) + + def __iter__(self) -> Iterator[ChunkDimProjection]: + dim_chunk_ix = self.dim_sel // self.dim_chunk_len + dim_offset = dim_chunk_ix * self.dim_chunk_len + dim_chunk_sel = self.dim_sel - dim_offset + dim_out_sel = None + is_complete_chunk = self.dim_chunk_len == 1 + yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) + + +@dataclass(frozen=True) +class SliceDimIndexer: + dim_len: int + dim_chunk_len: int + nitems: int + nchunks: int + + start: int + stop: int + step: int + + def __init__(self, dim_sel: slice, dim_len: int, dim_chunk_len: int) -> None: + # normalize + start, stop, step = dim_sel.indices(dim_len) + if step < 1: + raise NegativeStepError("only slices with step >= 1 are supported.") + + object.__setattr__(self, "start", start) + object.__setattr__(self, "stop", stop) + object.__setattr__(self, "step", step) + + object.__setattr__(self, "dim_len", dim_len) + object.__setattr__(self, "dim_chunk_len", dim_chunk_len) + object.__setattr__(self, "nitems", max(0, ceildiv((stop - start), step))) + object.__setattr__(self, "nchunks", ceildiv(dim_len, dim_chunk_len)) + + def __iter__(self) -> Iterator[ChunkDimProjection]: + # figure out the range of chunks we need to visit + dim_chunk_ix_from = 0 if self.start == 0 else self.start // self.dim_chunk_len + dim_chunk_ix_to = ceildiv(self.stop, self.dim_chunk_len) + + # iterate over chunks in range + for dim_chunk_ix in range(dim_chunk_ix_from, dim_chunk_ix_to): + # compute offsets for chunk within overall array + dim_offset = dim_chunk_ix * self.dim_chunk_len + dim_limit = min(self.dim_len, (dim_chunk_ix + 1) * self.dim_chunk_len) + + # determine chunk length, accounting for trailing chunk + dim_chunk_len = dim_limit - dim_offset + + if self.start < dim_offset: + # selection starts before current chunk + dim_chunk_sel_start = 0 + remainder = (dim_offset - self.start) % self.step + if remainder: + dim_chunk_sel_start += self.step - remainder + # compute number of previous items, provides offset into output array + dim_out_offset = ceildiv((dim_offset - self.start), self.step) + + else: + # selection starts within current chunk + dim_chunk_sel_start = self.start - dim_offset + dim_out_offset = 0 + + if self.stop > dim_limit: + # selection ends after current chunk + dim_chunk_sel_stop = dim_chunk_len + + else: + # selection ends within current chunk + dim_chunk_sel_stop = self.stop - dim_offset + + dim_chunk_sel = slice(dim_chunk_sel_start, dim_chunk_sel_stop, self.step) + dim_chunk_nitems = ceildiv((dim_chunk_sel_stop - dim_chunk_sel_start), self.step) + + # If there are no elements on the selection within this chunk, then skip + if dim_chunk_nitems == 0: + continue + + dim_out_sel = slice(dim_out_offset, dim_out_offset + dim_chunk_nitems) + + is_complete_chunk = ( + dim_chunk_sel_start == 0 and (self.stop >= dim_limit) and self.step in [1, None] + ) + yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) + + +def check_selection_length(selection: SelectionNormalized, shape: tuple[int, ...]) -> None: + if len(selection) > len(shape): + err_too_many_indices(selection, shape) + + +def replace_ellipsis(selection: Any, shape: tuple[int, ...]) -> SelectionNormalized: + selection = ensure_tuple(selection) + + # count number of ellipsis present + n_ellipsis = sum(1 for i in selection if i is Ellipsis) + + if n_ellipsis > 1: + # more than 1 is an error + raise IndexError("an index can only have a single ellipsis ('...')") + + elif n_ellipsis == 1: + # locate the ellipsis, count how many items to left and right + n_items_l = selection.index(Ellipsis) # items to left of ellipsis + n_items_r = len(selection) - (n_items_l + 1) # items to right of ellipsis + n_items = len(selection) - 1 # all non-ellipsis items + + if n_items >= len(shape): + # ellipsis does nothing, just remove it + selection = tuple(i for i in selection if i != Ellipsis) + + else: + # replace ellipsis with as many slices are needed for number of dims + new_item = selection[:n_items_l] + ((slice(None),) * (len(shape) - n_items)) + if n_items_r: + new_item += selection[-n_items_r:] + selection = new_item + + # fill out selection if not completely specified + if len(selection) < len(shape): + selection += (slice(None),) * (len(shape) - len(selection)) + + # check selection not too long + check_selection_length(selection, shape) + + return cast("SelectionNormalized", selection) + + +def replace_lists(selection: SelectionNormalized) -> SelectionNormalized: + return tuple( + np.asarray(dim_sel) if isinstance(dim_sel, list) else dim_sel for dim_sel in selection + ) + + +T = TypeVar("T") + + +def ensure_tuple(v: Any) -> SelectionNormalized: + if not isinstance(v, tuple): + v = (v,) + return cast("SelectionNormalized", v) + + +class ChunkProjection(NamedTuple): + """A mapping of items from chunk to output array. Can be used to extract items from the + chunk array for loading into an output array. Can also be used to extract items from a + value array for setting/updating in a chunk array. + + Attributes + ---------- + chunk_coords + Indices of chunk. + chunk_selection + Selection of items from chunk array. + out_selection + Selection of items in target (output) array. + is_complete_chunk: + True if a complete chunk is indexed + """ + + chunk_coords: tuple[int, ...] + chunk_selection: tuple[Selector, ...] | npt.NDArray[np.intp] + out_selection: tuple[Selector, ...] | npt.NDArray[np.intp] | slice + is_complete_chunk: bool + + +def is_slice(s: Any) -> TypeGuard[slice]: + return isinstance(s, slice) + + +def is_contiguous_slice(s: Any) -> TypeGuard[slice]: + return is_slice(s) and (s.step is None or s.step == 1) + + +def is_positive_slice(s: Any) -> TypeGuard[slice]: + return is_slice(s) and (s.step is None or s.step >= 1) + + +def is_contiguous_selection(selection: Any) -> TypeGuard[slice]: + selection = ensure_tuple(selection) + return all((is_integer_array(s) or is_contiguous_slice(s) or s == Ellipsis) for s in selection) + + +def is_basic_selection(selection: Any) -> TypeGuard[BasicSelection]: + selection = ensure_tuple(selection) + return all(is_integer(s) or is_positive_slice(s) for s in selection) + + +@dataclass(frozen=True) +class BasicIndexer(Indexer): + dim_indexers: list[IntDimIndexer | SliceDimIndexer] + shape: tuple[int, ...] + drop_axes: tuple[int, ...] + + def __init__( + self, + selection: BasicSelection, + shape: tuple[int, ...], + chunk_grid: ChunkGrid, + ) -> None: + chunk_shape = get_chunk_shape(chunk_grid) + # handle ellipsis + selection_normalized = replace_ellipsis(selection, shape) + + # setup per-dimension indexers + dim_indexers: list[IntDimIndexer | SliceDimIndexer] = [] + for dim_sel, dim_len, dim_chunk_len in zip( + selection_normalized, shape, chunk_shape, strict=True + ): + dim_indexer: IntDimIndexer | SliceDimIndexer + if is_integer(dim_sel): + dim_indexer = IntDimIndexer(dim_sel, dim_len, dim_chunk_len) + + elif is_slice(dim_sel): + dim_indexer = SliceDimIndexer(dim_sel, dim_len, dim_chunk_len) + + else: + raise IndexError( + "unsupported selection item for basic indexing; " + f"expected integer or slice, got {type(dim_sel)!r}" + ) + + dim_indexers.append(dim_indexer) + + object.__setattr__(self, "dim_indexers", dim_indexers) + object.__setattr__( + self, + "shape", + tuple(s.nitems for s in self.dim_indexers if not isinstance(s, IntDimIndexer)), + ) + object.__setattr__(self, "drop_axes", ()) + + def __iter__(self) -> Iterator[ChunkProjection]: + for dim_projections in itertools.product(*self.dim_indexers): + chunk_coords = tuple(p.dim_chunk_ix for p in dim_projections) + chunk_selection = tuple(p.dim_chunk_sel for p in dim_projections) + out_selection = tuple( + p.dim_out_sel for p in dim_projections if p.dim_out_sel is not None + ) + is_complete_chunk = all(p.is_complete_chunk for p in dim_projections) + yield ChunkProjection(chunk_coords, chunk_selection, out_selection, is_complete_chunk) + + +@dataclass(frozen=True) +class BoolArrayDimIndexer: + dim_sel: npt.NDArray[np.bool_] + dim_len: int + dim_chunk_len: int + nchunks: int + + chunk_nitems: npt.NDArray[Any] + chunk_nitems_cumsum: npt.NDArray[Any] + nitems: int + dim_chunk_ixs: npt.NDArray[np.intp] + + def __init__(self, dim_sel: npt.NDArray[np.bool_], dim_len: int, dim_chunk_len: int) -> None: + # check number of dimensions + if not is_bool_array(dim_sel, 1): + raise IndexError("Boolean arrays in an orthogonal selection must be 1-dimensional only") + + # check shape + if dim_sel.shape[0] != dim_len: + raise IndexError( + f"Boolean array has the wrong length for dimension; expected {dim_len}, got {dim_sel.shape[0]}" + ) + + # precompute number of selected items for each chunk + nchunks = ceildiv(dim_len, dim_chunk_len) + chunk_nitems = np.zeros(nchunks, dtype="i8") + for dim_chunk_ix in range(nchunks): + dim_offset = dim_chunk_ix * dim_chunk_len + chunk_nitems[dim_chunk_ix] = np.count_nonzero( + dim_sel[dim_offset : dim_offset + dim_chunk_len] + ) + chunk_nitems_cumsum = np.cumsum(chunk_nitems) + nitems = chunk_nitems_cumsum[-1] + dim_chunk_ixs = np.nonzero(chunk_nitems)[0] + + # store attributes + object.__setattr__(self, "dim_sel", dim_sel) + object.__setattr__(self, "dim_len", dim_len) + object.__setattr__(self, "dim_chunk_len", dim_chunk_len) + object.__setattr__(self, "nchunks", nchunks) + object.__setattr__(self, "chunk_nitems", chunk_nitems) + object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum) + object.__setattr__(self, "nitems", nitems) + object.__setattr__(self, "dim_chunk_ixs", dim_chunk_ixs) + + def __iter__(self) -> Iterator[ChunkDimProjection]: + # iterate over chunks with at least one item + for dim_chunk_ix in self.dim_chunk_ixs: + # find region in chunk + dim_offset = dim_chunk_ix * self.dim_chunk_len + dim_chunk_sel = self.dim_sel[dim_offset : dim_offset + self.dim_chunk_len] + + # pad out if final chunk + if dim_chunk_sel.shape[0] < self.dim_chunk_len: + tmp = np.zeros(self.dim_chunk_len, dtype=bool) + tmp[: dim_chunk_sel.shape[0]] = dim_chunk_sel + dim_chunk_sel = tmp + + # find region in output + if dim_chunk_ix == 0: + start = 0 + else: + start = self.chunk_nitems_cumsum[dim_chunk_ix - 1] + stop = self.chunk_nitems_cumsum[dim_chunk_ix] + dim_out_sel = slice(start, stop) + is_complete_chunk = False # TODO + + yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) + + +class Order(Enum): + """ + Enum for indexing order. + """ + + UNKNOWN = 0 + INCREASING = 1 + DECREASING = 2 + UNORDERED = 3 + + @staticmethod + def check(a: npt.NDArray[Any]) -> Order: + diff = np.diff(a) + diff_positive = diff >= 0 + n_diff_positive = np.count_nonzero(diff_positive) + all_increasing = n_diff_positive == len(diff_positive) + any_increasing = n_diff_positive > 0 + if all_increasing: + order = Order.INCREASING + elif any_increasing: + order = Order.UNORDERED + else: + order = Order.DECREASING + return order + + +def wraparound_indices(x: npt.NDArray[Any], dim_len: int) -> None: + loc_neg = x < 0 + if np.any(loc_neg): + x[loc_neg] += dim_len + + +def boundscheck_indices(x: npt.NDArray[Any], dim_len: int) -> None: + if np.any(x < 0) or np.any(x >= dim_len): + msg = f"index out of bounds for dimension with length {dim_len}" + raise BoundsCheckError(msg) + + +@dataclass(frozen=True) +class IntArrayDimIndexer: + """Integer array selection against a single dimension.""" + + dim_len: int + dim_chunk_len: int + nchunks: int + nitems: int + order: Order + dim_sel: npt.NDArray[np.intp] + dim_out_sel: npt.NDArray[np.intp] + chunk_nitems: int + dim_chunk_ixs: npt.NDArray[np.intp] + chunk_nitems_cumsum: npt.NDArray[np.intp] + + def __init__( + self, + dim_sel: npt.NDArray[np.intp], + dim_len: int, + dim_chunk_len: int, + wraparound: bool = True, + boundscheck: bool = True, + order: Order = Order.UNKNOWN, + ) -> None: + # ensure 1d array + dim_sel = np.asanyarray(dim_sel) + if not is_integer_array(dim_sel, 1): + raise IndexError("integer arrays in an orthogonal selection must be 1-dimensional only") + + nitems = len(dim_sel) + nchunks = ceildiv(dim_len, dim_chunk_len) + + # handle wraparound + if wraparound: + wraparound_indices(dim_sel, dim_len) + + # handle out of bounds + if boundscheck: + boundscheck_indices(dim_sel, dim_len) + + # determine which chunk is needed for each selection item + # note: for dense integer selections, the division operation here is the + # bottleneck + dim_sel_chunk = dim_sel // dim_chunk_len + + # determine order of indices + if order == Order.UNKNOWN: + order = Order.check(dim_sel) + order = Order(order) + + if order == Order.INCREASING: + dim_out_sel = None + elif order == Order.DECREASING: + dim_sel = dim_sel[::-1] + # TODO should be possible to do this without creating an arange + dim_out_sel = np.arange(nitems - 1, -1, -1) + else: + # sort indices to group by chunk + dim_out_sel = np.argsort(dim_sel_chunk) + dim_sel = np.take(dim_sel, dim_out_sel) + + # precompute number of selected items for each chunk + chunk_nitems = np.bincount(dim_sel_chunk, minlength=nchunks) + + # find chunks that we need to visit + dim_chunk_ixs = np.nonzero(chunk_nitems)[0] + + # compute offsets into the output array + chunk_nitems_cumsum = np.cumsum(chunk_nitems) + + # store attributes + object.__setattr__(self, "dim_len", dim_len) + object.__setattr__(self, "dim_chunk_len", dim_chunk_len) + object.__setattr__(self, "nchunks", nchunks) + object.__setattr__(self, "nitems", nitems) + object.__setattr__(self, "order", order) + object.__setattr__(self, "dim_sel", dim_sel) + object.__setattr__(self, "dim_out_sel", dim_out_sel) + object.__setattr__(self, "chunk_nitems", chunk_nitems) + object.__setattr__(self, "dim_chunk_ixs", dim_chunk_ixs) + object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum) + + def __iter__(self) -> Iterator[ChunkDimProjection]: + for dim_chunk_ix in self.dim_chunk_ixs: + dim_out_sel: slice | npt.NDArray[np.intp] + # find region in output + if dim_chunk_ix == 0: + start = 0 + else: + start = self.chunk_nitems_cumsum[dim_chunk_ix - 1] + stop = self.chunk_nitems_cumsum[dim_chunk_ix] + if self.order == Order.INCREASING: + dim_out_sel = slice(start, stop) + else: + dim_out_sel = self.dim_out_sel[start:stop] + + # find region in chunk + dim_offset = dim_chunk_ix * self.dim_chunk_len + dim_chunk_sel = self.dim_sel[start:stop] - dim_offset + is_complete_chunk = False # TODO + yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) + + +def slice_to_range(s: slice, length: int) -> range: + return range(*s.indices(length)) + + +def ix_(selection: Any, shape: tuple[int, ...]) -> npt.NDArray[np.intp]: + """Convert an orthogonal selection to a numpy advanced (fancy) selection, like ``numpy.ix_`` + but with support for slices and single ints.""" + + # normalisation + selection = replace_ellipsis(selection, shape) + + # replace slice and int as these are not supported by numpy.ix_ + selection = [ + slice_to_range(dim_sel, dim_len) + if isinstance(dim_sel, slice) + else [dim_sel] + if is_integer(dim_sel) + else dim_sel + for dim_sel, dim_len in zip(selection, shape, strict=True) + ] + + # now get numpy to convert to a coordinate selection + selection = np.ix_(*selection) + + return cast("npt.NDArray[np.intp]", selection) + + +def oindex(a: npt.NDArray[Any], selection: Selection) -> npt.NDArray[Any]: + """Implementation of orthogonal indexing with slices and ints.""" + selection = replace_ellipsis(selection, a.shape) + drop_axes = tuple(i for i, s in enumerate(selection) if is_integer(s)) + selection = ix_(selection, a.shape) + result = a[selection] + if drop_axes: + result = result.squeeze(axis=drop_axes) + return result + + +def oindex_set(a: npt.NDArray[Any], selection: Selection, value: Any) -> None: + selection = replace_ellipsis(selection, a.shape) + drop_axes = tuple(i for i, s in enumerate(selection) if is_integer(s)) + selection = ix_(selection, a.shape) + if not np.isscalar(value) and drop_axes: + value = np.asanyarray(value) + value_selection: list[Selector | None] = [slice(None)] * len(a.shape) + for i in drop_axes: + value_selection[i] = np.newaxis + value = value[tuple(value_selection)] + a[selection] = value + + +@dataclass(frozen=True) +class OrthogonalIndexer(Indexer): + dim_indexers: list[IntDimIndexer | SliceDimIndexer | IntArrayDimIndexer | BoolArrayDimIndexer] + shape: tuple[int, ...] + chunk_shape: tuple[int, ...] + is_advanced: bool + drop_axes: tuple[int, ...] + + def __init__(self, selection: Selection, shape: tuple[int, ...], chunk_grid: ChunkGrid) -> None: + chunk_shape = get_chunk_shape(chunk_grid) + + # handle ellipsis + selection = replace_ellipsis(selection, shape) + + # normalize list to array + selection = replace_lists(selection) + + # setup per-dimension indexers + dim_indexers: list[ + IntDimIndexer | SliceDimIndexer | IntArrayDimIndexer | BoolArrayDimIndexer + ] = [] + for dim_sel, dim_len, dim_chunk_len in zip(selection, shape, chunk_shape, strict=True): + dim_indexer: IntDimIndexer | SliceDimIndexer | IntArrayDimIndexer | BoolArrayDimIndexer + if is_integer(dim_sel): + dim_indexer = IntDimIndexer(dim_sel, dim_len, dim_chunk_len) + + elif isinstance(dim_sel, slice): + dim_indexer = SliceDimIndexer(dim_sel, dim_len, dim_chunk_len) + + elif is_integer_array(dim_sel): + dim_indexer = IntArrayDimIndexer(dim_sel, dim_len, dim_chunk_len) + + elif is_bool_array(dim_sel): + dim_indexer = BoolArrayDimIndexer(dim_sel, dim_len, dim_chunk_len) + + else: + raise IndexError( + "unsupported selection item for orthogonal indexing; " + "expected integer, slice, integer array or Boolean " + f"array, got {type(dim_sel)!r}" + ) + + dim_indexers.append(dim_indexer) + + shape = tuple(s.nitems for s in dim_indexers if not isinstance(s, IntDimIndexer)) + is_advanced = not is_basic_selection(selection) + if is_advanced: + drop_axes = tuple( + i + for i, dim_indexer in enumerate(dim_indexers) + if isinstance(dim_indexer, IntDimIndexer) + ) + else: + drop_axes = () + + object.__setattr__(self, "dim_indexers", dim_indexers) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "chunk_shape", chunk_shape) + object.__setattr__(self, "is_advanced", is_advanced) + object.__setattr__(self, "drop_axes", drop_axes) + + def __iter__(self) -> Iterator[ChunkProjection]: + for dim_projections in itertools.product(*self.dim_indexers): + chunk_coords = tuple(p.dim_chunk_ix for p in dim_projections) + chunk_selection: tuple[Selector, ...] | npt.NDArray[Any] = tuple( + p.dim_chunk_sel for p in dim_projections + ) + out_selection: tuple[Selector, ...] | npt.NDArray[Any] = tuple( + p.dim_out_sel for p in dim_projections if p.dim_out_sel is not None + ) + + # handle advanced indexing arrays orthogonally + if self.is_advanced: + # N.B., numpy doesn't support orthogonal indexing directly as yet, + # so need to work around via np.ix_. Also np.ix_ does not support a + # mixture of arrays and slices or integers, so need to convert slices + # and integers into ranges. + chunk_selection = ix_(chunk_selection, self.chunk_shape) + + # special case for non-monotonic indices + if not is_basic_selection(out_selection): + out_selection = ix_(out_selection, self.shape) + + is_complete_chunk = all(p.is_complete_chunk for p in dim_projections) + yield ChunkProjection(chunk_coords, chunk_selection, out_selection, is_complete_chunk) + + +@dataclass(frozen=True) +class OIndex: + array: AnyArray + + # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool + def __getitem__(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: + from zarr.core.array import Array + + # if input is a Zarr array, we materialize it now. + if isinstance(selection, Array): + selection = _zarr_array_to_int_or_bool_array(selection) + + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + return self.array.get_orthogonal_selection( + cast("OrthogonalSelection", new_selection), fields=fields + ) + + def __setitem__(self, selection: OrthogonalSelection, value: npt.ArrayLike) -> None: + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + return self.array.set_orthogonal_selection( + cast("OrthogonalSelection", new_selection), value, fields=fields + ) + + +@dataclass(frozen=True) +class AsyncOIndex(Generic[T_ArrayMetadata]): + array: AsyncArray[T_ArrayMetadata] + + async def getitem(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: + from zarr.core.array import Array + + # if input is a Zarr array, we materialize it now. + if isinstance(selection, Array): + selection = _zarr_array_to_int_or_bool_array(selection) + + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + return await self.array.get_orthogonal_selection( + cast(OrthogonalSelection, new_selection), fields=fields + ) + + +@dataclass(frozen=True) +class BlockIndexer(Indexer): + dim_indexers: list[SliceDimIndexer] + shape: tuple[int, ...] + drop_axes: tuple[int, ...] + + def __init__( + self, selection: BasicSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid + ) -> None: + chunk_shape = get_chunk_shape(chunk_grid) + + # handle ellipsis + selection_normalized = replace_ellipsis(selection, shape) + + # normalize list to array + selection_normalized = replace_lists(selection_normalized) + + # setup per-dimension indexers + dim_indexers = [] + for dim_sel, dim_len, dim_chunk_size in zip( + selection_normalized, shape, chunk_shape, strict=True + ): + dim_numchunks = int(np.ceil(dim_len / dim_chunk_size)) + + if is_integer(dim_sel): + if dim_sel < 0: + dim_sel = dim_numchunks + dim_sel + + start = dim_sel * dim_chunk_size + stop = start + dim_chunk_size + slice_ = slice(start, stop) + + elif is_slice(dim_sel): + start = dim_sel.start if dim_sel.start is not None else 0 + stop = dim_sel.stop if dim_sel.stop is not None else dim_numchunks + + if dim_sel.step not in {1, None}: + raise IndexError( + "unsupported selection item for block indexing; " + f"expected integer or slice with step=1, got {type(dim_sel)!r}" + ) + + # Can't reuse wraparound_indices because it expects a numpy array + # We have integers here. + if start < 0: + start = dim_numchunks + start + if stop < 0: + stop = dim_numchunks + stop + + start *= dim_chunk_size + stop *= dim_chunk_size + slice_ = slice(start, stop) + + else: + raise IndexError( + "unsupported selection item for block indexing; " + f"expected integer or slice, got {type(dim_sel)!r}" + ) + + dim_indexer = SliceDimIndexer(slice_, dim_len, dim_chunk_size) + dim_indexers.append(dim_indexer) + + if start >= dim_len or start < 0: + msg = f"index out of bounds for dimension with length {dim_len}" + raise BoundsCheckError(msg) + + shape = tuple(s.nitems for s in dim_indexers) + + object.__setattr__(self, "dim_indexers", dim_indexers) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "drop_axes", ()) + + def __iter__(self) -> Iterator[ChunkProjection]: + for dim_projections in itertools.product(*self.dim_indexers): + chunk_coords = tuple(p.dim_chunk_ix for p in dim_projections) + chunk_selection = tuple(p.dim_chunk_sel for p in dim_projections) + out_selection = tuple( + p.dim_out_sel for p in dim_projections if p.dim_out_sel is not None + ) + is_complete_chunk = all(p.is_complete_chunk for p in dim_projections) + yield ChunkProjection(chunk_coords, chunk_selection, out_selection, is_complete_chunk) + + +@dataclass(frozen=True) +class BlockIndex: + array: AnyArray + + def __getitem__(self, selection: BasicSelection) -> NDArrayLikeOrScalar: + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + return self.array.get_block_selection(cast("BasicSelection", new_selection), fields=fields) + + def __setitem__(self, selection: BasicSelection, value: npt.ArrayLike) -> None: + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + return self.array.set_block_selection( + cast("BasicSelection", new_selection), value, fields=fields + ) + + +def is_coordinate_selection( + selection: SelectionNormalized, shape: tuple[int, ...] +) -> TypeGuard[CoordinateSelectionNormalized]: + return ( + isinstance(selection, tuple) + and len(selection) == len(shape) + and all(is_integer(dim_sel) or is_integer_array(dim_sel) for dim_sel in selection) + ) + + +def is_mask_selection(selection: Selection, shape: tuple[int, ...]) -> TypeGuard[MaskSelection]: + return ( + isinstance(selection, tuple) + and len(selection) == 1 + and is_bool_array(selection[0]) + and selection[0].shape == shape + ) + + +@dataclass(frozen=True) +class CoordinateIndexer(Indexer): + sel_shape: tuple[int, ...] + selection: CoordinateSelectionNormalized + sel_sort: npt.NDArray[np.intp] | None + chunk_nitems_cumsum: npt.NDArray[np.intp] + chunk_rixs: npt.NDArray[np.intp] + chunk_mixs: tuple[npt.NDArray[np.intp], ...] + shape: tuple[int, ...] + chunk_shape: tuple[int, ...] + drop_axes: tuple[int, ...] + + def __init__( + self, selection: CoordinateSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid + ) -> None: + chunk_shape = get_chunk_shape(chunk_grid) + + cdata_shape: tuple[int, ...] + if shape == (): + cdata_shape = (1,) + else: + cdata_shape = tuple(math.ceil(s / c) for s, c in zip(shape, chunk_shape, strict=True)) + nchunks = reduce(operator.mul, cdata_shape, 1) + + # some initial normalization + selection_normalized = cast("CoordinateSelectionNormalized", ensure_tuple(selection)) + selection_normalized = tuple( + np.asarray([i]) if is_integer(i) else i for i in selection_normalized + ) + selection_normalized = cast( + "CoordinateSelectionNormalized", replace_lists(selection_normalized) + ) + + # validation + if not is_coordinate_selection(selection_normalized, shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + + # handle wraparound, boundscheck + for dim_sel, dim_len in zip(selection_normalized, shape, strict=True): + # handle wraparound + wraparound_indices(dim_sel, dim_len) + + # handle out of bounds + boundscheck_indices(dim_sel, dim_len) + + # compute chunk index for each point in the selection + chunks_multi_index = tuple( + dim_sel // dim_chunk_len + for (dim_sel, dim_chunk_len) in zip(selection_normalized, chunk_shape, strict=True) + ) + + # broadcast selection - this will raise error if array dimensions don't match + selection_broadcast = tuple(np.broadcast_arrays(*selection_normalized)) + chunks_multi_index_broadcast = np.broadcast_arrays(*chunks_multi_index) + + # remember shape of selection, because we will flatten indices for processing + sel_shape = selection_broadcast[0].shape or (1,) + + # flatten selection + selection_broadcast = tuple(dim_sel.reshape(-1) for dim_sel in selection_broadcast) + chunks_multi_index_broadcast = tuple( + dim_chunks.reshape(-1) for dim_chunks in chunks_multi_index_broadcast + ) + + # ravel chunk indices + chunks_raveled_indices = np.ravel_multi_index( + chunks_multi_index_broadcast, dims=cdata_shape + ) + + # group points by chunk + if np.any(np.diff(chunks_raveled_indices) < 0): + # optimisation, only sort if needed + sel_sort = np.argsort(chunks_raveled_indices) + selection_broadcast = tuple(dim_sel[sel_sort] for dim_sel in selection_broadcast) + else: + sel_sort = None + + shape = selection_broadcast[0].shape or (1,) + + # precompute number of selected items for each chunk + chunk_nitems = np.bincount(chunks_raveled_indices, minlength=nchunks) + chunk_nitems_cumsum = np.cumsum(chunk_nitems) + # locate the chunks we need to process + chunk_rixs = np.nonzero(chunk_nitems)[0] + + # unravel chunk indices + chunk_mixs = np.unravel_index(chunk_rixs, cdata_shape) + + object.__setattr__(self, "sel_shape", sel_shape) + object.__setattr__(self, "selection", selection_broadcast) + object.__setattr__(self, "sel_sort", sel_sort) + object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum) + object.__setattr__(self, "chunk_rixs", chunk_rixs) + object.__setattr__(self, "chunk_mixs", chunk_mixs) + object.__setattr__(self, "chunk_shape", chunk_shape) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "drop_axes", ()) + + def __iter__(self) -> Iterator[ChunkProjection]: + # iterate over chunks + for i, chunk_rix in enumerate(self.chunk_rixs): + chunk_coords = tuple(m[i] for m in self.chunk_mixs) + if chunk_rix == 0: + start = 0 + else: + start = self.chunk_nitems_cumsum[chunk_rix - 1] + stop = self.chunk_nitems_cumsum[chunk_rix] + out_selection: slice | npt.NDArray[np.intp] + if self.sel_sort is None: + out_selection = slice(start, stop) + else: + out_selection = self.sel_sort[start:stop] + + chunk_offsets = tuple( + dim_chunk_ix * dim_chunk_len + for dim_chunk_ix, dim_chunk_len in zip(chunk_coords, self.chunk_shape, strict=True) + ) + chunk_selection = tuple( + dim_sel[start:stop] - dim_chunk_offset + for (dim_sel, dim_chunk_offset) in zip(self.selection, chunk_offsets, strict=True) + ) + + is_complete_chunk = False # TODO + yield ChunkProjection(chunk_coords, chunk_selection, out_selection, is_complete_chunk) + + +@dataclass(frozen=True) +class MaskIndexer(CoordinateIndexer): + def __init__( + self, selection: MaskSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid + ) -> None: + # some initial normalization + selection_normalized = cast("tuple[MaskSelection]", ensure_tuple(selection)) + selection_normalized = cast("tuple[MaskSelection]", replace_lists(selection_normalized)) + + # validation + if not is_mask_selection(selection_normalized, shape): + raise IndexError( + "invalid mask selection; expected one Boolean (mask)" + f"array with the same shape as the target array, got {selection_normalized!r}" + ) + + # convert to indices + selection_indices = np.nonzero(selection_normalized[0]) + + # delegate the rest to superclass + super().__init__(selection_indices, shape, chunk_grid) + + +@dataclass(frozen=True) +class VIndex: + array: AnyArray + + # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool + def __getitem__( + self, selection: CoordinateSelection | MaskSelection | AnyArray + ) -> NDArrayLikeOrScalar: + from zarr.core.array import Array + + # if input is a Zarr array, we materialize it now. + if isinstance(selection, Array): + selection = _zarr_array_to_int_or_bool_array(selection) + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + if is_coordinate_selection(new_selection, self.array.shape): + return self.array.get_coordinate_selection(new_selection, fields=fields) + elif is_mask_selection(new_selection, self.array.shape): + return self.array.get_mask_selection(new_selection, fields=fields) + else: + msg = ( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {new_selection!r}" + ) + raise VindexInvalidSelectionError(msg) + + def __setitem__( + self, selection: CoordinateSelection | MaskSelection, value: npt.ArrayLike + ) -> None: + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + if is_coordinate_selection(new_selection, self.array.shape): + self.array.set_coordinate_selection(new_selection, value, fields=fields) + elif is_mask_selection(new_selection, self.array.shape): + self.array.set_mask_selection(new_selection, value, fields=fields) + else: + msg = ( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {new_selection!r}" + ) + raise VindexInvalidSelectionError(msg) + + +@dataclass(frozen=True) +class AsyncVIndex(Generic[T_ArrayMetadata]): + array: AsyncArray[T_ArrayMetadata] + + # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool + async def getitem( + self, selection: CoordinateSelection | MaskSelection | AnyArray + ) -> NDArrayLikeOrScalar: + # TODO deduplicate these internals with the sync version of getitem + # TODO requires solving this circular sync issue: https://github.com/zarr-developers/zarr-python/pull/3083#discussion_r2230737448 + from zarr.core.array import Array + + # if input is a Zarr array, we materialize it now. + if isinstance(selection, Array): + selection = _zarr_array_to_int_or_bool_array(selection) + fields, new_selection = pop_fields(selection) + new_selection = ensure_tuple(new_selection) + new_selection = replace_lists(new_selection) + if is_coordinate_selection(new_selection, self.array.shape): + return await self.array.get_coordinate_selection(new_selection, fields=fields) + elif is_mask_selection(new_selection, self.array.shape): + return await self.array.get_mask_selection(new_selection, fields=fields) + else: + msg = ( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {new_selection!r}" + ) + raise VindexInvalidSelectionError(msg) + + +def check_fields(fields: Fields | None, dtype: np.dtype[Any]) -> np.dtype[Any]: + # early out + if fields is None: + return dtype + # check type + if not isinstance(fields, str | list | tuple): + raise IndexError( + f"'fields' argument must be a string or list of strings; found {type(fields)!r}" + ) + if fields: + if dtype.names is None: + raise IndexError("invalid 'fields' argument, array does not have any fields") + try: + if isinstance(fields, str): + # single field selection + out_dtype = dtype[fields] + else: + # multiple field selection + out_dtype = np.dtype([(f, dtype[f]) for f in fields]) + except KeyError as e: + raise IndexError(f"invalid 'fields' argument, field not found: {e!r}") from e + else: + return out_dtype + else: + return dtype + + +def check_no_multi_fields(fields: Fields | None) -> Fields | None: + if isinstance(fields, list): + if len(fields) == 1: + return fields[0] + elif len(fields) > 1: + raise IndexError("multiple fields are not supported for this operation") + return fields + + +def pop_fields(selection: SelectionWithFields) -> tuple[Fields | None, Selection]: + if isinstance(selection, str): + # single field selection + return selection, () + elif not isinstance(selection, tuple): + # single selection item, no fields + # leave selection as-is + return None, cast("Selection", selection) + else: + # multiple items, split fields from selection items + fields: Fields = [f for f in selection if isinstance(f, str)] + fields = fields[0] if len(fields) == 1 else fields + selection_tuple = tuple(s for s in selection if not isinstance(s, str)) + selection = cast( + "Selection", selection_tuple[0] if len(selection_tuple) == 1 else selection_tuple + ) + return fields, selection + + +def make_slice_selection(selection: Any) -> list[slice]: + ls: list[slice] = [] + for dim_selection in selection: + if is_integer(dim_selection): + ls.append(slice(int(dim_selection), int(dim_selection) + 1, 1)) + elif isinstance(dim_selection, np.ndarray): + if len(dim_selection) == 1: + ls.append(slice(int(dim_selection[0]), int(dim_selection[0]) + 1, 1)) + else: + raise ArrayIndexError + else: + ls.append(dim_selection) + return ls + + +def decode_morton(z: int, chunk_shape: tuple[int, ...]) -> tuple[int, ...]: + # Inspired by compressed morton code as implemented in Neuroglancer + # https://github.com/google/neuroglancer/blob/master/src/neuroglancer/datasource/precomputed/volume.md#compressed-morton-code + bits = tuple((c - 1).bit_length() for c in chunk_shape) + max_coords_bits = max(bits) + input_bit = 0 + input_value = z + out = [0] * len(chunk_shape) + + for coord_bit in range(max_coords_bits): + for dim in range(len(chunk_shape)): + if coord_bit < bits[dim]: + bit = (input_value >> input_bit) & 1 + out[dim] |= bit << coord_bit + input_bit += 1 + return tuple(out) + + +def decode_morton_vectorized( + z: npt.NDArray[np.intp], chunk_shape: tuple[int, ...] +) -> npt.NDArray[np.intp]: + """Vectorized Morton code decoding for multiple z values. + + Parameters + ---------- + z : ndarray + 1D array of Morton codes to decode. + chunk_shape : tuple of int + Shape defining the coordinate space. + + Returns + ------- + ndarray + 2D array of shape (len(z), len(chunk_shape)) containing decoded coordinates. + """ + n_dims = len(chunk_shape) + bits = tuple((c - 1).bit_length() for c in chunk_shape) + + max_coords_bits = max(bits) if bits else 0 + out = np.zeros((len(z), n_dims), dtype=np.intp) + + input_bit = 0 + for coord_bit in range(max_coords_bits): + for dim in range(n_dims): + if coord_bit < bits[dim]: + # Extract bit at position input_bit from all z values + bit_values = (z >> input_bit) & 1 + # Place bit at coord_bit position in dimension dim + out[:, dim] |= bit_values << coord_bit + input_bit += 1 + + return out + + +@lru_cache(maxsize=16) +def _morton_order(chunk_shape: tuple[int, ...]) -> npt.NDArray[np.intp]: + n_total = product(chunk_shape) + n_dims = len(chunk_shape) + if n_total == 0: + out = np.empty((0, n_dims), dtype=np.intp) + out.flags.writeable = False + return out + + # Ceiling hypercube: smallest power-of-2 hypercube whose Morton codes span + # all valid coordinates in chunk_shape. (c-1).bit_length() gives the number + # of bits needed to index c values (0 for singleton dims). n_z = 2**total_bits + # is the size of this hypercube. + total_bits = sum((c - 1).bit_length() for c in chunk_shape) + n_z = 1 << total_bits if total_bits > 0 else 1 + + # Decode all Morton codes in the ceiling hypercube, then filter to valid coords. + # This is fully vectorized. For shapes with n_z >> n_total (e.g. (33,33,33): + # n_z=262144, n_total=35937), consider the argsort strategy below. + order: npt.NDArray[np.intp] + if n_z <= 4 * n_total: + # Ceiling strategy: decode all n_z codes vectorized, filter in-bounds. + # Works well when the overgeneration ratio n_z/n_total is small (≤4). + z_values = np.arange(n_z, dtype=np.intp) + all_coords = decode_morton_vectorized(z_values, chunk_shape) + shape_arr = np.array(chunk_shape, dtype=np.intp) + valid_mask = np.all(all_coords < shape_arr, axis=1) + order = all_coords[valid_mask] + else: + # Argsort strategy: enumerate all n_total valid coordinates directly, + # encode each to a Morton code, then sort by code. Avoids the 8x or + # larger overgeneration penalty for near-miss shapes like (33,33,33). + # Cost: O(n_total * bits) encode + O(n_total log n_total) sort, + # vs O(n_z * bits) = O(8 * n_total * bits) for ceiling. + grids = np.meshgrid(*[np.arange(c, dtype=np.intp) for c in chunk_shape], indexing="ij") + all_coords = np.stack([g.ravel() for g in grids], axis=1) + + # Encode all coordinates to Morton codes (vectorized). + bits_per_dim = tuple((c - 1).bit_length() for c in chunk_shape) + max_coord_bits = max(bits_per_dim) + z_codes = np.zeros(n_total, dtype=np.intp) + output_bit = 0 + for coord_bit in range(max_coord_bits): + for dim in range(n_dims): + if coord_bit < bits_per_dim[dim]: + z_codes |= ((all_coords[:, dim] >> coord_bit) & 1) << output_bit + output_bit += 1 + + sort_idx: npt.NDArray[np.intp] = np.argsort(z_codes, kind="stable") + order = np.asarray(all_coords[sort_idx], dtype=np.intp) + + order.flags.writeable = False + return order + + +@lru_cache(maxsize=16) +def _morton_order_keys(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]: + return tuple(tuple(int(x) for x in row) for row in _morton_order(chunk_shape)) + + +def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: + return iter(_morton_order_keys(tuple(chunk_shape))) + + +def c_order_iter(chunks_per_shard: tuple[int, ...]) -> Iterator[tuple[int, ...]]: + return itertools.product(*(range(x) for x in chunks_per_shard)) + + +def get_indexer( + selection: SelectionWithFields, shape: tuple[int, ...], chunk_grid: ChunkGrid +) -> Indexer: + _, pure_selection = pop_fields(selection) + if is_pure_fancy_indexing(pure_selection, len(shape)): + new_selection = ensure_tuple(selection) + new_selection = replace_lists(new_selection) + if is_coordinate_selection(new_selection, shape): + return CoordinateIndexer(cast("CoordinateSelection", selection), shape, chunk_grid) + elif is_mask_selection(new_selection, shape): + return MaskIndexer(cast("MaskSelection", selection), shape, chunk_grid) + else: + msg = ( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {new_selection!r}" + ) + raise VindexInvalidSelectionError(msg) + elif is_pure_orthogonal_indexing(pure_selection, len(shape)): + return OrthogonalIndexer(cast("OrthogonalSelection", selection), shape, chunk_grid) + else: + return BasicIndexer(cast("BasicSelection", selection), shape, chunk_grid) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__init__.py b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..57385386b6b9d1807af718d20565e139567a1736 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__init__.py @@ -0,0 +1,17 @@ +from typing import TypeAlias, TypeVar + +from .v2 import ArrayV2Metadata, ArrayV2MetadataDict +from .v3 import ArrayMetadataJSON_V3, ArrayV3Metadata + +ArrayMetadata: TypeAlias = ArrayV2Metadata | ArrayV3Metadata +ArrayMetadataDict: TypeAlias = ArrayV2MetadataDict | ArrayMetadataJSON_V3 +T_ArrayMetadata = TypeVar("T_ArrayMetadata", ArrayV2Metadata, ArrayV3Metadata, covariant=True) + +__all__ = [ + "ArrayMetadata", + "ArrayMetadataDict", + "ArrayMetadataJSON_V3", + "ArrayV2Metadata", + "ArrayV2MetadataDict", + "ArrayV3Metadata", +] diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00cbb61dac792c2aa6a7c6e9d3fd12e6f3d96f87 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/common.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/common.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f3e98bb1ac49b0a911f652626d33ee7c9aca743 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/common.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/io.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/io.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6f99f21e4d4f1b39cb0f2fecf483ac331edf2e5 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/io.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v2.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v2.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a35419f9e5f38bdec3a23b94abc809c0f51d83a Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v2.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v3.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v3.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1640eb2a7632fae082fbb873d6395c7ff45249ae Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/core/metadata/__pycache__/v3.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/common.py b/.venv/lib/python3.11/site-packages/zarr/core/metadata/common.py new file mode 100644 index 0000000000000000000000000000000000000000..44d3eb292bbab5c9f2b63695ffcd113c0d0ea508 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/metadata/common.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from zarr.core.common import JSON + + +def parse_attributes(data: dict[str, JSON] | None) -> dict[str, JSON]: + if data is None: + return {} + + return data diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/io.py b/.venv/lib/python3.11/site-packages/zarr/core/metadata/io.py new file mode 100644 index 0000000000000000000000000000000000000000..7b63f5493b36b43ec2dd483fddd8988db1add8dd --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/metadata/io.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from zarr.abc.store import set_or_delete +from zarr.core.buffer.core import default_buffer_prototype +from zarr.errors import ContainsArrayError +from zarr.storage._common import StorePath, ensure_no_existing_node + +if TYPE_CHECKING: + from zarr.core.common import ZarrFormat + from zarr.core.group import GroupMetadata + from zarr.core.metadata import ArrayMetadata + + +def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]: + from zarr.core.group import GroupMetadata + + path = store_path.path + if not path: + return {} + + required_parts = path.split("/")[:-1] + + # the root group + parents = {"": GroupMetadata(zarr_format=zarr_format)} + + for i, part in enumerate(required_parts): + parent_path = "/".join(required_parts[:i] + [part]) + parents[parent_path] = GroupMetadata(zarr_format=zarr_format) + + return parents + + +async def save_metadata( + store_path: StorePath, metadata: ArrayMetadata | GroupMetadata, ensure_parents: bool = False +) -> None: + """Asynchronously save the array or group metadata. + + Parameters + ---------- + store_path : StorePath + Location to save metadata. + metadata : ArrayMetadata | GroupMetadata + Metadata to save. + ensure_parents : bool, optional + Create any missing parent groups, and check no existing parents are arrays. + + Raises + ------ + ValueError + """ + to_save = metadata.to_buffer_dict(default_buffer_prototype()) + set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()] + + if ensure_parents: + # To enable zarr.create(store, path="a/b/c"), we need to create all the intermediate groups. + parents = _build_parents(store_path, metadata.zarr_format) + ensure_array_awaitables = [] + + for parent_path, parent_metadata in parents.items(): + parent_store_path = StorePath(store_path.store, parent_path) + + # Error if an array already exists at any parent location. Only groups can have child nodes. + ensure_array_awaitables.append( + ensure_no_existing_node( + parent_store_path, parent_metadata.zarr_format, node_type="array" + ) + ) + set_awaitables.extend( + [ + (parent_store_path / key).set_if_not_exists(value) + for key, value in parent_metadata.to_buffer_dict( + default_buffer_prototype() + ).items() + ] + ) + + # Checks for parent arrays must happen first, before any metadata is modified + try: + await asyncio.gather(*ensure_array_awaitables) + except ContainsArrayError as e: + # clear awaitables to avoid RuntimeWarning: coroutine was never awaited + for awaitable in set_awaitables: + awaitable.close() + + raise ValueError( + f"A parent of {store_path} is an array - only groups may have child nodes." + ) from e + + await asyncio.gather(*set_awaitables) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/v2.py b/.venv/lib/python3.11/site-packages/zarr/core/metadata/v2.py new file mode 100644 index 0000000000000000000000000000000000000000..32045434267d63626e3541ce1a41fc0fed5262ff --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/metadata/v2.py @@ -0,0 +1,330 @@ +from __future__ import annotations + +import warnings +from collections.abc import Iterable, Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict, cast + +from zarr.abc.metadata import Metadata +from zarr.abc.numcodec import Numcodec, _is_numcodec +from zarr.core.chunk_grids import RegularChunkGrid +from zarr.core.dtype import get_data_type_from_json +from zarr.core.dtype.common import OBJECT_CODEC_IDS, DTypeSpec_V2 +from zarr.errors import ZarrUserWarning +from zarr.registry import get_numcodec + +if TYPE_CHECKING: + from typing import Literal, Self + + import numpy.typing as npt + + from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.dtype.wrapper import ( + TBaseDType, + TBaseScalar, + TDType_co, + TScalar_co, + ZDType, + ) + +import json +from dataclasses import dataclass, field, fields, replace + +import numpy as np + +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.chunk_key_encodings import parse_separator +from zarr.core.common import ( + JSON, + ZARRAY_JSON, + ZATTRS_JSON, + MemoryOrder, + parse_shapelike, +) +from zarr.core.config import config, parse_indexing_order +from zarr.core.metadata.common import parse_attributes + + +class ArrayV2MetadataDict(TypedDict): + """ + A typed dictionary model for Zarr format 2 metadata. + """ + + zarr_format: Literal[2] + attributes: dict[str, JSON] + + +# Union of acceptable types for v2 compressors +CompressorLikev2: TypeAlias = dict[str, JSON] | Numcodec | None + + +@dataclass(frozen=True, kw_only=True) +class ArrayV2Metadata(Metadata): + shape: tuple[int, ...] + chunks: tuple[int, ...] + dtype: ZDType[TBaseDType, TBaseScalar] + fill_value: int | float | str | bytes | None = None + order: MemoryOrder = "C" + filters: tuple[Numcodec, ...] | None = None + dimension_separator: Literal[".", "/"] = "." + compressor: Numcodec | None + attributes: dict[str, JSON] = field(default_factory=dict) + zarr_format: Literal[2] = field(init=False, default=2) + + def __init__( + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TDType_co, TScalar_co], + chunks: tuple[int, ...], + fill_value: Any, + order: MemoryOrder, + dimension_separator: Literal[".", "/"] = ".", + compressor: CompressorLikev2 = None, + filters: Iterable[Numcodec | dict[str, JSON]] | None = None, + attributes: dict[str, JSON] | None = None, + ) -> None: + """ + Metadata for a Zarr format 2 array. + """ + shape_parsed = parse_shapelike(shape) + chunks_parsed = parse_shapelike(chunks) + compressor_parsed = parse_compressor(compressor) + order_parsed = parse_indexing_order(order) + dimension_separator_parsed = parse_separator(dimension_separator) + filters_parsed = parse_filters(filters) + fill_value_parsed: TBaseScalar | None + if fill_value is not None: + fill_value_parsed = dtype.cast_scalar(fill_value) + else: + fill_value_parsed = fill_value + attributes_parsed = parse_attributes(attributes) + + object.__setattr__(self, "shape", shape_parsed) + object.__setattr__(self, "dtype", dtype) + object.__setattr__(self, "chunks", chunks_parsed) + object.__setattr__(self, "compressor", compressor_parsed) + object.__setattr__(self, "order", order_parsed) + object.__setattr__(self, "dimension_separator", dimension_separator_parsed) + object.__setattr__(self, "filters", filters_parsed) + object.__setattr__(self, "fill_value", fill_value_parsed) + object.__setattr__(self, "attributes", attributes_parsed) + + # ensure that the metadata document is consistent + _ = parse_metadata(self) + + @property + def ndim(self) -> int: + return len(self.shape) + + @cached_property + def chunk_grid(self) -> RegularChunkGrid: + return RegularChunkGrid(chunk_shape=self.chunks) + + @property + def shards(self) -> tuple[int, ...] | None: + return None + + def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + zarray_dict = self.to_dict() + zattrs_dict = zarray_dict.pop("attributes", {}) + json_indent = config.get("json_indent") + return { + ZARRAY_JSON: prototype.buffer.from_bytes( + json.dumps(zarray_dict, indent=json_indent, allow_nan=True).encode() + ), + ZATTRS_JSON: prototype.buffer.from_bytes( + json.dumps(zattrs_dict, indent=json_indent, allow_nan=True).encode() + ), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ArrayV2Metadata: + # Make a copy to protect the original from modification. + _data = data.copy() + # Check that the zarr_format attribute is correct. + _ = parse_zarr_format(_data.pop("zarr_format")) + + # To resolve a numpy object dtype array, we need to search for an object codec, + # which could be in filters or as a compressor. + # we will reference a hard-coded collection of object codec ids for this search. + + _filters, _compressor = (data.get("filters"), data.get("compressor")) + if _filters is not None: + _filters = cast("tuple[dict[str, JSON], ...]", _filters) + object_codec_id = get_object_codec_id(tuple(_filters) + (_compressor,)) + else: + object_codec_id = get_object_codec_id((_compressor,)) + # we add a layer of indirection here around the dtype attribute of the array metadata + # because we also need to know the object codec id, if any, to resolve the data type + dtype_spec: DTypeSpec_V2 = { + "name": data["dtype"], + "object_codec_id": object_codec_id, + } + dtype = get_data_type_from_json(dtype_spec, zarr_format=2) + + _data["dtype"] = dtype + fill_value_encoded = _data.get("fill_value") + if fill_value_encoded is not None: + fill_value = dtype.from_json_scalar(fill_value_encoded, zarr_format=2) + _data["fill_value"] = fill_value + + # zarr v2 allowed arbitrary keys here. + # We don't want the ArrayV2Metadata constructor to fail just because someone put an + # extra key in the metadata. + expected = {x.name for x in fields(cls)} + expected |= {"dtype", "chunks"} + + # check if `filters` is an empty sequence; if so use None instead and raise a warning + filters = _data.get("filters") + if ( + isinstance(filters, Sequence) + and not isinstance(filters, (str, bytes)) + and len(filters) == 0 + ): + msg = ( + "Found an empty list of filters in the array metadata document. " + "This is contrary to the Zarr V2 specification, and will cause an error in the future. " + "Use None (or Null in a JSON document) instead of an empty list of filters." + ) + warnings.warn(msg, ZarrUserWarning, stacklevel=1) + _data["filters"] = None + + _data = {k: v for k, v in _data.items() if k in expected} + + return cls(**_data) + + def to_dict(self) -> dict[str, JSON]: + zarray_dict = super().to_dict() + if _is_numcodec(zarray_dict["compressor"]): + codec_config = zarray_dict["compressor"].get_config() + # Hotfix for https://github.com/zarr-developers/zarr-python/issues/2647 + if codec_config["id"] == "zstd" and not codec_config.get("checksum", False): + codec_config.pop("checksum") + zarray_dict["compressor"] = codec_config + + if zarray_dict["filters"] is not None: + raw_filters = zarray_dict["filters"] + # TODO: remove this when we can stratically type the output JSON data structure + # entirely + if not isinstance(raw_filters, list | tuple): + raise TypeError("Invalid type for filters. Expected a list or tuple.") + new_filters = [] + for f in raw_filters: + if _is_numcodec(f): + new_filters.append(f.get_config()) + else: + new_filters.append(f) + zarray_dict["filters"] = new_filters + + # serialize the fill value after dtype-specific JSON encoding + if self.fill_value is not None: + fill_value = self.dtype.to_json_scalar(self.fill_value, zarr_format=2) + zarray_dict["fill_value"] = fill_value + + # pull the "name" attribute out of the dtype spec returned by self.dtype.to_json + zarray_dict["dtype"] = self.dtype.to_json(zarr_format=2)["name"] + + return zarray_dict + + def get_chunk_spec( + self, _chunk_coords: tuple[int, ...], array_config: ArrayConfig, prototype: BufferPrototype + ) -> ArraySpec: + return ArraySpec( + shape=self.chunks, + dtype=self.dtype, + fill_value=self.fill_value, + config=array_config, + prototype=prototype, + ) + + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + chunk_identifier = self.dimension_separator.join(map(str, chunk_coords)) + return "0" if chunk_identifier == "" else chunk_identifier + + def update_shape(self, shape: tuple[int, ...]) -> Self: + return replace(self, shape=shape) + + def update_attributes(self, attributes: dict[str, JSON]) -> Self: + return replace(self, attributes=attributes) + + +def parse_dtype(data: npt.DTypeLike) -> np.dtype[Any]: + if isinstance(data, list): # this is a valid _VoidDTypeLike check + data = [tuple(d) for d in data] + return np.dtype(data) + + +def parse_zarr_format(data: object) -> Literal[2]: + if data == 2: + return 2 + raise ValueError(f"Invalid value. Expected 2. Got {data}.") + + +def parse_filters(data: object) -> tuple[Numcodec, ...] | None: + """ + Parse a potential tuple of filters + """ + out: list[Numcodec] = [] + + if data is None: + return data + if isinstance(data, Iterable): + for idx, val in enumerate(data): + if _is_numcodec(val): + out.append(val) + elif isinstance(val, dict): + out.append(get_numcodec(val)) # type: ignore[arg-type] + else: + msg = f"Invalid filter at index {idx}. Expected a numcodecs.abc.Codec or a dict representation of numcodecs.abc.Codec. Got {type(val)} instead." + raise TypeError(msg) + if len(out) == 0: + # Per the v2 spec, an empty tuple is not allowed -- use None to express "no filters" + return None + else: + return tuple(out) + # take a single codec instance and wrap it in a tuple + if _is_numcodec(data): + return (data,) + msg = f"Invalid filters. Expected None, an iterable of numcodecs.abc.Codec or dict representations of numcodecs.abc.Codec. Got {type(data)} instead." + raise TypeError(msg) + + +def parse_compressor(data: object) -> Numcodec | None: + """ + Parse a potential compressor. + """ + if data is None or _is_numcodec(data): + return data + if isinstance(data, dict): + return get_numcodec(data) # type: ignore[arg-type] + msg = f"Invalid compressor. Expected None, a numcodecs.abc.Codec, or a dict representation of a numcodecs.abc.Codec. Got {type(data)} instead." + raise ValueError(msg) + + +def parse_metadata(data: ArrayV2Metadata) -> ArrayV2Metadata: + if (l_chunks := len(data.chunks)) != (l_shape := len(data.shape)): + msg = ( + f"The `shape` and `chunks` attributes must have the same length. " + f"`chunks` has length {l_chunks}, but `shape` has length {l_shape}." + ) + raise ValueError(msg) + return data + + +def get_object_codec_id(maybe_object_codecs: Sequence[JSON]) -> str | None: + """ + Inspect a sequence of codecs / filters for an "object codec", i.e. a codec + that can serialize object arrays to contiguous bytes. Zarr python + maintains a hard-coded set of object codec ids. If any element from the input + has an id that matches one of the hard-coded object codec ids, that id + is returned immediately. + """ + object_codec_id = None + for maybe_object_codec in maybe_object_codecs: + if ( + isinstance(maybe_object_codec, dict) + and maybe_object_codec.get("id") in OBJECT_CODEC_IDS + ): + return cast("str", maybe_object_codec["id"]) + return object_codec_id diff --git a/.venv/lib/python3.11/site-packages/zarr/core/metadata/v3.py b/.venv/lib/python3.11/site-packages/zarr/core/metadata/v3.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce155bd9aabf14c3292e9f05ced12e20a927115 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/metadata/v3.py @@ -0,0 +1,442 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, NotRequired, TypedDict, TypeGuard, cast + +from zarr.abc.metadata import Metadata +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.dtype import VariableLengthUTF8, ZDType, get_data_type_from_json +from zarr.core.dtype.common import check_dtype_spec_v3 + +if TYPE_CHECKING: + from typing import Self + + from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.chunk_grids import ChunkGrid + from zarr.core.common import JSON + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar + + +import json +from collections.abc import Iterable +from dataclasses import dataclass, field, replace +from typing import Any, Literal + +from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid +from zarr.core.chunk_key_encodings import ( + ChunkKeyEncoding, + ChunkKeyEncodingLike, + parse_chunk_key_encoding, +) +from zarr.core.common import ( + JSON, + ZARR_JSON, + DimensionNames, + NamedConfig, + parse_named_configuration, + parse_shapelike, +) +from zarr.core.config import config +from zarr.core.metadata.common import parse_attributes +from zarr.errors import MetadataValidationError, NodeTypeValidationError, UnknownCodecError +from zarr.registry import get_codec_class + + +def parse_zarr_format(data: object) -> Literal[3]: + if data == 3: + return 3 + msg = f"Invalid value for 'zarr_format'. Expected '3'. Got '{data}'." + raise MetadataValidationError(msg) + + +def parse_node_type_array(data: object) -> Literal["array"]: + if data == "array": + return "array" + msg = f"Invalid value for 'node_type'. Expected 'array'. Got '{data}'." + raise NodeTypeValidationError(msg) + + +def parse_codecs(data: object) -> tuple[Codec, ...]: + out: tuple[Codec, ...] = () + + if not isinstance(data, Iterable): + raise TypeError(f"Expected iterable, got {type(data)}") + + for c in data: + if isinstance( + c, ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec + ): # Can't use Codec here because of mypy limitation + out += (c,) + else: + name_parsed, _ = parse_named_configuration(c, require_configuration=False) + + try: + out += (get_codec_class(name_parsed).from_dict(c),) + except KeyError as e: + raise UnknownCodecError(f"Unknown codec: {e.args[0]!r}") from e + + return out + + +def validate_array_bytes_codec(codecs: tuple[Codec, ...]) -> ArrayBytesCodec: + # ensure that we have at least one ArrayBytesCodec + abcs: list[ArrayBytesCodec] = [codec for codec in codecs if isinstance(codec, ArrayBytesCodec)] + if len(abcs) == 0: + raise ValueError("At least one ArrayBytesCodec is required.") + elif len(abcs) > 1: + raise ValueError("Only one ArrayBytesCodec is allowed.") + + return abcs[0] + + +def validate_codecs(codecs: tuple[Codec, ...], dtype: ZDType[TBaseDType, TBaseScalar]) -> None: + """Check that the codecs are valid for the given dtype""" + from zarr.codecs.sharding import ShardingCodec + + abc = validate_array_bytes_codec(codecs) + + # Recursively resolve array-bytes codecs within sharding codecs + while isinstance(abc, ShardingCodec): + abc = validate_array_bytes_codec(abc.codecs) + + # we need to have special codecs if we are decoding vlen strings or bytestrings + # TODO: use codec ID instead of class name + codec_class_name = abc.__class__.__name__ + # TODO: Fix typing here + if isinstance(dtype, VariableLengthUTF8) and not codec_class_name == "VLenUTF8Codec": # type: ignore[unreachable] + raise ValueError( + f"For string dtype, ArrayBytesCodec must be `VLenUTF8Codec`, got `{codec_class_name}`." + ) + + +def parse_dimension_names(data: object) -> tuple[str | None, ...] | None: + if data is None: + return data + elif isinstance(data, Iterable) and all(isinstance(x, type(None) | str) for x in data): + return tuple(data) + else: + msg = f"Expected either None or an iterable of str, got {type(data)}" + raise TypeError(msg) + + +def parse_storage_transformers(data: object) -> tuple[dict[str, JSON], ...]: + """ + Parse storage_transformers. Zarr python cannot use storage transformers + at this time, so this function doesn't attempt to validate them. + """ + if data is None: + return () + if isinstance(data, Iterable): + if len(tuple(data)) >= 1: + return data # type: ignore[return-value] + else: + return () + raise TypeError( + f"Invalid storage_transformers. Expected an iterable of dicts. Got {type(data)} instead." + ) + + +class AllowedExtraField(TypedDict): + """ + This class models allowed extra fields in array metadata. + They are ignored by Zarr Python. + """ + + must_understand: Literal[False] + + +def check_allowed_extra_field(data: object) -> TypeGuard[AllowedExtraField]: + """ + Check if the extra field is allowed according to the Zarr v3 spec. The object + must be a mapping with a "must_understand" key set to `False`. + """ + return isinstance(data, Mapping) and data.get("must_understand") is False + + +def parse_extra_fields( + data: Mapping[str, AllowedExtraField] | None, +) -> dict[str, AllowedExtraField]: + if data is None: + return {} + else: + conflict_keys = ARRAY_METADATA_KEYS & set(data.keys()) + if len(conflict_keys) > 0: + msg = ( + "Invalid extra fields. " + "The following keys: " + f"{sorted(conflict_keys)} " + "are invalid because they collide with keys reserved for use by the " + "array metadata document." + ) + raise ValueError(msg) + return dict(data) + + +class ArrayMetadataJSON_V3(TypedDict): + """ + A typed dictionary model for zarr v3 metadata. + """ + + zarr_format: Literal[3] + node_type: Literal["array"] + data_type: str | NamedConfig[str, Mapping[str, object]] + shape: tuple[int, ...] + chunk_grid: NamedConfig[str, Mapping[str, object]] + chunk_key_encoding: NamedConfig[str, Mapping[str, object]] + fill_value: object + codecs: tuple[str | NamedConfig[str, Mapping[str, object]], ...] + attributes: NotRequired[Mapping[str, JSON]] + storage_transformers: NotRequired[tuple[NamedConfig[str, Mapping[str, object]], ...]] + dimension_names: NotRequired[tuple[str | None]] + + +ARRAY_METADATA_KEYS = set(ArrayMetadataJSON_V3.__annotations__.keys()) + + +@dataclass(frozen=True, kw_only=True) +class ArrayV3Metadata(Metadata): + shape: tuple[int, ...] + data_type: ZDType[TBaseDType, TBaseScalar] + chunk_grid: ChunkGrid + chunk_key_encoding: ChunkKeyEncoding + fill_value: Any + codecs: tuple[Codec, ...] + attributes: dict[str, Any] = field(default_factory=dict) + dimension_names: tuple[str | None, ...] | None = None + zarr_format: Literal[3] = field(default=3, init=False) + node_type: Literal["array"] = field(default="array", init=False) + storage_transformers: tuple[dict[str, JSON], ...] + extra_fields: dict[str, AllowedExtraField] + + def __init__( + self, + *, + shape: Iterable[int], + data_type: ZDType[TBaseDType, TBaseScalar], + chunk_grid: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any], + chunk_key_encoding: ChunkKeyEncodingLike, + fill_value: object, + codecs: Iterable[Codec | dict[str, JSON] | NamedConfig[str, Any] | str], + attributes: dict[str, JSON] | None, + dimension_names: DimensionNames, + storage_transformers: Iterable[dict[str, JSON]] | None = None, + extra_fields: Mapping[str, AllowedExtraField] | None = None, + ) -> None: + """ + Because the class is a frozen dataclass, we set attributes using object.__setattr__ + """ + + shape_parsed = parse_shapelike(shape) + chunk_grid_parsed = ChunkGrid.from_dict(chunk_grid) + chunk_key_encoding_parsed = parse_chunk_key_encoding(chunk_key_encoding) + dimension_names_parsed = parse_dimension_names(dimension_names) + # Note: relying on a type method is numpy-specific + fill_value_parsed = data_type.cast_scalar(fill_value) + attributes_parsed = parse_attributes(attributes) + codecs_parsed_partial = parse_codecs(codecs) + storage_transformers_parsed = parse_storage_transformers(storage_transformers) + extra_fields_parsed = parse_extra_fields(extra_fields) + array_spec = ArraySpec( + shape=shape_parsed, + dtype=data_type, + fill_value=fill_value_parsed, + config=ArrayConfig.from_dict({}), # TODO: config is not needed here. + prototype=default_buffer_prototype(), # TODO: prototype is not needed here. + ) + codecs_parsed = tuple(c.evolve_from_array_spec(array_spec) for c in codecs_parsed_partial) + validate_codecs(codecs_parsed_partial, data_type) + + object.__setattr__(self, "shape", shape_parsed) + object.__setattr__(self, "data_type", data_type) + object.__setattr__(self, "chunk_grid", chunk_grid_parsed) + object.__setattr__(self, "chunk_key_encoding", chunk_key_encoding_parsed) + object.__setattr__(self, "codecs", codecs_parsed) + object.__setattr__(self, "dimension_names", dimension_names_parsed) + object.__setattr__(self, "fill_value", fill_value_parsed) + object.__setattr__(self, "attributes", attributes_parsed) + object.__setattr__(self, "storage_transformers", storage_transformers_parsed) + object.__setattr__(self, "extra_fields", extra_fields_parsed) + + self._validate_metadata() + + def _validate_metadata(self) -> None: + if isinstance(self.chunk_grid, RegularChunkGrid) and len(self.shape) != len( + self.chunk_grid.chunk_shape + ): + raise ValueError( + "`chunk_shape` and `shape` need to have the same number of dimensions." + ) + if self.dimension_names is not None and len(self.shape) != len(self.dimension_names): + raise ValueError( + "`dimension_names` and `shape` need to have the same number of dimensions." + ) + if self.fill_value is None: + raise ValueError("`fill_value` is required.") + for codec in self.codecs: + codec.validate(shape=self.shape, dtype=self.data_type, chunk_grid=self.chunk_grid) + + @property + def ndim(self) -> int: + return len(self.shape) + + @property + def dtype(self) -> ZDType[TBaseDType, TBaseScalar]: + return self.data_type + + @property + def chunks(self) -> tuple[int, ...]: + if isinstance(self.chunk_grid, RegularChunkGrid): + from zarr.codecs.sharding import ShardingCodec + + if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + sharding_codec = self.codecs[0] + assert isinstance(sharding_codec, ShardingCodec) # for mypy + return sharding_codec.chunk_shape + else: + return self.chunk_grid.chunk_shape + + msg = ( + f"The `chunks` attribute is only defined for arrays using `RegularChunkGrid`." + f"This array has a {self.chunk_grid} instead." + ) + raise NotImplementedError(msg) + + @property + def shards(self) -> tuple[int, ...] | None: + if isinstance(self.chunk_grid, RegularChunkGrid): + from zarr.codecs.sharding import ShardingCodec + + if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + return self.chunk_grid.chunk_shape + else: + return None + + msg = ( + f"The `shards` attribute is only defined for arrays using `RegularChunkGrid`." + f"This array has a {self.chunk_grid} instead." + ) + raise NotImplementedError(msg) + + @property + def inner_codecs(self) -> tuple[Codec, ...]: + if isinstance(self.chunk_grid, RegularChunkGrid): + from zarr.codecs.sharding import ShardingCodec + + if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + return self.codecs[0].codecs + return self.codecs + + def get_chunk_spec( + self, _chunk_coords: tuple[int, ...], array_config: ArrayConfig, prototype: BufferPrototype + ) -> ArraySpec: + assert isinstance(self.chunk_grid, RegularChunkGrid), ( + "Currently, only regular chunk grid is supported" + ) + return ArraySpec( + shape=self.chunk_grid.chunk_shape, + dtype=self.dtype, + fill_value=self.fill_value, + config=array_config, + prototype=prototype, + ) + + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + return self.chunk_key_encoding.encode_chunk_key(chunk_coords) + + def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + json_indent = config.get("json_indent") + d = self.to_dict() + return { + ZARR_JSON: prototype.buffer.from_bytes( + json.dumps(d, allow_nan=True, indent=json_indent).encode() + ) + } + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + # make a copy because we are modifying the dict + _data = data.copy() + + # check that the zarr_format attribute is correct + _ = parse_zarr_format(_data.pop("zarr_format")) + # check that the node_type attribute is correct + _ = parse_node_type_array(_data.pop("node_type")) + + data_type_json = _data.pop("data_type") + if not check_dtype_spec_v3(data_type_json): + raise ValueError(f"Invalid data_type: {data_type_json!r}") + data_type = get_data_type_from_json(data_type_json, zarr_format=3) + + # check that the fill value is consistent with the data type + try: + fill = _data.pop("fill_value") + fill_value_parsed = data_type.from_json_scalar(fill, zarr_format=3) + except ValueError as e: + raise TypeError(f"Invalid fill_value: {fill!r}") from e + + # check if there are extra keys + extra_keys = set(_data.keys()) - ARRAY_METADATA_KEYS + allowed_extra_fields: dict[str, AllowedExtraField] = {} + invalid_extra_fields = {} + for key in extra_keys: + val = _data[key] + if check_allowed_extra_field(val): + allowed_extra_fields[key] = val + else: + invalid_extra_fields[key] = val + if len(invalid_extra_fields) > 0: + msg = ( + "Got a Zarr V3 metadata document with the following disallowed extra fields:" + f"{sorted(invalid_extra_fields.keys())}." + 'Extra fields are not allowed unless they are a dict with a "must_understand" key' + "which is assigned the value `False`." + ) + raise MetadataValidationError(msg) + # TODO: replace this with a real type check! + _data_typed = cast(ArrayMetadataJSON_V3, _data) + + return cls( + shape=_data_typed["shape"], + chunk_grid=_data_typed["chunk_grid"], + chunk_key_encoding=_data_typed["chunk_key_encoding"], + codecs=_data_typed["codecs"], + attributes=_data_typed.get("attributes", {}), # type: ignore[arg-type] + dimension_names=_data_typed.get("dimension_names", None), + fill_value=fill_value_parsed, + data_type=data_type, + extra_fields=allowed_extra_fields, + storage_transformers=_data_typed.get("storage_transformers", ()), # type: ignore[arg-type] + ) + + def to_dict(self) -> dict[str, JSON]: + out_dict = super().to_dict() + extra_fields = out_dict.pop("extra_fields") + out_dict = out_dict | extra_fields # type: ignore[operator] + + out_dict["fill_value"] = self.data_type.to_json_scalar( + self.fill_value, zarr_format=self.zarr_format + ) + if not isinstance(out_dict, dict): + raise TypeError(f"Expected dict. Got {type(out_dict)}.") + + # if `dimension_names` is `None`, we do not include it in + # the metadata document + if out_dict["dimension_names"] is None: + out_dict.pop("dimension_names") + + # TODO: replace the `to_dict` / `from_dict` on the `Metadata`` class with + # to_json, from_json, and have ZDType inherit from `Metadata` + # until then, we have this hack here, which relies on the fact that to_dict will pass through + # any non-`Metadata` fields as-is. + dtype_meta = out_dict["data_type"] + if isinstance(dtype_meta, ZDType): + out_dict["data_type"] = dtype_meta.to_json(zarr_format=3) # type: ignore[unreachable] + return out_dict + + def update_shape(self, shape: tuple[int, ...]) -> Self: + return replace(self, shape=shape) + + def update_attributes(self, attributes: dict[str, JSON]) -> Self: + return replace(self, attributes=attributes) diff --git a/.venv/lib/python3.11/site-packages/zarr/core/sync.py b/.venv/lib/python3.11/site-packages/zarr/core/sync.py new file mode 100644 index 0000000000000000000000000000000000000000..fe435cc2b80eec8ecaae7893247a92a33797a6e7 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/sync.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import asyncio +import atexit +import logging +import os +import threading +from concurrent.futures import ThreadPoolExecutor, wait +from typing import TYPE_CHECKING, TypeVar + +from typing_extensions import ParamSpec + +from zarr.core.config import config + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine + from typing import Any + +logger = logging.getLogger(__name__) + + +P = ParamSpec("P") +T = TypeVar("T") + +# From https://github.com/fsspec/filesystem_spec/blob/master/fsspec/asyn.py + +iothread: list[threading.Thread | None] = [None] # dedicated IO thread +loop: list[asyncio.AbstractEventLoop | None] = [ + None +] # global event loop for any non-async instance +_lock: threading.Lock | None = None # global lock placeholder +_executor: ThreadPoolExecutor | None = None # global executor placeholder + + +class SyncError(Exception): + pass + + +def _get_lock() -> threading.Lock: + """Allocate or return a threading lock. + + The lock is allocated on first use to allow setting one lock per forked process. + """ + global _lock + if not _lock: + _lock = threading.Lock() + return _lock + + +def _get_executor() -> ThreadPoolExecutor: + """Return Zarr Thread Pool Executor + + The executor is allocated on first use. + """ + global _executor + if not _executor: + max_workers = config.get("threading.max_workers", None) + logger.debug("Creating Zarr ThreadPoolExecutor with max_workers=%s", max_workers) + _executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="zarr_pool") + _get_loop().set_default_executor(_executor) + return _executor + + +def cleanup_resources() -> None: + global _executor + if _executor: + _executor.shutdown(wait=True, cancel_futures=True) + _executor = None + + if loop[0] is not None: + with _get_lock(): + # Stop the event loop safely + loop[0].call_soon_threadsafe(loop[0].stop) # Stop loop from another thread + if iothread[0] is not None: + iothread[0].join(timeout=0.2) # Add a timeout to avoid hanging + + if iothread[0].is_alive(): + logger.warning( + "Thread did not finish cleanly; forcefully closing the event loop." + ) + + # Forcefully close the event loop to release resources + loop[0].close() + + # dereference the loop and iothread + loop[0] = None + iothread[0] = None + + +atexit.register(cleanup_resources) + + +def reset_resources_after_fork() -> None: + """ + Ensure that global resources are reset after a fork. Without this function, + forked processes will retain invalid references to the parent process's resources. + """ + global loop, iothread, _executor + # These lines are excluded from coverage because this function only runs in a child process, + # which is not observed by the test coverage instrumentation. Despite the apparent lack of + # test coverage, this function should be adequately tested by any test that uses Zarr IO with + # multiprocessing. + loop[0] = None # pragma: no cover + iothread[0] = None # pragma: no cover + _executor = None # pragma: no cover + + +# this is only available on certain operating systems +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=reset_resources_after_fork) + + +async def _runner(coro: Coroutine[Any, Any, T]) -> T | BaseException: + """ + Await a coroutine and return the result of running it. If awaiting the coroutine raises an + exception, the exception will be returned. + """ + try: + return await coro + except Exception as ex: + return ex + + +def sync( + coro: Coroutine[Any, Any, T], + loop: asyncio.AbstractEventLoop | None = None, + timeout: float | None = None, +) -> T: + """ + Make loop run coroutine until it returns. Runs in other thread + """ + if loop is None: + # NB: if the loop is not running *yet*, it is OK to submit work + # and we will wait for it + loop = _get_loop() + if _executor is None and config.get("threading.max_workers", None) is not None: + # trigger executor creation and attach to loop + _ = _get_executor() + if not isinstance(loop, asyncio.AbstractEventLoop): + raise TypeError(f"loop cannot be of type {type(loop)}") + if loop.is_closed(): + raise RuntimeError("Loop is not running") + try: + loop0 = asyncio.events.get_running_loop() + if loop0 is loop: + raise SyncError("Calling sync() from within a running loop") + except RuntimeError: + pass + + future = asyncio.run_coroutine_threadsafe(_runner(coro), loop) + + finished, unfinished = wait([future], return_when=asyncio.ALL_COMPLETED, timeout=timeout) + if len(unfinished) > 0: + raise TimeoutError(f"Coroutine {coro} failed to finish within {timeout} s") + assert len(finished) == 1 + return_result = next(iter(finished)).result() + + if isinstance(return_result, BaseException): + raise return_result + else: + return return_result + + +def _get_loop() -> asyncio.AbstractEventLoop: + """Create or return the default fsspec IO loop + + The loop will be running on a separate thread. + """ + if loop[0] is None: + with _get_lock(): + # repeat the check just in case the loop got filled between the + # previous two calls from another thread + if loop[0] is None: + logger.debug("Creating Zarr event loop") + new_loop = asyncio.new_event_loop() + loop[0] = new_loop + iothread[0] = threading.Thread(target=new_loop.run_forever, name="zarr_io") + assert iothread[0] is not None + iothread[0].daemon = True + iothread[0].start() + assert loop[0] is not None + return loop[0] + + +async def _collect_aiterator(data: AsyncIterator[T]) -> tuple[T, ...]: + """ + Collect an entire async iterator into a tuple + """ + result = [x async for x in data] + return tuple(result) + + +def collect_aiterator(data: AsyncIterator[T]) -> tuple[T, ...]: + """ + Synchronously collect an entire async iterator into a tuple. + """ + return sync(_collect_aiterator(data)) + + +class SyncMixin: + def _sync(self, coroutine: Coroutine[Any, Any, T]) -> T: + # TODO: refactor this to to take *args and **kwargs and pass those to the method + # this should allow us to better type the sync wrapper + return sync( + coroutine, + timeout=config.get("async.timeout"), + ) + + def _sync_iter(self, async_iterator: AsyncIterator[T]) -> list[T]: + async def iter_to_list() -> list[T]: + return [item async for item in async_iterator] + + return self._sync(iter_to_list()) + + +async def _with_semaphore( + func: Callable[[], Awaitable[T]], semaphore: asyncio.Semaphore | None = None +) -> T: + """ + Await the result of invoking the no-argument-callable ``func`` within the context manager + provided by a Semaphore, if one is provided. Otherwise, just await the result of invoking + ``func``. + """ + if semaphore is None: + return await func() + async with semaphore: + return await func() diff --git a/.venv/lib/python3.11/site-packages/zarr/core/sync_group.py b/.venv/lib/python3.11/site-packages/zarr/core/sync_group.py new file mode 100644 index 0000000000000000000000000000000000000000..8af514e93840b02608cab1a46599aa5407ac6ff1 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/core/sync_group.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from zarr.core.group import Group, GroupMetadata, _parse_async_node +from zarr.core.group import create_hierarchy as create_hierarchy_async +from zarr.core.group import create_nodes as create_nodes_async +from zarr.core.group import create_rooted_hierarchy as create_rooted_hierarchy_async +from zarr.core.group import get_node as get_node_async +from zarr.core.sync import _collect_aiterator, sync + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr.abc.store import Store + from zarr.core.common import ZarrFormat + from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + from zarr.types import AnyArray + + +def create_nodes( + *, store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata] +) -> Iterator[tuple[str, Group | AnyArray]]: + """Create a collection of arrays and / or groups concurrently. + + Note: no attempt is made to validate that these arrays and / or groups collectively form a + valid Zarr hierarchy. It is the responsibility of the caller of this function to ensure that + the ``nodes`` parameter satisfies any correctness constraints. + + Parameters + ---------- + store : Store + The storage backend to use. + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes + in the hierarchy, and the values are the metadata of the nodes. The + metadata must be either an instance of GroupMetadata, ArrayV3Metadata + or ArrayV2Metadata. + + Yields + ------ + Group | Array + The created nodes. + """ + coro = create_nodes_async(store=store, nodes=nodes) + + for key, value in sync(_collect_aiterator(coro)): + yield key, _parse_async_node(value) + + +def create_hierarchy( + *, + store: Store, + nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], + overwrite: bool = False, +) -> Iterator[tuple[str, Group | AnyArray]]: + """ + Create a complete zarr hierarchy from a collection of metadata objects. + + This function will parse its input to ensure that the hierarchy is complete. Any implicit groups + will be inserted as needed. For example, an input like + ```{'a/b': GroupMetadata}``` will be parsed to + ```{'': GroupMetadata, 'a': GroupMetadata, 'b': Groupmetadata}``` + + After input parsing, this function then creates all the nodes in the hierarchy concurrently. + + Arrays and Groups are yielded in the order they are created. This order is not stable and + should not be relied on. + + Parameters + ---------- + store : Store + The storage backend to use. + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes in the hierarchy, + relative to the root of the ``Store``. The root of the store can be specified with the empty + string ``''``. The values are instances of ``GroupMetadata`` or ``ArrayMetadata``. Note that + all values must have the same ``zarr_format`` -- it is an error to mix zarr versions in the + same hierarchy. + + Leading "/" characters from keys will be removed. + overwrite : bool + Whether to overwrite existing nodes. Defaults to ``False``, in which case an error is + raised instead of overwriting an existing array or group. + + This function will not erase an existing group unless that group is explicitly named in + ``nodes``. If ``nodes`` defines implicit groups, e.g. ``{`'a/b/c': GroupMetadata}``, and a + group already exists at path ``a``, then this function will leave the group at ``a`` as-is. + + Yields + ------ + tuple[str, Group | Array] + This function yields (path, node) pairs, in the order the nodes were created. + + Examples + -------- + ```python + from zarr import create_hierarchy + from zarr.storage import MemoryStore + from zarr.core.group import GroupMetadata + + store = MemoryStore() + nodes = {'a': GroupMetadata(attributes={'name': 'leaf'})} + nodes_created = dict(create_hierarchy(store=store, nodes=nodes)) + print(nodes) + # {'a': GroupMetadata(attributes={'name': 'leaf'}, zarr_format=3, consolidated_metadata=None, node_type='group')} + ``` + """ + coro = create_hierarchy_async(store=store, nodes=nodes, overwrite=overwrite) + + for key, value in sync(_collect_aiterator(coro)): + yield key, _parse_async_node(value) + + +def create_rooted_hierarchy( + *, + store: Store, + nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], + overwrite: bool = False, +) -> Group | AnyArray: + """ + Create a Zarr hierarchy with a root, and return the root node, which could be a ``Group`` + or ``Array`` instance. + + Parameters + ---------- + store : Store + The storage backend to use. + nodes : dict[str, GroupMetadata | ArrayV3Metadata | ArrayV2Metadata] + A dictionary defining the hierarchy. The keys are the paths of the nodes + in the hierarchy, and the values are the metadata of the nodes. The + metadata must be either an instance of GroupMetadata, ArrayV3Metadata + or ArrayV2Metadata. + overwrite : bool + Whether to overwrite existing nodes. Default is ``False``. + + Returns + ------- + Group | Array + """ + async_node = sync(create_rooted_hierarchy_async(store=store, nodes=nodes, overwrite=overwrite)) + return _parse_async_node(async_node) + + +def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyArray | Group: + """ + Get an Array or Group from a path in a Store. + + Parameters + ---------- + store : Store + The store-like object to read from. + path : str + The path to the node to read. + zarr_format : {2, 3} + The zarr format of the node to read. + + Returns + ------- + Array | Group + """ + + return _parse_async_node(sync(get_node_async(store=store, path=path, zarr_format=zarr_format))) diff --git a/.venv/lib/python3.11/site-packages/zarr/creation.py b/.venv/lib/python3.11/site-packages/zarr/creation.py new file mode 100644 index 0000000000000000000000000000000000000000..605b5af5de6d26bc80064eef83f2849a1a870706 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/creation.py @@ -0,0 +1,47 @@ +""" +Helpers for creating arrays. + +!!! warning "Deprecated" + This sub-module is deprecated. All functions here are defined in the top level zarr namespace instead. + +""" + +import warnings + +from zarr.api.synchronous import ( + array, + create, + empty, + empty_like, + full, + full_like, + ones, + ones_like, + open_array, + open_like, + zeros, + zeros_like, +) +from zarr.errors import ZarrDeprecationWarning + +__all__ = [ + "array", + "create", + "empty", + "empty_like", + "full", + "full_like", + "ones", + "ones_like", + "open_array", + "open_like", + "zeros", + "zeros_like", +] + +warnings.warn( + "zarr.creation is deprecated. " + "Import these functions from the top level zarr. namespace instead.", + ZarrDeprecationWarning, + stacklevel=2, +) diff --git a/.venv/lib/python3.11/site-packages/zarr/dtype.py b/.venv/lib/python3.11/site-packages/zarr/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..2c7eb651b0e76266f3ac571dd7a27c83ff4ce4a2 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/dtype.py @@ -0,0 +1,88 @@ +from zarr.core.dtype import ( + Bool, + Complex64, + Complex128, + DataTypeValidationError, + DateTime64, + DateTime64JSON_V2, + DateTime64JSON_V3, + FixedLengthUTF32, + FixedLengthUTF32JSON_V2, + FixedLengthUTF32JSON_V3, + Float16, + Float32, + Float64, + Int8, + Int16, + Int32, + Int64, + NullTerminatedBytes, + NullterminatedBytesJSON_V2, + NullTerminatedBytesJSON_V3, + RawBytes, + RawBytesJSON_V2, + RawBytesJSON_V3, + Structured, + StructuredJSON_V2, + StructuredJSON_V3, + TimeDelta64, + TimeDelta64JSON_V2, + TimeDelta64JSON_V3, + UInt8, + UInt16, + UInt32, + UInt64, + VariableLengthBytes, + VariableLengthBytesJSON_V2, + VariableLengthUTF8, + VariableLengthUTF8JSON_V2, + ZDType, + data_type_registry, + # Import for backwards compatibility, but not included in __all__ + # so it doesn't show up in the docs + parse_data_type, # noqa: F401 + parse_dtype, +) + +__all__ = [ + "Bool", + "Complex64", + "Complex128", + "DataTypeValidationError", + "DateTime64", + "DateTime64JSON_V2", + "DateTime64JSON_V3", + "FixedLengthUTF32", + "FixedLengthUTF32JSON_V2", + "FixedLengthUTF32JSON_V3", + "Float16", + "Float32", + "Float64", + "Int8", + "Int16", + "Int32", + "Int64", + "NullTerminatedBytes", + "NullTerminatedBytesJSON_V3", + "NullterminatedBytesJSON_V2", + "RawBytes", + "RawBytesJSON_V2", + "RawBytesJSON_V3", + "Structured", + "StructuredJSON_V2", + "StructuredJSON_V3", + "TimeDelta64", + "TimeDelta64JSON_V2", + "TimeDelta64JSON_V3", + "UInt8", + "UInt16", + "UInt32", + "UInt64", + "VariableLengthBytes", + "VariableLengthBytesJSON_V2", + "VariableLengthUTF8", + "VariableLengthUTF8JSON_V2", + "ZDType", + "data_type_registry", + "parse_dtype", +] diff --git a/.venv/lib/python3.11/site-packages/zarr/errors.py b/.venv/lib/python3.11/site-packages/zarr/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..bcd6a08deb57ffa3051d6503857836c54dd14b65 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/errors.py @@ -0,0 +1,146 @@ +__all__ = [ + "ArrayIndexError", + "ArrayNotFoundError", + "BaseZarrError", + "BoundsCheckError", + "ContainsArrayAndGroupError", + "ContainsArrayError", + "ContainsGroupError", + "GroupNotFoundError", + "MetadataValidationError", + "NegativeStepError", + "NodeTypeValidationError", + "UnstableSpecificationWarning", + "VindexInvalidSelectionError", + "ZarrDeprecationWarning", + "ZarrFutureWarning", + "ZarrRuntimeWarning", +] + + +class BaseZarrError(ValueError): + """ + Base error which all zarr errors are sub-classed from. + """ + + _msg: str = "{}" + + def __init__(self, *args: object) -> None: + """ + If a single argument is passed, treat it as a pre-formatted message. + + If multiple arguments are passed, they are used as arguments for a template string class + variable. This behavior is deprecated. + """ + if len(args) == 1: + super().__init__(args[0]) + else: + super().__init__(self._msg.format(*args)) + + +class NodeNotFoundError(BaseZarrError, FileNotFoundError): + """ + Raised when a node (array or group) is not found at a certain path. + """ + + +class ArrayNotFoundError(NodeNotFoundError): + """ + Raised when an array isn't found at a certain path. + """ + + _msg = "No array found in store {!r} at path {!r}" + + +class GroupNotFoundError(NodeNotFoundError): + """ + Raised when a group isn't found at a certain path. + """ + + _msg = "No group found in store {!r} at path {!r}" + + +class ContainsGroupError(BaseZarrError): + """Raised when a group already exists at a certain path.""" + + _msg = "A group exists in store {!r} at path {!r}." + + +class ContainsArrayError(BaseZarrError): + """Raised when an array already exists at a certain path.""" + + _msg = "An array exists in store {!r} at path {!r}." + + +class ContainsArrayAndGroupError(BaseZarrError): + """Raised when both array and group metadata are found at the same path.""" + + _msg = ( + "Array and group metadata documents (.zarray and .zgroup) were both found in store " + "{!r} at path {!r}. " + "Only one of these files may be present in a given directory / prefix. " + "Remove the .zarray file, or the .zgroup file, or both." + ) + + +class MetadataValidationError(BaseZarrError): + """Raised when the Zarr metadata is invalid in some way""" + + _msg = "Invalid value for '{}'. Expected '{}'. Got '{}'." + + +class UnknownCodecError(BaseZarrError): + """ + Raised when an unknown codec was used. + """ + + +class NodeTypeValidationError(MetadataValidationError): + """ + Specialized exception when the node_type of the metadata document is incorrect. + + This can be raised when the value is invalid or unexpected given the context, + for example an 'array' node when we expected a 'group'. + """ + + +class ZarrFutureWarning(FutureWarning): + """ + A warning intended for end users raised to indicate deprecated features. + """ + + +class UnstableSpecificationWarning(ZarrFutureWarning): + """ + A warning raised to indicate that a feature is outside the Zarr specification. + """ + + +class ZarrDeprecationWarning(DeprecationWarning): + """ + A warning raised to indicate that a feature will be removed in a future release. + """ + + +class ZarrUserWarning(UserWarning): + """ + A warning raised to report problems with user code. + """ + + +class ZarrRuntimeWarning(RuntimeWarning): + """ + A warning for dubious runtime behavior. + """ + + +class VindexInvalidSelectionError(IndexError): ... + + +class NegativeStepError(IndexError): ... + + +class BoundsCheckError(IndexError): ... + + +class ArrayIndexError(IndexError): ... diff --git a/.venv/lib/python3.11/site-packages/zarr/experimental/__init__.py b/.venv/lib/python3.11/site-packages/zarr/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3863510c65937bbf6a6fd76bb745754e6f60e6c4 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/experimental/__init__.py @@ -0,0 +1 @@ +"""The experimental module is a site for exporting new or experimental Zarr features.""" diff --git a/.venv/lib/python3.11/site-packages/zarr/experimental/cache_store.py b/.venv/lib/python3.11/site-packages/zarr/experimental/cache_store.py new file mode 100644 index 0000000000000000000000000000000000000000..1535b42f677c6d0f633fb928d91717ebf86e938a --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/experimental/cache_store.py @@ -0,0 +1,450 @@ +from __future__ import annotations + +import asyncio +import logging +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal, Self + +from zarr.abc.store import ByteRequest, Store +from zarr.storage._wrapper import WrapperStore + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from zarr.core.buffer.core import Buffer, BufferPrototype + +# A cache entry identifier. Plain ``str`` for full-key entries that live in +# the Store-backed cache; ``(str, ByteRequest)`` for byte-range entries that +# live in the in-memory range cache. +_CacheEntryKey = str | tuple[str, ByteRequest] + + +@dataclass(slots=True) +class _CacheState: + cache_order: OrderedDict[_CacheEntryKey, None] = field(default_factory=OrderedDict) + current_size: int = 0 + key_sizes: dict[_CacheEntryKey, int] = field(default_factory=dict) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + hits: int = 0 + misses: int = 0 + evictions: int = 0 + key_insert_times: dict[_CacheEntryKey, float] = field(default_factory=dict) + range_cache: dict[str, dict[ByteRequest, Buffer]] = field(default_factory=dict) + + +class CacheStore(WrapperStore[Store]): + """ + A dual-store caching implementation for Zarr stores. + + This cache wraps any Store implementation and uses a separate Store instance + as the cache backend. This provides persistent caching capabilities with + time-based expiration, size-based eviction, and flexible cache storage options. + + Full-key reads are cached in the Store-backed cache. Byte-range reads are + cached in a separate in-memory dictionary so that partial reads never + pollute the filesystem (or other persistent backend). Both caches share + the same ``max_size`` budget and LRU eviction policy. + + Parameters + ---------- + store : Store + The underlying store to wrap with caching + cache_store : Store + The store to use for caching (can be any Store implementation) + max_age_seconds : int | None, optional + Maximum age of cached entries in seconds. None means no expiration. + Default is None. + max_size : int | None, optional + Maximum size of the cache in bytes. When exceeded, least recently used + items are evicted. None means unlimited size. Default is None. + Note: Individual values larger than max_size will not be cached. + cache_set_data : bool, optional + Whether to cache data when it's written to the store. Default is True. + + Examples + -------- + ```python + import zarr + from zarr.storage import MemoryStore + from zarr.experimental.cache_store import CacheStore + + # Create a cached store + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + store=source_store, + cache_store=cache_store, + max_age_seconds=60, + max_size=1024*1024 + ) + + # Use it like any other store + array = zarr.create(shape=(100,), store=cached_store) + array[:] = 42 + ``` + + """ + + _cache: Store + max_age_seconds: int | Literal["infinity"] + max_size: int | None + cache_set_data: bool + _state: _CacheState + + def __init__( + self, + store: Store, + *, + cache_store: Store, + max_age_seconds: int | str = "infinity", + max_size: int | None = None, + cache_set_data: bool = True, + ) -> None: + super().__init__(store) + + if not cache_store.supports_deletes: + msg = ( + f"The provided cache store {cache_store} does not support deletes. " + "The cache_store must support deletes for CacheStore to function properly." + ) + raise ValueError(msg) + + self._cache = cache_store + # Validate and set max_age_seconds + if isinstance(max_age_seconds, str): + if max_age_seconds != "infinity": + raise ValueError("max_age_seconds string value must be 'infinity'") + self.max_age_seconds = "infinity" + else: + self.max_age_seconds = max_age_seconds + self.max_size = max_size + self.cache_set_data = cache_set_data + self._state = _CacheState() + + def _with_store(self, store: Store) -> Self: + # Cannot support this operation because it would share a cache, but have a new store + # So cache keys would conflict + raise NotImplementedError("CacheStore does not support this operation.") + + def with_read_only(self, read_only: bool = False) -> Self: + # Create a new cache store that shares the same cache and mutable state + store = type(self)( + store=self._store.with_read_only(read_only), + cache_store=self._cache, + max_age_seconds=self.max_age_seconds, + max_size=self.max_size, + cache_set_data=self.cache_set_data, + ) + store._state = self._state + return store + + def _is_key_fresh(self, entry_key: _CacheEntryKey) -> bool: + """Check if a cached entry is still fresh based on max_age_seconds. + + Uses monotonic time for accurate elapsed time measurement. + """ + if self.max_age_seconds == "infinity": + return True + now = time.monotonic() + elapsed = now - self._state.key_insert_times.get(entry_key, 0) + return elapsed < self.max_age_seconds + + async def _accommodate_value(self, value_size: int) -> None: + """Ensure there is enough space in the cache for a new value. + + Must be called while holding self._state.lock. + """ + if self.max_size is None: + return + + # Remove least recently used items until we have enough space + while self._state.current_size + value_size > self.max_size and self._state.cache_order: + # Get the least recently used key (first in OrderedDict) + lru_key = next(iter(self._state.cache_order)) + await self._evict_key(lru_key) + + async def _evict_key(self, entry_key: _CacheEntryKey) -> None: + """Evict a cache entry. + + Must be called while holding self._state.lock. + + For ``str`` keys the entry is deleted from the Store-backed cache. + For ``(str, ByteRequest)`` keys the entry is removed from the + in-memory range cache. + """ + key_size = self._state.key_sizes.get(entry_key, 0) + + if isinstance(entry_key, str): + await self._cache.delete(entry_key) + else: + base_key, byte_range = entry_key + per_key = self._state.range_cache.get(base_key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[base_key] + + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.key_sizes.pop(entry_key, None) + self._state.current_size = max(0, self._state.current_size - key_size) + self._state.evictions += 1 + + async def _track_entry(self, entry_key: _CacheEntryKey, value: Buffer) -> bool: + """Register *entry_key* in the shared size / LRU tracking. + + Returns ``True`` if the entry was tracked, ``False`` if the value + exceeds ``max_size`` and was skipped. Callers should roll back any + data they already stored when this returns ``False``. + + This method holds the lock for the entire operation to ensure atomicity. + """ + value_size = len(value) + + # Check if value exceeds max size + if self.max_size is not None and value_size > self.max_size: + return False + + async with self._state.lock: + # If key already exists, subtract old size first + if entry_key in self._state.key_sizes: + old_size = self._state.key_sizes[entry_key] + self._state.current_size -= old_size + + # Make room for the new value + await self._accommodate_value(value_size) + + # Update tracking atomically + self._state.cache_order[entry_key] = None + self._state.current_size += value_size + self._state.key_sizes[entry_key] = value_size + self._state.key_insert_times[entry_key] = time.monotonic() + + return True + + async def _update_access_order(self, entry_key: _CacheEntryKey) -> None: + """Update the access order for LRU tracking.""" + if entry_key in self._state.cache_order: + async with self._state.lock: + self._state.cache_order.move_to_end(entry_key) + + def _remove_from_tracking(self, entry_key: _CacheEntryKey) -> None: + """Remove an entry from all tracking structures. + + Must be called while holding self._state.lock. + """ + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.key_sizes.pop(entry_key, None) + + def _invalidate_range_entries(self, key: str) -> None: + """Remove all byte-range entries for *key* from the range cache and tracking. + + Must be called while holding self._state.lock. + """ + per_key = self._state.range_cache.pop(key, None) + if per_key is not None: + for byte_range in per_key: + entry_key: _CacheEntryKey = (key, byte_range) + entry_size = self._state.key_sizes.pop(entry_key, 0) + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.current_size = max(0, self._state.current_size - entry_size) + + # ------------------------------------------------------------------ + # get helpers + # ------------------------------------------------------------------ + + async def _cache_miss( + self, key: str, byte_range: ByteRequest | None, result: Buffer | None + ) -> None: + """Handle a cache miss by storing or cleaning up after a source-store fetch.""" + if result is None: + if byte_range is None: + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + else: + entry_key: _CacheEntryKey = (key, byte_range) + async with self._state.lock: + per_key = self._state.range_cache.get(key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[key] + self._remove_from_tracking(entry_key) + else: + if byte_range is None: + await self._cache.set(key, result) + await self._track_entry(key, result) + else: + entry_key = (key, byte_range) + self._state.range_cache.setdefault(key, {})[byte_range] = result + tracked = await self._track_entry(entry_key, result) + if not tracked: + # Value too large for the cache — roll back the insertion + per_key = self._state.range_cache.get(key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[key] + + async def _get_try_cache( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """Try to get data from cache first, falling back to source store.""" + if byte_range is None: + # Full-key read — use Store-backed cache + maybe_cached = await self._cache.get(key, prototype) + if maybe_cached is not None: + self._state.hits += 1 + await self._update_access_order(key) + return maybe_cached + else: + # Byte-range read — use in-memory range cache + entry_key: _CacheEntryKey = (key, byte_range) + per_key = self._state.range_cache.get(key) + if per_key is not None: + cached_buf = per_key.get(byte_range) + if cached_buf is not None: + self._state.hits += 1 + await self._update_access_order(entry_key) + return cached_buf + + # Cache miss — fetch from source store + self._state.misses += 1 + result = await super().get(key, prototype, byte_range) + await self._cache_miss(key, byte_range, result) + return result + + async def _get_no_cache( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """Get data directly from source store and update cache.""" + self._state.misses += 1 + result = await super().get(key, prototype, byte_range) + await self._cache_miss(key, byte_range, result) + return result + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """ + Retrieve data from the store, using cache when appropriate. + + Parameters + ---------- + key : str + The key to retrieve + prototype : BufferPrototype + Buffer prototype for creating the result buffer + byte_range : ByteRequest, optional + Byte range to retrieve + + Returns + ------- + Buffer | None + The retrieved data, or None if not found + """ + entry_key: _CacheEntryKey = (key, byte_range) if byte_range is not None else key + if not self._is_key_fresh(entry_key): + return await self._get_no_cache(key, prototype, byte_range) + else: + return await self._get_try_cache(key, prototype, byte_range) + + async def set(self, key: str, value: Buffer) -> None: + """ + Store data in the underlying store and optionally in cache. + + Parameters + ---------- + key : str + The key to store under + value : Buffer + The data to store + """ + await super().set(key, value) + # Invalidate all cached byte-range entries (source data changed) + async with self._state.lock: + self._invalidate_range_entries(key) + if self.cache_set_data: + await self._cache.set(key, value) + await self._track_entry(key, value) + else: + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + + async def delete(self, key: str) -> None: + """ + Delete data from both the underlying store and cache. + + Parameters + ---------- + key : str + The key to delete + """ + await super().delete(key) + # Invalidate all cached byte-range entries + async with self._state.lock: + self._invalidate_range_entries(key) + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + + def cache_info(self) -> dict[str, Any]: + """Return information about the cache state.""" + return { + "cache_store_type": type(self._cache).__name__, + "max_age_seconds": "infinity" + if self.max_age_seconds == "infinity" + else self.max_age_seconds, + "max_size": self.max_size, + "current_size": self._state.current_size, + "cache_set_data": self.cache_set_data, + "tracked_keys": len(self._state.key_insert_times), + "cached_keys": len(self._state.cache_order), + } + + def cache_stats(self) -> dict[str, Any]: + """Return cache performance statistics.""" + total_requests = self._state.hits + self._state.misses + hit_rate = self._state.hits / total_requests if total_requests > 0 else 0.0 + return { + "hits": self._state.hits, + "misses": self._state.misses, + "evictions": self._state.evictions, + "total_requests": total_requests, + "hit_rate": hit_rate, + } + + async def clear_cache(self) -> None: + """Clear all cached data and tracking information.""" + # Clear the cache store if it supports clear + if hasattr(self._cache, "clear"): + await self._cache.clear() + + # Reset tracking + async with self._state.lock: + self._state.key_insert_times.clear() + self._state.cache_order.clear() + self._state.key_sizes.clear() + self._state.range_cache.clear() + self._state.current_size = 0 + + def __repr__(self) -> str: + """Return string representation of the cache store.""" + return ( + f"{self.__class__.__name__}(" + f"store={self._store!r}, " + f"cache_store={self._cache!r}, " + f"max_age_seconds={self.max_age_seconds}, " + f"max_size={self.max_size}, " + f"current_size={self._state.current_size}, " + f"cached_keys={len(self._state.cache_order)})" + ) diff --git a/.venv/lib/python3.11/site-packages/zarr/metadata/__init__.py b/.venv/lib/python3.11/site-packages/zarr/metadata/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/zarr/metadata/migrate_v3.py b/.venv/lib/python3.11/site-packages/zarr/metadata/migrate_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..a72939100d02432a4207f4cabdc272f1c79a2b6b --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/metadata/migrate_v3.py @@ -0,0 +1,295 @@ +import asyncio +import logging +from typing import cast + +import numcodecs.abc + +import zarr +from zarr import Group +from zarr.abc.codec import ArrayArrayCodec, BytesBytesCodec, Codec +from zarr.abc.store import Store +from zarr.codecs.blosc import BloscCodec, BloscShuffle +from zarr.codecs.bytes import BytesCodec +from zarr.codecs.gzip import GzipCodec +from zarr.codecs.transpose import TransposeCodec +from zarr.codecs.zstd import ZstdCodec +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding +from zarr.core.common import ( + ZARR_JSON, + ZARRAY_JSON, + ZATTRS_JSON, + ZGROUP_JSON, + ZMETADATA_V2_JSON, + ZarrFormat, +) +from zarr.core.dtype.common import HasEndianness +from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType +from zarr.core.group import GroupMetadata +from zarr.core.metadata.v2 import ArrayV2Metadata +from zarr.core.metadata.v3 import ArrayV3Metadata +from zarr.core.sync import sync +from zarr.registry import get_codec_class +from zarr.storage import StorePath +from zarr.types import AnyArray + +_logger = logging.getLogger(__name__) + + +def migrate_v2_to_v3( + *, + input_store: Store, + output_store: Store | None = None, + dry_run: bool = False, +) -> None: + """Migrate all v2 metadata in a Zarr store to v3. + + This will create a zarr.json file at each level of a Zarr hierarchy (for every group / array). + v2 files (.zarray, .zattrs etc.) will be left as-is. + + Parameters + ---------- + input_store : Store + Input Zarr to migrate. + output_store : Store, optional + Output location to write v3 metadata (no array data will be copied). If not provided, v3 metadata will be + written to input_store. + dry_run : bool, optional + Enable a 'dry run' - files that would be created are logged, but no files are created or changed. + """ + + zarr_v2 = zarr.open(store=input_store, mode="r+") + + if output_store is not None: + # w- access to not allow overwrite of existing data + output_path = sync(StorePath.open(output_store, path="", mode="w-")) + else: + output_path = zarr_v2.store_path + + migrate_to_v3(zarr_v2, output_path, dry_run=dry_run) + + +def migrate_to_v3(zarr_v2: AnyArray | Group, output_path: StorePath, dry_run: bool = False) -> None: + """Migrate all v2 metadata in a Zarr array/group to v3. + + Note - if a group is provided, then all arrays / groups within this group will also be converted. + A zarr.json file will be created for each level and written to output_path, with any v2 files + (.zarray, .zattrs etc.) left as-is. + + Parameters + ---------- + zarr_v2 : Array | Group + An array or group with zarr_format = 2 + output_path : StorePath + The store path to write generated v3 metadata to. + dry_run : bool, optional + Enable a 'dry run' - files that would be created are logged, but no files are created or changed. + """ + if not zarr_v2.metadata.zarr_format == 2: + raise TypeError("Only arrays / groups with zarr v2 metadata can be converted") + + if isinstance(zarr_v2.metadata, GroupMetadata): + _convert_group(zarr_v2, output_path, dry_run) + else: + _convert_array(zarr_v2, output_path, dry_run) + + +async def remove_metadata( + store: Store, + zarr_format: ZarrFormat, + force: bool = False, + dry_run: bool = False, +) -> None: + """Remove all v2 (.zarray, .zattrs, .zgroup, .zmetadata) or v3 (zarr.json) metadata files from the given Zarr. + + Note - this will remove metadata files at all levels of the hierarchy (every group and array). + + Parameters + ---------- + store : Store + Zarr to remove metadata from. + zarr_format : ZarrFormat + Which format's metadata to remove - 2 or 3. + force : bool, optional + When False, metadata can only be removed if a valid alternative exists e.g. deletion of v2 metadata will + only be allowed when v3 metadata is also present. When True, metadata can be removed when there is no + alternative. + dry_run : bool, optional + Enable a 'dry run' - files that would be deleted are logged, but no files are removed or changed. + """ + + if not store.supports_deletes: + raise ValueError("Store must support deletes to remove metadata") + store_path = await StorePath.open(store, path="", mode="r+") + + metadata_files_all = { + 2: [ZARRAY_JSON, ZATTRS_JSON, ZGROUP_JSON, ZMETADATA_V2_JSON], + 3: [ZARR_JSON], + } + + if zarr_format == 2: + alternative_metadata = 3 + else: + alternative_metadata = 2 + + awaitables = [] + async for file_path in store.list(): + parent_path, _, file_name = file_path.rpartition("/") + + if file_name not in metadata_files_all[zarr_format]: + continue + + if force or await _metadata_exists( + cast(ZarrFormat, alternative_metadata), store_path / parent_path + ): + _logger.info("Deleting metadata at %s", store_path / file_path) + if not dry_run: + awaitables.append((store_path / file_path).delete()) + else: + raise ValueError( + f"Cannot remove v{zarr_format} metadata at {store_path / file_path} - no v{alternative_metadata} " + "metadata exists. To delete anyway, use the 'force' option." + ) + + await asyncio.gather(*awaitables) + + +def _convert_group(zarr_v2: Group, output_path: StorePath, dry_run: bool) -> None: + if zarr_v2.metadata.consolidated_metadata is not None: + raise NotImplementedError("Migration of consolidated metadata isn't supported.") + + # process members of the group + for key in zarr_v2: + migrate_to_v3(zarr_v2[key], output_path=output_path / key, dry_run=dry_run) + + # write group's converted metadata + group_metadata_v3 = GroupMetadata( + attributes=zarr_v2.metadata.attributes, zarr_format=3, consolidated_metadata=None + ) + sync(_save_v3_metadata(group_metadata_v3, output_path, dry_run=dry_run)) + + +def _convert_array(zarr_v2: AnyArray, output_path: StorePath, dry_run: bool) -> None: + array_metadata_v3 = _convert_array_metadata(cast(ArrayV2Metadata, zarr_v2.metadata)) + sync(_save_v3_metadata(array_metadata_v3, output_path, dry_run=dry_run)) + + +async def _metadata_exists(zarr_format: ZarrFormat, store_path: StorePath) -> bool: + metadata_files_required = {2: [ZARRAY_JSON, ZGROUP_JSON], 3: [ZARR_JSON]} + + for metadata_file in metadata_files_required[zarr_format]: + if await (store_path / metadata_file).exists(): + return True + + return False + + +def _convert_array_metadata(metadata_v2: ArrayV2Metadata) -> ArrayV3Metadata: + chunk_key_encoding = V2ChunkKeyEncoding(separator=metadata_v2.dimension_separator) + + codecs: list[Codec] = [] + + # array-array codecs + if metadata_v2.order == "F": + # F is equivalent to order: n-1, ... 1, 0 + codecs.append(TransposeCodec(order=list(range(len(metadata_v2.shape) - 1, -1, -1)))) + + if metadata_v2.filters is not None: + codecs.extend(_convert_filters(metadata_v2.filters)) + + # array-bytes codecs + if not isinstance(metadata_v2.dtype, HasEndianness): + codecs.append(BytesCodec(endian=None)) + else: + codecs.append(BytesCodec(endian=metadata_v2.dtype.endianness)) + + # bytes-bytes codecs + if metadata_v2.compressor is not None: + bytes_bytes_codec = _convert_compressor(metadata_v2.compressor, metadata_v2.dtype) + codecs.append(bytes_bytes_codec) + + return ArrayV3Metadata( + shape=metadata_v2.shape, + data_type=metadata_v2.dtype, + chunk_grid=metadata_v2.chunk_grid, + chunk_key_encoding=chunk_key_encoding, + fill_value=metadata_v2.fill_value, + codecs=codecs, + attributes=metadata_v2.attributes, + dimension_names=None, + storage_transformers=None, + ) + + +def _convert_filters(filters: tuple[numcodecs.abc.Codec, ...]) -> list[ArrayArrayCodec]: + filters_codecs = [_find_numcodecs_zarr3(filter) for filter in filters] + for codec in filters_codecs: + if not isinstance(codec, ArrayArrayCodec): + raise TypeError(f"Filter {type(codec)} is not an ArrayArrayCodec") + + return cast(list[ArrayArrayCodec], filters_codecs) + + +def _convert_compressor( + compressor: numcodecs.abc.Codec, dtype: ZDType[TBaseDType, TBaseScalar] +) -> BytesBytesCodec: + match compressor.codec_id: + case "blosc": + return BloscCodec( + typesize=dtype.to_native_dtype().itemsize, + cname=compressor.cname, + clevel=compressor.clevel, + shuffle=BloscShuffle.from_int(compressor.shuffle), + blocksize=compressor.blocksize, + ) + + case "zstd": + return ZstdCodec( + level=compressor.level, + checksum=compressor.checksum, + ) + + case "gzip": + return GzipCodec(level=compressor.level) + + case _: + # If possible, find matching zarr.codecs.numcodecs codec + compressor_codec = _find_numcodecs_zarr3(compressor) + + if not isinstance(compressor_codec, BytesBytesCodec): + raise TypeError(f"Compressor {type(compressor_codec)} is not a BytesBytesCodec") + + return compressor_codec + + +def _find_numcodecs_zarr3(numcodecs_codec: numcodecs.abc.Codec) -> Codec: + """Find matching zarr.codecs.numcodecs codec (if it exists)""" + + numcodec_name = f"numcodecs.{numcodecs_codec.codec_id}" + numcodec_dict = { + "name": numcodec_name, + "configuration": numcodecs_codec.get_config(), + } + + try: + codec_v3 = get_codec_class(numcodec_name) + except KeyError as exc: + raise ValueError( + f"Couldn't find corresponding zarr.codecs.numcodecs codec for {numcodecs_codec.codec_id}" + ) from exc + + return codec_v3.from_dict(numcodec_dict) + + +async def _save_v3_metadata( + metadata_v3: ArrayV3Metadata | GroupMetadata, output_path: StorePath, dry_run: bool = False +) -> None: + zarr_json_path = output_path / ZARR_JSON + if await zarr_json_path.exists(): + raise ValueError(f"{ZARR_JSON} already exists at {zarr_json_path}") + + _logger.info("Saving metadata to %s", zarr_json_path) + to_save = metadata_v3.to_buffer_dict(default_buffer_prototype()) + + if not dry_run: + await zarr_json_path.set_if_not_exists(to_save[ZARR_JSON]) diff --git a/.venv/lib/python3.11/site-packages/zarr/py.typed b/.venv/lib/python3.11/site-packages/zarr/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.venv/lib/python3.11/site-packages/zarr/registry.py b/.venv/lib/python3.11/site-packages/zarr/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..d0850a1387bb9ad0c00fa6624e4da79ff8e8cbec --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/registry.py @@ -0,0 +1,338 @@ +from __future__ import annotations + +import warnings +from collections import defaultdict +from importlib.metadata import entry_points as get_entry_points +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from zarr.core.config import BadConfigError, config +from zarr.core.dtype import data_type_registry +from zarr.errors import ZarrUserWarning + +if TYPE_CHECKING: + from importlib.metadata import EntryPoint + + from zarr.abc.codec import ( + ArrayArrayCodec, + ArrayBytesCodec, + BytesBytesCodec, + Codec, + CodecJSON_V2, + CodecPipeline, + ) + from zarr.abc.numcodec import Numcodec + from zarr.core.buffer import Buffer, NDBuffer + from zarr.core.chunk_key_encodings import ChunkKeyEncoding + from zarr.core.common import JSON + +__all__ = [ + "Registry", + "get_buffer_class", + "get_chunk_key_encoding_class", + "get_codec_class", + "get_ndbuffer_class", + "get_pipeline_class", + "register_buffer", + "register_chunk_key_encoding", + "register_codec", + "register_ndbuffer", + "register_pipeline", +] + +T = TypeVar("T") + + +class Registry(dict[str, type[T]], Generic[T]): + def __init__(self) -> None: + super().__init__() + self.lazy_load_list: list[EntryPoint] = [] + + def lazy_load(self, use_entrypoint_name: bool = False) -> None: + for e in self.lazy_load_list: + self.register(e.load(), qualname=e.name if use_entrypoint_name else None) + + self.lazy_load_list.clear() + + def register(self, cls: type[T], qualname: str | None = None) -> None: + if qualname is None: + qualname = fully_qualified_name(cls) + self[qualname] = cls + + +__codec_registries: dict[str, Registry[Codec]] = defaultdict(Registry) +__pipeline_registry: Registry[CodecPipeline] = Registry() +__buffer_registry: Registry[Buffer] = Registry() +__ndbuffer_registry: Registry[NDBuffer] = Registry() +__chunk_key_encoding_registry: Registry[ChunkKeyEncoding] = Registry() + +""" +The registry module is responsible for managing implementations of codecs, +pipelines, buffers, ndbuffers, and chunk key encodings and collecting them from entrypoints. +The implementation used is determined by the config. + +The registry module is also responsible for managing dtypes. +""" + + +def _collect_entrypoints() -> list[Registry[Any]]: + """ + Collects codecs, pipelines, dtypes, buffers and ndbuffers from entrypoints. + Entry points can either be single items or groups of items. + Allowed syntax for entry_points.txt is e.g. + + [zarr.codecs] + gzip = package:EntrypointGzipCodec1 + [zarr.codecs.gzip] + some_name = package:EntrypointGzipCodec2 + another = package:EntrypointGzipCodec3 + + [zarr] + buffer = package:TestBuffer1 + [zarr.buffer] + xyz = package:TestBuffer2 + abc = package:TestBuffer3 + ... + """ + entry_points = get_entry_points() + + __buffer_registry.lazy_load_list.extend(entry_points.select(group="zarr.buffer")) + __buffer_registry.lazy_load_list.extend(entry_points.select(group="zarr", name="buffer")) + __ndbuffer_registry.lazy_load_list.extend(entry_points.select(group="zarr.ndbuffer")) + __ndbuffer_registry.lazy_load_list.extend(entry_points.select(group="zarr", name="ndbuffer")) + + data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr.data_type")) + data_type_registry._lazy_load_list.extend(entry_points.select(group="zarr", name="data_type")) + + __chunk_key_encoding_registry.lazy_load_list.extend( + entry_points.select(group="zarr.chunk_key_encoding") + ) + __chunk_key_encoding_registry.lazy_load_list.extend( + entry_points.select(group="zarr", name="chunk_key_encoding") + ) + + __pipeline_registry.lazy_load_list.extend(entry_points.select(group="zarr.codec_pipeline")) + __pipeline_registry.lazy_load_list.extend( + entry_points.select(group="zarr", name="codec_pipeline") + ) + for e in entry_points.select(group="zarr.codecs"): + __codec_registries[e.name].lazy_load_list.append(e) + for group in entry_points.groups: + if group.startswith("zarr.codecs."): + codec_name = group.split(".")[2] + __codec_registries[codec_name].lazy_load_list.extend(entry_points.select(group=group)) + return [ + *__codec_registries.values(), + __pipeline_registry, + __buffer_registry, + __ndbuffer_registry, + __chunk_key_encoding_registry, + ] + + +def _reload_config() -> None: + config.refresh() + + +def fully_qualified_name(cls: type) -> str: + module = cls.__module__ + return module + "." + cls.__qualname__ + + +def register_codec(key: str, codec_cls: type[Codec], *, qualname: str | None = None) -> None: + if key not in __codec_registries: + __codec_registries[key] = Registry() + __codec_registries[key].register(codec_cls, qualname=qualname) + + +def register_pipeline(pipe_cls: type[CodecPipeline]) -> None: + __pipeline_registry.register(pipe_cls) + + +def register_ndbuffer(cls: type[NDBuffer], qualname: str | None = None) -> None: + __ndbuffer_registry.register(cls, qualname) + + +def register_buffer(cls: type[Buffer], qualname: str | None = None) -> None: + __buffer_registry.register(cls, qualname) + + +def register_chunk_key_encoding(key: str, cls: type) -> None: + __chunk_key_encoding_registry.register(cls, key) + + +def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: + if reload_config: + _reload_config() + + if key in __codec_registries: + # logger.debug("Auto loading codec '%s' from entrypoint", codec_id) + __codec_registries[key].lazy_load() + + codec_classes = __codec_registries[key] + if not codec_classes: + raise KeyError(key) + config_entry = config.get("codecs", {}).get(key) + if config_entry is None: + if len(codec_classes) == 1: + return next(iter(codec_classes.values())) + warnings.warn( + f"Codec '{key}' not configured in config. Selecting any implementation.", + stacklevel=2, + category=ZarrUserWarning, + ) + return list(codec_classes.values())[-1] + selected_codec_cls = codec_classes[config_entry] + + if selected_codec_cls: + return selected_codec_cls + raise KeyError(key) + + +def _resolve_codec(data: dict[str, JSON]) -> Codec: + """ + Get a codec instance from a dict representation of that codec. + """ + # TODO: narrow the type of the input to only those dicts that map on to codec class instances. + return get_codec_class(data["name"]).from_dict(data) # type: ignore[arg-type] + + +def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: + """ + Normalize the input to a ``BytesBytesCodec`` instance. + If the input is already a ``BytesBytesCodec``, it is returned as is. If the input is a dict, it + is converted to a ``BytesBytesCodec`` instance via the ``_resolve_codec`` function. + """ + from zarr.abc.codec import BytesBytesCodec + + if isinstance(data, dict): + result = _resolve_codec(data) + if not isinstance(result, BytesBytesCodec): + msg = f"Expected a dict representation of a BytesBytesCodec; got a dict representation of a {type(result)} instead." + raise TypeError(msg) + else: + if not isinstance(data, BytesBytesCodec): + raise TypeError(f"Expected a BytesBytesCodec. Got {type(data)} instead.") + result = data + return result + + +def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: + """ + Normalize the input to a ``ArrayBytesCodec`` instance. + If the input is already a ``ArrayBytesCodec``, it is returned as is. If the input is a dict, it + is converted to a ``ArrayBytesCodec`` instance via the ``_resolve_codec`` function. + """ + from zarr.abc.codec import ArrayBytesCodec + + if isinstance(data, dict): + result = _resolve_codec(data) + if not isinstance(result, ArrayBytesCodec): + msg = f"Expected a dict representation of an ArrayBytesCodec; got a dict representation of a {type(result)} instead." + raise TypeError(msg) + else: + if not isinstance(data, ArrayBytesCodec): + raise TypeError(f"Expected an ArrayBytesCodec. Got {type(data)} instead.") + result = data + return result + + +def _parse_array_array_codec(data: dict[str, JSON] | Codec) -> ArrayArrayCodec: + """ + Normalize the input to a ``ArrayArrayCodec`` instance. + If the input is already a ``ArrayArrayCodec``, it is returned as is. If the input is a dict, it + is converted to a ``ArrayArrayCodec`` instance via the ``_resolve_codec`` function. + """ + from zarr.abc.codec import ArrayArrayCodec + + if isinstance(data, dict): + result = _resolve_codec(data) + if not isinstance(result, ArrayArrayCodec): + msg = f"Expected a dict representation of an ArrayArrayCodec; got a dict representation of a {type(result)} instead." + raise TypeError(msg) + else: + if not isinstance(data, ArrayArrayCodec): + raise TypeError(f"Expected an ArrayArrayCodec. Got {type(data)} instead.") + result = data + return result + + +def get_pipeline_class(reload_config: bool = False) -> type[CodecPipeline]: + if reload_config: + _reload_config() + __pipeline_registry.lazy_load() + path = config.get("codec_pipeline.path") + pipeline_class = __pipeline_registry.get(path) + if pipeline_class: + return pipeline_class + raise BadConfigError( + f"Pipeline class '{path}' not found in registered pipelines: {list(__pipeline_registry)}." + ) + + +def get_buffer_class(reload_config: bool = False) -> type[Buffer]: + if reload_config: + _reload_config() + __buffer_registry.lazy_load() + + path = config.get("buffer") + buffer_class = __buffer_registry.get(path) + if buffer_class: + return buffer_class + raise BadConfigError( + f"Buffer class '{path}' not found in registered buffers: {list(__buffer_registry)}." + ) + + +def get_ndbuffer_class(reload_config: bool = False) -> type[NDBuffer]: + if reload_config: + _reload_config() + __ndbuffer_registry.lazy_load() + path = config.get("ndbuffer") + ndbuffer_class = __ndbuffer_registry.get(path) + if ndbuffer_class: + return ndbuffer_class + raise BadConfigError( + f"NDBuffer class '{path}' not found in registered buffers: {list(__ndbuffer_registry)}." + ) + + +def get_chunk_key_encoding_class(key: str) -> type[ChunkKeyEncoding]: + __chunk_key_encoding_registry.lazy_load(use_entrypoint_name=True) + if key not in __chunk_key_encoding_registry: + raise KeyError( + f"Chunk key encoding '{key}' not found in registered chunk key encodings: {list(__chunk_key_encoding_registry)}." + ) + return __chunk_key_encoding_registry[key] + + +_collect_entrypoints() + + +def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec: + """ + Resolve a numcodec codec from the numcodecs registry. + + This requires the Numcodecs package to be installed. + + Parameters + ---------- + data : CodecJSON_V2 + The JSON metadata for the codec. + + Returns + ------- + codec : Numcodec + + Examples + -------- + ```python + from zarr.registry import get_numcodec + codec = get_numcodec({'id': 'zlib', 'level': 1}) + codec + # Zlib(level=1) + ``` + """ + + from numcodecs.registry import get_codec + + return get_codec(data) # type: ignore[no-any-return] diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__init__.py b/.venv/lib/python3.11/site-packages/zarr/storage/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..00df50214fc2e397a3eb548e2c6c0de7ae89e22e --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/__init__.py @@ -0,0 +1,44 @@ +import sys +import warnings +from types import ModuleType +from typing import Any + +from zarr.errors import ZarrDeprecationWarning +from zarr.storage._common import StoreLike, StorePath +from zarr.storage._fsspec import FsspecStore +from zarr.storage._local import LocalStore +from zarr.storage._logging import LoggingStore +from zarr.storage._memory import GpuMemoryStore, MemoryStore +from zarr.storage._obstore import ObjectStore +from zarr.storage._wrapper import WrapperStore +from zarr.storage._zip import ZipStore + +__all__ = [ + "FsspecStore", + "GpuMemoryStore", + "LocalStore", + "LoggingStore", + "MemoryStore", + "ObjectStore", + "StoreLike", + "StorePath", + "WrapperStore", + "ZipStore", +] + + +class VerboseModule(ModuleType): + def __setattr__(self, attr: str, value: Any) -> None: + if attr == "default_compressor": + warnings.warn( + "setting zarr.storage.default_compressor is deprecated, use " + "zarr.config to configure array.v2_default_compressor " + "e.g. config.set({'codecs.zstd':'numcodecs.Zstd', 'array.v2_default_compressor.numeric': 'zstd'})", + ZarrDeprecationWarning, + stacklevel=1, + ) + else: + super().__setattr__(attr, value) + + +sys.modules[__name__].__class__ = VerboseModule diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/__init__.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a67f68b1f5ac333285ad97a1fff39d8ac172ab69 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/__init__.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_common.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_common.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9ea60b48f79d60abae4681dae6a5132d0b4a3a3 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_common.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_fsspec.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_fsspec.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a14bf04fca192adf1fd7f7c2a4f3057d7fb69839 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_fsspec.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_local.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_local.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f7221232314baea52d29854e6d0e9faeba49182 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_local.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_logging.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_logging.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc7496c57b14134a5c7570e22a5a58948768f493 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_logging.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_memory.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_memory.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..344baf9c54995ed3ec2319dadb0b9b6d808decf7 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_memory.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_obstore.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_obstore.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d117ed76b3cb703c35351ecbc2dac758194e555d Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_obstore.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_utils.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba09b6ea2bc05e8db038f6ec346508cd4d04794f Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_utils.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_wrapper.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_wrapper.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2efe6d26c45128f88f848bbe78990c5aaaa24e8 Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_wrapper.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_zip.cpython-311.pyc b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_zip.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..839df9e9560eb7737ed98afd87643745a618769d Binary files /dev/null and b/.venv/lib/python3.11/site-packages/zarr/storage/__pycache__/_zip.cpython-311.pyc differ diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_common.py b/.venv/lib/python3.11/site-packages/zarr/storage/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..08c05864aab730346d87e88232396dea23de736d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_common.py @@ -0,0 +1,667 @@ +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias + +from zarr.abc.store import ( + ByteRequest, + Store, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, +) +from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.common import ( + ANY_ACCESS_MODE, + ZARR_JSON, + ZARRAY_JSON, + ZGROUP_JSON, + AccessModeLiteral, + ZarrFormat, +) +from zarr.errors import ContainsArrayAndGroupError, ContainsArrayError, ContainsGroupError +from zarr.storage._local import LocalStore +from zarr.storage._memory import MemoryStore +from zarr.storage._utils import normalize_path + +_has_fsspec = importlib.util.find_spec("fsspec") +if _has_fsspec: + from fsspec.mapping import FSMap +else: + FSMap = None + +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype + + +def _dereference_path(root: str, path: str) -> str: + if not isinstance(root, str): + msg = f"{root=} is not a string ({type(root)=})" # type: ignore[unreachable] + raise TypeError(msg) + if not isinstance(path, str): + msg = f"{path=} is not a string ({type(path)=})" # type: ignore[unreachable] + raise TypeError(msg) + root = root.rstrip("/") + path = f"{root}/{path}" if root else path + return path.rstrip("/") + + +class StorePath: + """ + Path-like interface for a Store. + + Parameters + ---------- + store : Store + The store to use. + path : str + The path within the store. + """ + + store: Store + path: str + + def __init__(self, store: Store, path: str = "") -> None: + self.store = store + self.path = normalize_path(path) + + @property + def read_only(self) -> bool: + return self.store.read_only + + @classmethod + async def _create_open_instance(cls, store: Store, path: str) -> Self: + """Helper to create and return a StorePath instance.""" + await store._ensure_open() + return cls(store, path) + + @classmethod + async def open(cls, store: Store, path: str, mode: AccessModeLiteral | None = None) -> Self: + """ + Open StorePath based on the provided mode. + + * If the mode is None, return an opened version of the store with no changes. + * If the mode is 'r+', 'w-', 'w', or 'a' and the store is read-only, raise a ValueError. + * If the mode is 'r' and the store is not read-only, return a copy of the store with read_only set to True. + * If the mode is 'w-' and the store is not read-only and the StorePath contains keys, raise a FileExistsError. + * If the mode is 'w' and the store is not read-only, delete all keys nested within the StorePath. + + Parameters + ---------- + mode : AccessModeLiteral + The mode to use when initializing the store path. + + The accepted values are: + + - ``'r'``: read only (must exist) + - ``'r+'``: read/write (must exist) + - ``'a'``: read/write (create if doesn't exist) + - ``'w'``: read/write (overwrite if exists) + - ``'w-'``: read/write (create if doesn't exist). + + Raises + ------ + FileExistsError + If the mode is 'w-' and the store path already exists. + ValueError + If the mode is not "r" and the store is read-only, or + """ + + # fastpath if mode is None + if mode is None: + return await cls._create_open_instance(store, path) + + if mode not in ANY_ACCESS_MODE: + raise ValueError(f"Invalid mode: {mode}, expected one of {ANY_ACCESS_MODE}") + + if store.read_only: + # Don't allow write operations on a read-only store + if mode != "r": + raise ValueError( + f"Store is read-only but mode is {mode!r}. Create a writable store or use 'r' mode." + ) + self = await cls._create_open_instance(store, path) + elif mode == "r": + # Create read-only copy for read mode on writable store + try: + read_only_store = store.with_read_only(True) + except NotImplementedError as e: + raise ValueError( + "Store is not read-only but mode is 'r'. Unable to create a read-only copy of the store. " + "Please use a read-only store or a storage class that implements .with_read_only()." + ) from e + self = await cls._create_open_instance(read_only_store, path) + else: + # writable store and writable mode + self = await cls._create_open_instance(store, path) + + # Handle mode-specific operations + match mode: + case "w-": + if not await self.is_empty(): + raise FileExistsError( + f"Cannot create '{path}' with mode 'w-' because it already contains data. " + f"Use mode 'w' to overwrite or 'a' to append." + ) + case "w": + await self.delete_dir() + return self + + async def get( + self, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """ + Read bytes from the store. + + Parameters + ---------- + prototype : BufferPrototype, optional + The buffer prototype to use when reading the bytes. + byte_range : ByteRequest, optional + The range of bytes to read. + + Returns + ------- + buffer : Buffer or None + The read bytes, or None if the key does not exist. + """ + if prototype is None: + prototype = default_buffer_prototype() + return await self.store.get(self.path, prototype=prototype, byte_range=byte_range) + + async def set(self, value: Buffer) -> None: + """ + Write bytes to the store. + + Parameters + ---------- + value : Buffer + The buffer to write. + """ + await self.store.set(self.path, value) + + async def delete(self) -> None: + """ + Delete the key from the store. + + Raises + ------ + NotImplementedError + If the store does not support deletion. + """ + await self.store.delete(self.path) + + async def delete_dir(self) -> None: + """ + Delete all keys with the given prefix from the store. + """ + await self.store.delete_dir(self.path) + + async def set_if_not_exists(self, default: Buffer) -> None: + """ + Store a key to ``value`` if the key is not already present. + + Parameters + ---------- + default : Buffer + The buffer to store if the key is not already present. + """ + await self.store.set_if_not_exists(self.path, default) + + async def exists(self) -> bool: + """ + Check if the key exists in the store. + + Returns + ------- + bool + True if the key exists in the store, False otherwise. + """ + return await self.store.exists(self.path) + + async def is_empty(self) -> bool: + """ + Check if any keys exist in the store with the given prefix. + + Returns + ------- + bool + True if no keys exist in the store with the given prefix, False otherwise. + """ + return await self.store.is_empty(self.path) + + # ------------------------------------------------------------------- + # Synchronous IO delegation + # ------------------------------------------------------------------- + + def get_sync( + self, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Synchronous read — delegates to ``self.store.get_sync(self.path, ...)``.""" + if not isinstance(self.store, SupportsGetSync): + raise TypeError(f"Store {type(self.store).__name__} does not support synchronous get.") + if prototype is None: + prototype = default_buffer_prototype() + return self.store.get_sync(self.path, prototype=prototype, byte_range=byte_range) + + def set_sync(self, value: Buffer) -> None: + """Synchronous write — delegates to ``self.store.set_sync(self.path, value)``.""" + if not isinstance(self.store, SupportsSetSync): + raise TypeError(f"Store {type(self.store).__name__} does not support synchronous set.") + self.store.set_sync(self.path, value) + + def delete_sync(self) -> None: + """Synchronous delete — delegates to ``self.store.delete_sync(self.path)``.""" + if not isinstance(self.store, SupportsDeleteSync): + raise TypeError( + f"Store {type(self.store).__name__} does not support synchronous delete." + ) + self.store.delete_sync(self.path) + + def __truediv__(self, other: str) -> StorePath: + """Combine this store path with another path""" + return self.__class__(self.store, _dereference_path(self.path, other)) + + def __str__(self) -> str: + return _dereference_path(str(self.store), self.path) + + def __repr__(self) -> str: + return f"StorePath({self.store.__class__.__name__}, '{self}')" + + def __eq__(self, other: object) -> bool: + """ + Check if two StorePath objects are equal. + + Returns + ------- + bool + True if the two objects are equal, False otherwise. + + Notes + ----- + Two StorePath objects are considered equal if their stores are equal + and their paths are equal. + """ + try: + return self.store == other.store and self.path == other.path # type: ignore[attr-defined, no-any-return] + except Exception: + pass + return False + + +StoreLike: TypeAlias = Store | StorePath | FSMap | Path | str | dict[str, Buffer] + + +async def make_store( + store_like: StoreLike | None, + *, + mode: AccessModeLiteral | None = None, + storage_options: dict[str, Any] | None = None, +) -> Store: + """ + Convert a `StoreLike` object into a Store object. + + `StoreLike` objects are converted to `Store` as follows: + + - `Store` or `StorePath` = `Store` object. + - `Path` or `str` = `LocalStore` object. + - `str` that starts with a protocol = `FsspecStore` object. + - `dict[str, Buffer]` = `MemoryStore` object. + - `None` = `MemoryStore` object. + - `FSMap` = `FsspecStore` object. + + Parameters + ---------- + store_like : StoreLike | None + The `StoreLike` object to convert to a `Store` object. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + mode : StoreAccessMode | None, optional + The mode to use when creating the `Store` object. If None, the + default mode is 'r'. + storage_options : dict[str, Any] | None, optional + The storage options to use when creating the `RemoteStore` object. If + None, the default storage options are used. + + Returns + ------- + Store + The converted Store object. + + Raises + ------ + TypeError + If the StoreLike object is not one of the supported types, or if storage_options is provided but not used. + """ + from zarr.storage._fsspec import FsspecStore # circular import + + if ( + not (isinstance(store_like, str) and _is_fsspec_uri(store_like)) + and storage_options is not None + ): + raise TypeError( + "'storage_options' was provided but unused. " + "'storage_options' is only used when the store is passed as an FSSpec URI string.", + ) + + assert mode in (None, "r", "r+", "a", "w", "w-") + _read_only = mode == "r" + + if isinstance(store_like, StorePath): + # Get underlying store + return store_like.store + + elif isinstance(store_like, Store): + # Already a Store + return store_like + + elif isinstance(store_like, dict): + # Already a dictionary that can be a MemoryStore + # + # We deliberate only consider dict[str, Buffer] here, and not arbitrary mutable mappings. + # By only allowing dictionaries, which are in-memory, we know that MemoryStore appropriate. + return await MemoryStore.open(store_dict=store_like, read_only=_read_only) + + elif store_like is None: + # Create a new in-memory store + return await make_store({}, mode=mode, storage_options=storage_options) + + elif isinstance(store_like, Path): + # Create a new LocalStore + return await LocalStore.open(root=store_like, mode=mode, read_only=_read_only) + + elif isinstance(store_like, str): + # Either an FSSpec URI or a local filesystem path + if _is_fsspec_uri(store_like): + return FsspecStore.from_url( + store_like, storage_options=storage_options, read_only=_read_only + ) + else: + # Assume a filesystem path + return await make_store(Path(store_like), mode=mode, storage_options=storage_options) + + elif _has_fsspec and isinstance(store_like, FSMap): + return FsspecStore.from_mapper(store_like, read_only=_read_only) + + else: + raise TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'") + + +async def make_store_path( + store_like: StoreLike | None, + *, + path: str | None = "", + mode: AccessModeLiteral | None = None, + storage_options: dict[str, Any] | None = None, +) -> StorePath: + """ + Convert a `StoreLike` object into a StorePath object. + + This function takes a `StoreLike` object and returns a `StorePath` object. See `make_store` for details + of which `Store` is used for each type of `store_like` object. + + Parameters + ---------- + store_like : StoreLike or None, default=None + The `StoreLike` object to convert to a `StorePath` object. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + path : str | None, optional + The path to use when creating the `StorePath` object. If None, the + default path is the empty string. + mode : StoreAccessMode | None, optional + The mode to use when creating the `StorePath` object. If None, the + default mode is 'r'. + storage_options : dict[str, Any] | None, optional + The storage options to use when creating the `RemoteStore` object. If + None, the default storage options are used. + + Returns + ------- + StorePath + The converted StorePath object. + + Raises + ------ + TypeError + If the StoreLike object is not one of the supported types, or if storage_options is provided but not used. + ValueError + If path is provided for a store that does not support it. + + See Also + -------- + make_store + """ + path_normalized = normalize_path(path) + + if isinstance(store_like, StorePath): + # Already a StorePath + if storage_options: + raise TypeError( + "'storage_options' was provided but unused. " + "'storage_options' is only used when the store is passed as an FSSpec URI string.", + ) + return store_like / path_normalized + + elif _has_fsspec and isinstance(store_like, FSMap) and path: + raise ValueError( + "'path' was provided but is not used for FSMap store_like objects. Specify the path when creating the FSMap instance instead." + ) + + else: + store = await make_store(store_like, mode=mode, storage_options=storage_options) + return await StorePath.open(store, path=path_normalized, mode=mode) + + +def _is_fsspec_uri(uri: str) -> bool: + """ + Check if a URI looks like a non-local fsspec URI. + + Examples + -------- + ```python + from zarr.storage._common import _is_fsspec_uri + _is_fsspec_uri("s3://bucket") + # True + _is_fsspec_uri("my-directory") + # False + _is_fsspec_uri("local://my-directory") + # False + ``` + """ + return "://" in uri or ("::" in uri and "local://" not in uri) + + +async def ensure_no_existing_node( + store_path: StorePath, + zarr_format: ZarrFormat, + node_type: Literal["array", "group"] | None = None, +) -> None: + """ + Check if a store_path is safe for array / group creation. + Returns `None` or raises an exception. + + Parameters + ---------- + store_path : StorePath + The storage location to check. + zarr_format : ZarrFormat + The Zarr format to check. + node_type : str | None, optional + Raise an error if an "array", or "group" exists. By default (when None), raises an error for either. + + Raises + ------ + ContainsArrayError, ContainsGroupError, ContainsArrayAndGroupError + """ + if zarr_format == 2: + extant_node = await _contains_node_v2(store_path) + elif zarr_format == 3: + extant_node = await _contains_node_v3(store_path) + + match extant_node: + case "array": + if node_type != "group": + msg = f"An array exists in store {store_path.store!r} at path {store_path.path!r}." + raise ContainsArrayError(msg) + + case "group": + if node_type != "array": + msg = f"A group exists in store {store_path.store!r} at path {store_path.path!r}." + raise ContainsGroupError(msg) + + case "nothing": + return + + case _: + msg = f"Invalid value for extant_node: {extant_node}" # type: ignore[unreachable] + raise ValueError(msg) + + +async def _contains_node_v3(store_path: StorePath) -> Literal["array", "group", "nothing"]: + """ + Check if a store_path contains nothing, an array, or a group. This function + returns the string "array", "group", or "nothing" to denote containing an array, a group, or + nothing. + + Parameters + ---------- + store_path : StorePath + The location in storage to check. + + Returns + ------- + Literal["array", "group", "nothing"] + A string representing the zarr node found at store_path. + """ + result: Literal["array", "group", "nothing"] = "nothing" + extant_meta_bytes = await (store_path / ZARR_JSON).get() + # if no metadata document could be loaded, then we just return "nothing" + if extant_meta_bytes is not None: + try: + extant_meta_json = json.loads(extant_meta_bytes.to_bytes()) + # avoid constructing a full metadata document here in the name of speed. + if extant_meta_json["node_type"] == "array": + result = "array" + elif extant_meta_json["node_type"] == "group": + result = "group" + except (KeyError, json.JSONDecodeError): + # either of these errors is consistent with no array or group present. + pass + return result + + +async def _contains_node_v2(store_path: StorePath) -> Literal["array", "group", "nothing"]: + """ + Check if a store_path contains nothing, an array, a group, or both. If both an array and a + group are detected, a `ContainsArrayAndGroup` exception is raised. Otherwise, this function + returns the string "array", "group", or "nothing" to denote containing an array, a group, or + nothing. + + Parameters + ---------- + store_path : StorePath + The location in storage to check. + + Returns + ------- + Literal["array", "group", "nothing"] + A string representing the zarr node found at store_path. + """ + _array = await contains_array(store_path=store_path, zarr_format=2) + _group = await contains_group(store_path=store_path, zarr_format=2) + + if _array and _group: + msg = ( + "Array and group metadata documents (.zarray and .zgroup) were both found in store " + f"{store_path.store!r} at path {store_path.path!r}. " + "Only one of these files may be present in a given directory / prefix. " + "Remove the .zarray file, or the .zgroup file, or both." + ) + raise ContainsArrayAndGroupError(msg) + elif _array: + return "array" + elif _group: + return "group" + else: + return "nothing" + + +async def contains_array(store_path: StorePath, zarr_format: ZarrFormat) -> bool: + """ + Check if an array exists at a given StorePath. + + Parameters + ---------- + store_path : StorePath + The StorePath to check for an existing group. + zarr_format : + The zarr format to check for. + + Returns + ------- + bool + True if the StorePath contains a group, False otherwise. + + """ + if zarr_format == 3: + extant_meta_bytes = await (store_path / ZARR_JSON).get() + if extant_meta_bytes is None: + return False + else: + try: + extant_meta_json = json.loads(extant_meta_bytes.to_bytes()) + # we avoid constructing a full metadata document here in the name of speed. + if extant_meta_json["node_type"] == "array": + return True + except (ValueError, KeyError): + return False + elif zarr_format == 2: + return await (store_path / ZARRAY_JSON).exists() + msg = f"Invalid zarr_format provided. Got {zarr_format}, expected 2 or 3" + raise ValueError(msg) + + +async def contains_group(store_path: StorePath, zarr_format: ZarrFormat) -> bool: + """ + Check if a group exists at a given StorePath. + + Parameters + ---------- + + store_path : StorePath + The StorePath to check for an existing group. + zarr_format : + The zarr format to check for. + + Returns + ------- + + bool + True if the StorePath contains a group, False otherwise + + """ + if zarr_format == 3: + extant_meta_bytes = await (store_path / ZARR_JSON).get() + if extant_meta_bytes is None: + return False + else: + try: + extant_meta_json = json.loads(extant_meta_bytes.to_bytes()) + # we avoid constructing a full metadata document here in the name of speed. + result: bool = extant_meta_json["node_type"] == "group" + except (ValueError, KeyError): + return False + else: + return result + elif zarr_format == 2: + return await (store_path / ZGROUP_JSON).exists() + msg = f"Invalid zarr_format provided. Got {zarr_format}, expected 2 or 3" # type: ignore[unreachable] + raise ValueError(msg) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_fsspec.py b/.venv/lib/python3.11/site-packages/zarr/storage/_fsspec.py new file mode 100644 index 0000000000000000000000000000000000000000..f9e4ed375dd862c79f0334dfec55bd50f5587279 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_fsspec.py @@ -0,0 +1,442 @@ +from __future__ import annotations + +import json +import warnings +from contextlib import suppress +from typing import TYPE_CHECKING, Any + +from packaging.version import parse as parse_version + +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr.core.buffer import Buffer +from zarr.errors import ZarrUserWarning +from zarr.storage._common import _dereference_path + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable + + from fsspec import AbstractFileSystem + from fsspec.asyn import AsyncFileSystem + from fsspec.mapping import FSMap + + from zarr.core.buffer import BufferPrototype + + +ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( + FileNotFoundError, + IsADirectoryError, + NotADirectoryError, +) + + +def _make_async(fs: AbstractFileSystem) -> AsyncFileSystem: + """Convert a sync FSSpec filesystem to an async FFSpec filesystem + + If the filesystem class supports async operations, a new async instance is created + from the existing instance. + + If the filesystem class does not support async operations, the existing instance + is wrapped with AsyncFileSystemWrapper. + """ + import fsspec + + fsspec_version = parse_version(fsspec.__version__) + if fs.async_impl and fs.asynchronous: + # Already an async instance of an async filesystem, nothing to do + return fs + if fs.async_impl: + # Convert sync instance of an async fs to an async instance + fs_dict = json.loads(fs.to_json()) + fs_dict["asynchronous"] = True + return fsspec.AbstractFileSystem.from_json(json.dumps(fs_dict)) + + if fsspec_version < parse_version("2024.12.0"): + raise ImportError( + f"The filesystem '{fs}' is synchronous, and the required " + "AsyncFileSystemWrapper is not available. Upgrade fsspec to version " + "2024.12.0 or later to enable this functionality." + ) + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + + return AsyncFileSystemWrapper(fs, asynchronous=True) + + +class FsspecStore(Store): + """ + Store for remote data based on FSSpec. + + Parameters + ---------- + fs : AsyncFileSystem + The Async FSSpec filesystem to use with this store. + read_only : bool + Whether the store is read-only + path : str + The root path of the store. This should be a relative path and must not include the + filesystem scheme. + allowed_exceptions : tuple[type[Exception], ...] + When fetching data, these cases will be deemed to correspond to missing keys. + + Attributes + ---------- + fs + allowed_exceptions + supports_writes + supports_deletes + supports_listing + + Raises + ------ + TypeError + If the Filesystem does not support async operations. + ValueError + If the path argument includes a scheme. + + Warns + ----- + ZarrUserWarning + If the file system (fs) was not created with `asynchronous=True`. + + See Also + -------- + FsspecStore.from_upath + FsspecStore.from_url + """ + + # based on FSSpec + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + fs: AsyncFileSystem + allowed_exceptions: tuple[type[Exception], ...] + path: str + + def __init__( + self, + fs: AsyncFileSystem, + read_only: bool = False, + path: str = "/", + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> None: + super().__init__(read_only=read_only) + self.fs = fs + self.path = path + self.allowed_exceptions = allowed_exceptions + + if not self.fs.async_impl: + raise TypeError("Filesystem needs to support async operations.") + if not self.fs.asynchronous: + warnings.warn( + f"fs ({fs}) was not created with `asynchronous=True`, this may lead to surprising behavior", + category=ZarrUserWarning, + stacklevel=2, + ) + + @classmethod + def from_upath( + cls, + upath: Any, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from a upath object. + + Parameters + ---------- + upath : UPath + The upath to the root of the store. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + return cls( + fs=upath.fs, + path=upath.path.rstrip("/"), + read_only=read_only, + allowed_exceptions=allowed_exceptions, + ) + + @classmethod + def from_mapper( + cls, + fs_map: FSMap, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from an FSMap object. + + Parameters + ---------- + fs_map : FSMap + Fsspec mutable mapping object. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + fs = _make_async(fs_map.fs) + return cls( + fs=fs, + path=fs_map.root, + read_only=read_only, + allowed_exceptions=allowed_exceptions, + ) + + @classmethod + def from_url( + cls, + url: str, + storage_options: dict[str, Any] | None = None, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from a URL. The type of store is determined from the URL scheme. + + Parameters + ---------- + url : str + The URL to the root of the store. + storage_options : dict, optional + The options to pass to fsspec when creating the filesystem. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + try: + from fsspec import url_to_fs + except ImportError: + # before fsspec==2024.3.1 + from fsspec.core import url_to_fs + + opts = storage_options or {} + opts = {"asynchronous": True, **opts} + + fs, path = url_to_fs(url, **opts) + if not fs.async_impl: + fs = _make_async(fs) + + return cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions) + + def with_read_only(self, read_only: bool = False) -> FsspecStore: + # docstring inherited + return type(self)( + fs=self.fs, + path=self.path, + allowed_exceptions=self.allowed_exceptions, + read_only=read_only, + ) + + async def clear(self) -> None: + # docstring inherited + try: + for subpath in await self.fs._find(self.path, withdirs=True): + if subpath != self.path: + await self.fs._rm(subpath, recursive=True) + except FileNotFoundError: + pass + + def __repr__(self) -> str: + return f"" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self.path == other.path + and self.read_only == other.read_only + and self.fs == other.fs + ) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if not self._is_open: + await self._open() + path = _dereference_path(self.path, key) + + try: + if byte_range is None: + value = prototype.buffer.from_bytes(await self.fs._cat_file(path)) + elif isinstance(byte_range, RangeByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file( + path, + start=byte_range.start, + end=byte_range.end, + ) + ) + elif isinstance(byte_range, OffsetByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file(path, start=byte_range.offset, end=None) + ) + elif isinstance(byte_range, SuffixByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file(path, start=-byte_range.suffix, end=None) + ) + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + except self.allowed_exceptions: + return None + except OSError as e: + if "not satisfiable" in str(e): + # this is an s3-specific condition we probably don't want to leak + return prototype.buffer.from_bytes(b"") + raise + else: + return value + + async def set( + self, + key: str, + value: Buffer, + byte_range: tuple[int, int] | None = None, + ) -> None: + # docstring inherited + if not self._is_open: + await self._open() + self._check_writable() + if not isinstance(value, Buffer): + raise TypeError( + f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + path = _dereference_path(self.path, key) + # write data + if byte_range: + raise NotImplementedError + await self.fs._pipe_file(path, value.to_bytes()) + + async def delete(self, key: str) -> None: + # docstring inherited + self._check_writable() + path = _dereference_path(self.path, key) + try: + await self.fs._rm(path) + except FileNotFoundError: + pass + except self.allowed_exceptions: + pass + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + if not self.supports_deletes: + raise NotImplementedError( + "This method is only available for stores that support deletes." + ) + self._check_writable() + + path_to_delete = _dereference_path(self.path, prefix) + + with suppress(*self.allowed_exceptions): + await self.fs._rm(path_to_delete, recursive=True) + + async def exists(self, key: str) -> bool: + # docstring inherited + path = _dereference_path(self.path, key) + exists: bool = await self.fs._exists(path) + return exists + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + if key_ranges: + # _cat_ranges expects a list of paths, start, and end ranges, so we need to reformat each ByteRequest. + key_ranges = list(key_ranges) + paths: list[str] = [] + starts: list[int | None] = [] + stops: list[int | None] = [] + for key, byte_range in key_ranges: + paths.append(_dereference_path(self.path, key)) + if byte_range is None: + starts.append(None) + stops.append(None) + elif isinstance(byte_range, RangeByteRequest): + starts.append(byte_range.start) + stops.append(byte_range.end) + elif isinstance(byte_range, OffsetByteRequest): + starts.append(byte_range.offset) + stops.append(None) + elif isinstance(byte_range, SuffixByteRequest): + starts.append(-byte_range.suffix) + stops.append(None) + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + else: + return [] + # TODO: expectations for exceptions or missing keys? + res = await self.fs._cat_ranges(paths, starts, stops, on_error="return") + # the following is an s3-specific condition we probably don't want to leak + res = [b"" if (isinstance(r, OSError) and "not satisfiable" in str(r)) else r for r in res] + for r in res: + if isinstance(r, Exception) and not isinstance(r, self.allowed_exceptions): + raise r + + return [None if isinstance(r, Exception) else prototype.buffer.from_bytes(r) for r in res] + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + allfiles = await self.fs._find(self.path, detail=False, withdirs=False) + for onefile in (a.removeprefix(self.path + "/") for a in allfiles): + yield onefile + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + prefix = f"{self.path}/{prefix.rstrip('/')}" + try: + allfiles = await self.fs._ls(prefix, detail=False) + except FileNotFoundError: + return + for onefile in (a.replace(prefix + "/", "") for a in allfiles): + yield onefile.removeprefix(self.path).removeprefix("/") + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + for onefile in await self.fs._find( + f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False + ): + yield onefile.removeprefix(f"{self.path}/") + + async def getsize(self, key: str) -> int: + path = _dereference_path(self.path, key) + info = await self.fs._info(path) + + size = info.get("size") + + if size is None: + # Not all filesystems support size. Fall back to reading the entire object + return await super().getsize(key) + else: + # fsspec doesn't have typing. We'll need to assume or verify this is true + return int(size) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_local.py b/.venv/lib/python3.11/site-packages/zarr/storage/_local.py new file mode 100644 index 0000000000000000000000000000000000000000..96f1e617469927f7474b3715b18d507c63176547 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_local.py @@ -0,0 +1,602 @@ +from __future__ import annotations + +import asyncio +import contextlib +import io +import os +import shutil +import sys +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any, BinaryIO, Literal, Self + +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr.core.buffer import Buffer +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.common import AccessModeLiteral, concurrent_map + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, Iterator + + from zarr.core.buffer import BufferPrototype + + +def _get(path: Path, prototype: BufferPrototype, byte_range: ByteRequest | None) -> Buffer: + if byte_range is None: + return prototype.buffer.from_bytes(path.read_bytes()) + with path.open("rb") as f: + size = f.seek(0, io.SEEK_END) + if isinstance(byte_range, RangeByteRequest): + f.seek(byte_range.start) + return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell())) + elif isinstance(byte_range, OffsetByteRequest): + f.seek(byte_range.offset) + elif isinstance(byte_range, SuffixByteRequest): + f.seek(max(0, size - byte_range.suffix)) + else: + raise TypeError(f"Unexpected byte_range, got {byte_range}.") + return prototype.buffer.from_bytes(f.read()) + + +if sys.platform == "win32": + # Per the os.rename docs: + # On Windows, if dst exists a FileExistsError is always raised. + _safe_move = os.rename +else: + # On Unix, os.rename silently replace files, so instead we use os.link like + # atomicwrites: + # https://github.com/untitaker/python-atomicwrites/blob/1.4.1/atomicwrites/__init__.py#L59-L60 + # This also raises FileExistsError if dst exists. + def _safe_move(src: Path, dst: Path) -> None: + os.link(src, dst) + os.unlink(src) + + +@contextlib.contextmanager +def _atomic_write( + path: Path, + mode: Literal["r+b", "wb"], + exclusive: bool = False, +) -> Iterator[BinaryIO]: + tmp_path = path.with_suffix(f".{uuid.uuid4().hex}.partial") + try: + with tmp_path.open(mode) as f: + yield f + if exclusive: + _safe_move(tmp_path, path) + else: + tmp_path.replace(path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def _put(path: Path, value: Buffer, exclusive: bool = False) -> int: + path.parent.mkdir(parents=True, exist_ok=True) + # write takes any object supporting the buffer protocol + view = value.as_buffer_like() + with _atomic_write(path, "wb", exclusive=exclusive) as f: + return f.write(view) + + +class LocalStore(Store): + """ + Store for the local file system. + + Parameters + ---------- + root : str or Path + Directory to use as root of store. + read_only : bool + Whether the store is read-only + + Attributes + ---------- + supports_writes + supports_deletes + supports_listing + root + """ + + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + root: Path + + def __init__(self, root: Path | str, *, read_only: bool = False) -> None: + super().__init__(read_only=read_only) + if isinstance(root, str): + root = Path(root) + if not isinstance(root, Path): + raise TypeError( + f"'root' must be a string or Path instance. Got an instance of {type(root)} instead." + ) + self.root = root + + def with_read_only(self, read_only: bool = False) -> Self: + # docstring inherited + return type(self)( + root=self.root, + read_only=read_only, + ) + + @classmethod + async def open( + cls, root: Path | str, *, read_only: bool = False, mode: AccessModeLiteral | None = None + ) -> Self: + """ + Create and open the store. + + Parameters + ---------- + root : str or Path + Directory to use as root of store. + read_only : bool + Whether the store is read-only + mode : + Mode in which to create the store. This only affects opening the store, + and the final read-only state of the store is controlled through the + read_only parameter. + + Returns + ------- + Store + The opened store instance. + """ + # If mode = 'r+', want to open in read only mode (fail if exists), + # but return a writeable store + if mode is not None: + read_only_creation = mode in ["r", "r+"] + else: + read_only_creation = read_only + store = cls(root, read_only=read_only_creation) + await store._open() + + # Set read_only state + store = store.with_read_only(read_only) + await store._open() + return store + + async def _open(self, *, mode: AccessModeLiteral | None = None) -> None: + if not self.read_only: + self.root.mkdir(parents=True, exist_ok=True) + + if not self.root.exists(): + raise FileNotFoundError(f"{self.root} does not exist") + return await super()._open() + + async def clear(self) -> None: + # docstring inherited + self._check_writable() + shutil.rmtree(self.root) + self.root.mkdir() + + def __str__(self) -> str: + return f"file://{self.root.as_posix()}" + + def __repr__(self) -> str: + return f"LocalStore('{self}')" + + def __eq__(self, other: object) -> bool: + return isinstance(other, type(self)) and self.root == other.root + + # ------------------------------------------------------------------- + # Synchronous store methods + # ------------------------------------------------------------------- + + def _ensure_open_sync(self) -> None: + if not self._is_open: + if not self.read_only: + self.root.mkdir(parents=True, exist_ok=True) + if not self.root.exists(): + raise FileNotFoundError(f"{self.root} does not exist") + self._is_open = True + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if prototype is None: + prototype = default_buffer_prototype() + self._ensure_open_sync() + assert isinstance(key, str) + path = self.root / key + try: + return _get(path, prototype, byte_range) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + return None + + def set_sync(self, key: str, value: Buffer) -> None: + self._ensure_open_sync() + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"LocalStore.set(): `value` must be a Buffer instance. " + f"Got an instance of {type(value)} instead." + ) + path = self.root / key + _put(path, value) + + def delete_sync(self, key: str) -> None: + self._ensure_open_sync() + self._check_writable() + path = self.root / key + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + await self._open() + assert isinstance(key, str) + path = self.root / key + + try: + return await asyncio.to_thread(_get, path, prototype, byte_range) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + args = [] + for key, byte_range in key_ranges: + assert isinstance(key, str) + path = self.root / key + args.append((_get, path, prototype, byte_range)) + return await concurrent_map(args, asyncio.to_thread, limit=None) # TODO: fix limit + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + return await self._set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + try: + return await self._set(key, value, exclusive=True) + except FileExistsError: + pass + + async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None: + if not self._is_open: + await self._open() + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"LocalStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + path = self.root / key + await asyncio.to_thread(_put, path, value, exclusive=exclusive) + + async def delete(self, key: str) -> None: + """ + Remove a key from the store. + + Parameters + ---------- + key : str + + Notes + ----- + If ``key`` is a directory within this store, the entire directory + at ``store.root / key`` is deleted. + """ + # docstring inherited + self._check_writable() + path = self.root / key + if path.is_dir(): # TODO: support deleting directories? shutil.rmtree? + shutil.rmtree(path) + else: + await asyncio.to_thread(path.unlink, True) # Q: we may want to raise if path is missing + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + self._check_writable() + path = self.root / prefix + if path.is_dir(): + shutil.rmtree(path) + elif path.is_file(): + raise ValueError(f"delete_dir was passed a {prefix=!r} that is a file") + else: + # Non-existent directory + # This path is tested by test_group:test_create_creates_parents for one + pass + + async def exists(self, key: str) -> bool: + # docstring inherited + path = self.root / key + return await asyncio.to_thread(path.is_file) + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + to_strip = self.root.as_posix() + "/" + for p in list(self.root.rglob("*")): + if p.is_file(): + yield p.as_posix().replace(to_strip, "") + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + to_strip = self.root.as_posix() + "/" + prefix = prefix.rstrip("/") + for p in (self.root / prefix).rglob("*"): + if p.is_file(): + yield p.as_posix().replace(to_strip, "") + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + base = self.root / prefix + try: + key_iter = base.iterdir() + for key in key_iter: + yield key.relative_to(base).as_posix() + except (FileNotFoundError, NotADirectoryError): + pass + + async def _get_bytes( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the local store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + See Also + -------- + Store.get_bytes : Base implementation with full documentation. + get_bytes_sync : Synchronous version of this method. + + Examples + -------- + >>> store = await LocalStore.open("data") + >>> await store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for LocalStore + >>> data = await store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super()._get_bytes(key, prototype=prototype, byte_range=byte_range) + + def _get_bytes_sync( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the local store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + Warnings + -------- + Do not call this method from async functions. Use ``get_bytes()`` instead. + + See Also + -------- + Store.get_bytes_sync : Base implementation with full documentation. + get_bytes : Asynchronous version of this method. + + Examples + -------- + >>> store = LocalStore("data") + >>> store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for LocalStore + >>> data = store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return super()._get_bytes_sync(key, prototype=prototype, byte_range=byte_range) + + async def _get_json( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the local store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + Store.get_json : Base implementation with full documentation. + get_json_sync : Synchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = await LocalStore.open("data") + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for LocalStore + >>> data = await store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super()._get_json(key, prototype=prototype, byte_range=byte_range) + + def _get_json_sync( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the local store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + Warnings + -------- + Do not call this method from async functions. Use ``get_json()`` instead. + + See Also + -------- + Store.get_json_sync : Base implementation with full documentation. + get_json : Asynchronous version of this method. + get_bytes_sync : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = LocalStore("data") + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for LocalStore + >>> data = store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return super()._get_json_sync(key, prototype=prototype, byte_range=byte_range) + + async def move(self, dest_root: Path | str) -> None: + """ + Move the store to another path. The old root directory is deleted. + """ + if isinstance(dest_root, str): + dest_root = Path(dest_root) + os.makedirs(dest_root.parent, exist_ok=True) + if os.path.exists(dest_root): + raise FileExistsError(f"Destination root {dest_root} already exists.") + shutil.move(self.root, dest_root) + self.root = dest_root + + async def getsize(self, key: str) -> int: + return os.path.getsize(self.root / key) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_logging.py b/.venv/lib/python3.11/site-packages/zarr/storage/_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..98dca6b23d1b3226518002215344bce986883566 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_logging.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import inspect +import logging +import sys +import time +from collections import defaultdict +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Self, TypeVar + +from zarr.abc.store import Store +from zarr.storage._wrapper import WrapperStore + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Generator, Iterable + + from zarr.abc.store import ByteRequest + from zarr.core.buffer import Buffer, BufferPrototype + + counter: defaultdict[str, int] + +T_Store = TypeVar("T_Store", bound=Store) + + +class LoggingStore(WrapperStore[T_Store]): + """ + Store that logs all calls to another wrapped store. + + Parameters + ---------- + store : Store + Store to wrap + log_level : str + Log level + log_handler : logging.Handler + Log handler + + Attributes + ---------- + counter : dict + Counter of number of times each method has been called + """ + + counter: defaultdict[str, int] + + def __init__( + self, + store: T_Store, + log_level: str = "DEBUG", + log_handler: logging.Handler | None = None, + ) -> None: + super().__init__(store) + self.counter = defaultdict(int) + self.log_level = log_level + self.log_handler = log_handler + self._configure_logger(log_level, log_handler) + + def _configure_logger( + self, log_level: str = "DEBUG", log_handler: logging.Handler | None = None + ) -> None: + self.log_level = log_level + self.logger = logging.getLogger(f"LoggingStore({self._store})") + self.logger.setLevel(log_level) + + if not self.logger.hasHandlers(): + if not log_handler: + log_handler = self._default_handler() + # Add handler to logger + self.logger.addHandler(log_handler) + + def _default_handler(self) -> logging.Handler: + """Define a default log handler""" + handler = logging.StreamHandler(stream=sys.stdout) + handler.setLevel(self.log_level) + handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) + return handler + + def _with_store(self, store: T_Store) -> Self: + return type(self)(store=store, log_level=self.log_level, log_handler=self.log_handler) + + @contextmanager + def log(self, hint: Any = "") -> Generator[None, None, None]: + """Context manager to log method calls + + Each call to the wrapped store is logged to the configured logger and added to + the counter dict. + """ + method = inspect.stack()[2].function + op = f"{type(self._store).__name__}.{method}" + if hint: + op = f"{op}({hint})" + self.logger.info(" Calling %s", op) + start_time = time.time() + try: + self.counter[method] += 1 + yield + finally: + end_time = time.time() + self.logger.info("Finished %s [%.2f s]", op, end_time - start_time) + + @classmethod + async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self: + log_level = kwargs.pop("log_level", "DEBUG") + log_handler = kwargs.pop("log_handler", None) + store = store_cls(*args, **kwargs) + await store._open() + return cls(store=store, log_level=log_level, log_handler=log_handler) + + @property + def supports_writes(self) -> bool: + with self.log(): + return self._store.supports_writes + + @property + def supports_deletes(self) -> bool: + with self.log(): + return self._store.supports_deletes + + @property + def supports_listing(self) -> bool: + with self.log(): + return self._store.supports_listing + + @property + def read_only(self) -> bool: + with self.log(): + return self._store.read_only + + @property + def _is_open(self) -> bool: + with self.log(): + return self._store._is_open + + @_is_open.setter + def _is_open(self, value: bool) -> None: + raise NotImplementedError("LoggingStore must be opened via the `_open` method") + + async def _open(self) -> None: + with self.log(): + return await self._store._open() + + async def _ensure_open(self) -> None: + with self.log(): + return await self._store._ensure_open() + + async def is_empty(self, prefix: str = "") -> bool: + # docstring inherited + with self.log(): + return await self._store.is_empty(prefix=prefix) + + async def clear(self) -> None: + # docstring inherited + with self.log(): + return await self._store.clear() + + def __str__(self) -> str: + return f"logging-{self._store}" + + def __repr__(self) -> str: + return f"LoggingStore({self._store.__class__.__name__}, '{self._store}')" + + def __eq__(self, other: object) -> bool: + with self.log(other): + return type(self) is type(other) and self._store.__eq__(other._store) # type: ignore[attr-defined] + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + with self.log(key): + return await self._store.get(key=key, prototype=prototype, byte_range=byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + keys = ",".join([k[0] for k in key_ranges]) + with self.log(keys): + return await self._store.get_partial_values(prototype=prototype, key_ranges=key_ranges) + + async def exists(self, key: str) -> bool: + # docstring inherited + with self.log(key): + return await self._store.exists(key) + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return await self._store.set(key=key, value=value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return await self._store.set_if_not_exists(key=key, value=value) + + async def delete(self, key: str) -> None: + # docstring inherited + with self.log(key): + return await self._store.delete(key=key) + + async def list(self) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(): + async for key in self._store.list(): + yield key + + async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(prefix): + async for key in self._store.list_prefix(prefix=prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(prefix): + async for key in self._store.list_dir(prefix=prefix): + yield key + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + with self.log(prefix): + await self._store.delete_dir(prefix=prefix) + + async def getsize(self, key: str) -> int: + with self.log(key): + return await self._store.getsize(key) + + async def getsize_prefix(self, prefix: str) -> int: + with self.log(prefix): + return await self._store.getsize_prefix(prefix) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_memory.py b/.venv/lib/python3.11/site-packages/zarr/storage/_memory.py new file mode 100644 index 0000000000000000000000000000000000000000..1194894b9d0766aa9ae567b43e0ae93f146ea0e8 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_memory.py @@ -0,0 +1,519 @@ +from __future__ import annotations + +from logging import getLogger +from typing import TYPE_CHECKING, Any, Self + +from zarr.abc.store import ByteRequest, Store +from zarr.core.buffer import Buffer, gpu +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.common import concurrent_map +from zarr.storage._utils import _normalize_byte_range_index + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, MutableMapping + + from zarr.core.buffer import BufferPrototype + + +logger = getLogger(__name__) + + +class MemoryStore(Store): + """ + Store for local memory. + + Parameters + ---------- + store_dict : dict + Initial data + read_only : bool + Whether the store is read-only + + Attributes + ---------- + supports_writes + supports_deletes + supports_listing + """ + + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + _store_dict: MutableMapping[str, Buffer] + + def __init__( + self, + store_dict: MutableMapping[str, Buffer] | None = None, + *, + read_only: bool = False, + ) -> None: + super().__init__(read_only=read_only) + if store_dict is None: + store_dict = {} + self._store_dict = store_dict + + def with_read_only(self, read_only: bool = False) -> MemoryStore: + # docstring inherited + return type(self)( + store_dict=self._store_dict, + read_only=read_only, + ) + + async def clear(self) -> None: + # docstring inherited + self._store_dict.clear() + + def __str__(self) -> str: + return f"memory://{id(self._store_dict)}" + + def __repr__(self) -> str: + return f"MemoryStore('{self}')" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self._store_dict == other._store_dict + and self.read_only == other.read_only + ) + + # ------------------------------------------------------------------- + # Synchronous store methods + # ------------------------------------------------------------------- + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + self._is_open = True + assert isinstance(key, str) + try: + value = self._store_dict[key] + start, stop = _normalize_byte_range_index(value, byte_range) + return prototype.buffer.from_buffer(value[start:stop]) + except KeyError: + return None + + def set_sync(self, key: str, value: Buffer) -> None: + self._check_writable() + if not self._is_open: + self._is_open = True + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + self._store_dict[key] = value + + def delete_sync(self, key: str) -> None: + self._check_writable() + if not self._is_open: + self._is_open = True + try: + del self._store_dict[key] + except KeyError: + logger.debug("Key %s does not exist.", key) + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + await self._open() + assert isinstance(key, str) + try: + value = self._store_dict[key] + start, stop = _normalize_byte_range_index(value, byte_range) + return prototype.buffer.from_buffer(value[start:stop]) + except KeyError: + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + + # All the key-ranges arguments goes with the same prototype + async def _get(key: str, byte_range: ByteRequest | None) -> Buffer | None: + return await self.get(key, prototype=prototype, byte_range=byte_range) + + return await concurrent_map(key_ranges, _get, limit=None) + + async def exists(self, key: str) -> bool: + # docstring inherited + return key in self._store_dict + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + # docstring inherited + self._check_writable() + await self._ensure_open() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + if byte_range is not None: + buf = self._store_dict[key] + buf[byte_range[0] : byte_range[1]] = value + self._store_dict[key] = buf + else: + self._store_dict[key] = value + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + self._check_writable() + await self._ensure_open() + self._store_dict.setdefault(key, value) + + async def delete(self, key: str) -> None: + # docstring inherited + self._check_writable() + try: + del self._store_dict[key] + except KeyError: + logger.debug("Key %s does not exist.", key) + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + for key in self._store_dict: + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix) + for key in list(self._store_dict): + if key.startswith(prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + prefix = prefix.rstrip("/") + + if prefix == "": + keys_unique = {k.split("/")[0] for k in self._store_dict} + else: + # Our dictionary doesn't contain directory markers, but we want to include + # a pseudo directory when there's a nested item and we're listing an + # intermediate level. + keys_unique = { + key.removeprefix(prefix + "/").split("/")[0] + for key in self._store_dict + if key.startswith(prefix + "/") and key != prefix + } + + for key in keys_unique: + yield key + + async def _get_bytes( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the memory store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + See Also + -------- + Store.get_bytes : Base implementation with full documentation. + get_bytes_sync : Synchronous version of this method. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> await store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for MemoryStore + >>> data = await store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super()._get_bytes(key, prototype=prototype, byte_range=byte_range) + + def _get_bytes_sync( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the memory store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + Warnings + -------- + Do not call this method from async functions. Use ``get_bytes()`` instead. + + See Also + -------- + Store.get_bytes_sync : Base implementation with full documentation. + get_bytes : Asynchronous version of this method. + + Examples + -------- + >>> store = MemoryStore() + >>> store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for MemoryStore + >>> data = store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return super()._get_bytes_sync(key, prototype=prototype, byte_range=byte_range) + + async def _get_json( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the memory store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + Store.get_json : Base implementation with full documentation. + get_json_sync : Synchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for MemoryStore + >>> data = await store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super()._get_json(key, prototype=prototype, byte_range=byte_range) + + def _get_json_sync( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the memory store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + Warnings + -------- + Do not call this method from async functions. Use ``get_json()`` instead. + + See Also + -------- + Store.get_json_sync : Base implementation with full documentation. + get_json : Asynchronous version of this method. + get_bytes_sync : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = MemoryStore() + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for MemoryStore + >>> data = store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return super()._get_json_sync(key, prototype=prototype, byte_range=byte_range) + + +class GpuMemoryStore(MemoryStore): + """ + Store for GPU memory. + + Stores every chunk in GPU memory irrespective of the original location. + + The dictionary of buffers to initialize this memory store with *must* be + GPU Buffers. + + Writing data to this store through ``.set`` will move the buffer to the GPU + if necessary. + + Parameters + ---------- + store_dict : MutableMapping, optional + A mutable mapping with string keys and [zarr.core.buffer.gpu.Buffer][] + values. + read_only : bool + Whether to open the store in read-only mode. + """ + + _store_dict: MutableMapping[str, gpu.Buffer] # type: ignore[assignment] + + def __init__( + self, + store_dict: MutableMapping[str, gpu.Buffer] | None = None, + *, + read_only: bool = False, + ) -> None: + super().__init__(store_dict=store_dict, read_only=read_only) # type: ignore[arg-type] + + def __str__(self) -> str: + return f"gpumemory://{id(self._store_dict)}" + + def __repr__(self) -> str: + return f"GpuMemoryStore('{self}')" + + @classmethod + def from_dict(cls, store_dict: MutableMapping[str, Buffer]) -> Self: + """ + Create a GpuMemoryStore from a dictionary of buffers at any location. + + The dictionary backing the newly created ``GpuMemoryStore`` will not be + the same as ``store_dict``. + + Parameters + ---------- + store_dict : mapping + A mapping of strings keys to arbitrary Buffers. The buffer data + will be moved into a [`gpu.Buffer`][zarr.core.buffer.gpu.Buffer]. + + Returns + ------- + GpuMemoryStore + """ + gpu_store_dict = {k: gpu.Buffer.from_buffer(v) for k, v in store_dict.items()} + return cls(gpu_store_dict) + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + # docstring inherited + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + # Convert to gpu.Buffer + gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value) + await super().set(key, gpu_value, byte_range=byte_range) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_obstore.py b/.venv/lib/python3.11/site-packages/zarr/storage/_obstore.py new file mode 100644 index 0000000000000000000000000000000000000000..6e4011da5953a1fdc2a2256f4ec805a4099f0d79 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_obstore.py @@ -0,0 +1,512 @@ +from __future__ import annotations + +import asyncio +import contextlib +import pickle +from collections import defaultdict +from itertools import chain +from operator import itemgetter +from typing import TYPE_CHECKING, Generic, Self, TypedDict, TypeVar + +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr.core.common import concurrent_map +from zarr.core.config import config +from zarr.storage._utils import _relativize_path + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence + from typing import Any + + from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange + from obstore.store import ObjectStore as _UpstreamObjectStore + + from zarr.core.buffer import Buffer, BufferPrototype + +__all__ = ["ObjectStore"] + +_ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( + FileNotFoundError, + IsADirectoryError, + NotADirectoryError, +) + + +T_Store = TypeVar("T_Store", bound="_UpstreamObjectStore") + + +class ObjectStore(Store, Generic[T_Store]): + """ + Store that uses obstore for fast read/write from AWS, GCP, Azure. + + Parameters + ---------- + store : obstore.store.ObjectStore + An obstore store instance that is set up with the proper credentials. + read_only : bool + Whether to open the store in read-only mode. + + Warnings + -------- + ObjectStore is experimental and subject to API changes without notice. Please + raise an issue with any comments/concerns about the store. + """ + + store: T_Store + """The underlying obstore instance.""" + + def __eq__(self, value: object) -> bool: + if not isinstance(value, ObjectStore): + return False + + if not self.read_only == value.read_only: + return False + + return self.store == value.store # type: ignore[no-any-return] + + def __init__(self, store: T_Store, *, read_only: bool = False) -> None: + if not store.__class__.__module__.startswith("obstore"): + raise TypeError(f"expected ObjectStore class, got {store!r}") + super().__init__(read_only=read_only) + self.store = store + + def with_read_only(self, read_only: bool = False) -> Self: + # docstring inherited + return type(self)( + store=self.store, + read_only=read_only, + ) + + def __str__(self) -> str: + return f"object_store://{self.store}" + + def __repr__(self) -> str: + return f"{type(self).__name__}({self})" + + def __getstate__(self) -> dict[Any, Any]: + state = self.__dict__.copy() + state["store"] = pickle.dumps(self.store) + return state + + def __setstate__(self, state: dict[Any, Any]) -> None: + state["store"] = pickle.loads(state["store"]) + self.__dict__.update(state) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + # docstring inherited + import obstore as obs + + try: + if byte_range is None: + resp = await obs.get_async(self.store, key) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + elif isinstance(byte_range, RangeByteRequest): + bytes = await obs.get_range_async( + self.store, key, start=byte_range.start, end=byte_range.end + ) + return prototype.buffer.from_bytes(bytes) # type: ignore[arg-type] + elif isinstance(byte_range, OffsetByteRequest): + resp = await obs.get_async( + self.store, key, options={"range": {"offset": byte_range.offset}} + ) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + elif isinstance(byte_range, SuffixByteRequest): + # some object stores (Azure) don't support suffix requests. In this + # case, our workaround is to first get the length of the object and then + # manually request the byte range at the end. + try: + resp = await obs.get_async( + self.store, key, options={"range": {"suffix": byte_range.suffix}} + ) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + except obs.exceptions.NotSupportedError: + head_resp = await obs.head_async(self.store, key) + file_size = head_resp["size"] + suffix_len = byte_range.suffix + buffer = await obs.get_range_async( + self.store, + key, + start=file_size - suffix_len, + length=suffix_len, + ) + return prototype.buffer.from_bytes(buffer) # type: ignore[arg-type] + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}") + except _ALLOWED_EXCEPTIONS: + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + return await _get_partial_values(self.store, prototype=prototype, key_ranges=key_ranges) + + async def exists(self, key: str) -> bool: + # docstring inherited + import obstore as obs + + try: + await obs.head_async(self.store, key) + except FileNotFoundError: + return False + else: + return True + + @property + def supports_writes(self) -> bool: + # docstring inherited + return True + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + + buf = value.as_buffer_like() + await obs.put_async(self.store, key, buf) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + buf = value.as_buffer_like() + with contextlib.suppress(obs.exceptions.AlreadyExistsError): + await obs.put_async(self.store, key, buf, mode="create") + + @property + def supports_deletes(self) -> bool: + # docstring inherited + return True + + async def delete(self, key: str) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + + # Some obstore stores such as local filesystems, GCP and Azure raise an error + # when deleting a non-existent key, while others such as S3 and in-memory do + # not. We suppress the error to make the behavior consistent across all obstore + # stores. This is also in line with the behavior of the other Zarr store adapters. + with contextlib.suppress(FileNotFoundError): + await obs.delete_async(self.store, key) + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + + metas = await obs.list(self.store, prefix).collect_async() + keys = [(m["path"],) for m in metas] + await concurrent_map(keys, self.delete, limit=config.get("async.concurrency")) + + @property + def supports_listing(self) -> bool: + # docstring inherited + return True + + async def _list(self, prefix: str | None = None) -> AsyncGenerator[ObjectMeta, None]: + import obstore as obs + + objects: ListStream[Sequence[ObjectMeta]] = obs.list(self.store, prefix=prefix) + async for batch in objects: + for item in batch: + yield item + + def list(self) -> AsyncGenerator[str, None]: + # docstring inherited + return (obj["path"] async for obj in self._list()) + + def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + return (obj["path"] async for obj in self._list(prefix)) + + def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + import obstore as obs + + coroutine = obs.list_with_delimiter_async(self.store, prefix=prefix) + return _transform_list_dir(coroutine, prefix) + + async def getsize(self, key: str) -> int: + # docstring inherited + import obstore as obs + + resp = await obs.head_async(self.store, key) + return resp["size"] + + async def getsize_prefix(self, prefix: str) -> int: + # docstring inherited + sizes = [obj["size"] async for obj in self._list(prefix=prefix)] + return sum(sizes) + + +async def _transform_list_dir( + list_result_coroutine: Coroutine[Any, Any, ListResult[Sequence[ObjectMeta]]], prefix: str +) -> AsyncGenerator[str, None]: + """ + Transform the result of list_with_delimiter into an async generator of paths. + """ + list_result = await list_result_coroutine + + # We assume that the underlying object-store implementation correctly handles the + # prefix, so we don't double-check that the returned results actually start with the + # given prefix. + prefix = prefix.rstrip("/") + for path in chain( + list_result["common_prefixes"], map(itemgetter("path"), list_result["objects"]) + ): + yield _relativize_path(path=path, prefix=prefix) + + +class _BoundedRequest(TypedDict): + """Range request with a known start and end byte. + + These requests can be multiplexed natively on the Rust side with + `obstore.get_ranges_async`. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + start: int + """Start byte offset.""" + + end: int + """End byte offset.""" + + +class _OtherRequest(TypedDict): + """Offset or suffix range requests. + + These requests cannot be concurrent on the Rust side, and each need their own call + to `obstore.get_async`, passing in the `range` parameter. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + path: str + """The path to request from.""" + + range: OffsetRange | None + # Note: suffix requests are handled separately because some object stores (Azure) + # don't support them + """The range request type.""" + + +class _SuffixRequest(TypedDict): + """Offset or suffix range requests. + + These requests cannot be concurrent on the Rust side, and each need their own call + to `obstore.get_async`, passing in the `range` parameter. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + path: str + """The path to request from.""" + + range: SuffixRange + """The suffix range.""" + + +class _Response(TypedDict): + """A response buffer associated with the original index that it should be restored to.""" + + original_request_index: int + """The positional index in the original key_ranges input""" + + buffer: Buffer + """The buffer returned from obstore's range request.""" + + +async def _make_bounded_requests( + store: _UpstreamObjectStore, + path: str, + requests: list[_BoundedRequest], + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make all bounded requests for a specific file. + + `obstore.get_ranges_async` allows for making concurrent requests for multiple ranges + within a single file, and will e.g. merge concurrent requests. This only uses one + single Python coroutine. + """ + import obstore as obs + + starts = [r["start"] for r in requests] + ends = [r["end"] for r in requests] + async with semaphore: + responses = await obs.get_ranges_async(store, path=path, starts=starts, ends=ends) + + buffer_responses: list[_Response] = [] + for request, response in zip(requests, responses, strict=True): + buffer_responses.append( + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(response), # type: ignore[arg-type] + } + ) + + return buffer_responses + + +async def _make_other_request( + store: _UpstreamObjectStore, + request: _OtherRequest, + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make offset or full-file requests. + + We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all + futures can be gathered together. + """ + import obstore as obs + + async with semaphore: + if request["range"] is None: + resp = await obs.get_async(store, request["path"]) + else: + resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) + buffer = await resp.bytes_async() + + return [ + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(buffer), # type: ignore[arg-type] + } + ] + + +async def _make_suffix_request( + store: _UpstreamObjectStore, + request: _SuffixRequest, + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make suffix requests. + + This is separated out from `_make_other_request` because some object stores (Azure) + don't support suffix requests. In this case, our workaround is to first get the + length of the object and then manually request the byte range at the end. + + We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all + futures can be gathered together. + """ + import obstore as obs + + async with semaphore: + try: + resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) + buffer = await resp.bytes_async() + except obs.exceptions.NotSupportedError: + head_resp = await obs.head_async(store, request["path"]) + file_size = head_resp["size"] + suffix_len = request["range"]["suffix"] + buffer = await obs.get_range_async( + store, + request["path"], + start=file_size - suffix_len, + length=suffix_len, + ) + + return [ + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(buffer), # type: ignore[arg-type] + } + ] + + +async def _get_partial_values( + store: _UpstreamObjectStore, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], +) -> list[Buffer | None]: + """Make multiple range requests. + + ObjectStore has a `get_ranges` method that will additionally merge nearby ranges, + but it's _per_ file. So we need to split these key_ranges into **per-file** key + ranges, and then reassemble the results in the original order. + + We separate into different requests: + + - One call to `obstore.get_ranges_async` **per target file** + - One call to `obstore.get_async` for each other request. + """ + key_ranges = list(key_ranges) + per_file_bounded_requests: dict[str, list[_BoundedRequest]] = defaultdict(list) + other_requests: list[_OtherRequest] = [] + suffix_requests: list[_SuffixRequest] = [] + + for idx, (path, byte_range) in enumerate(key_ranges): + if byte_range is None: + other_requests.append( + { + "original_request_index": idx, + "path": path, + "range": None, + } + ) + elif isinstance(byte_range, RangeByteRequest): + per_file_bounded_requests[path].append( + {"original_request_index": idx, "start": byte_range.start, "end": byte_range.end} + ) + elif isinstance(byte_range, OffsetByteRequest): + other_requests.append( + { + "original_request_index": idx, + "path": path, + "range": {"offset": byte_range.offset}, + } + ) + elif isinstance(byte_range, SuffixByteRequest): + suffix_requests.append( + { + "original_request_index": idx, + "path": path, + "range": {"suffix": byte_range.suffix}, + } + ) + else: + raise ValueError(f"Unsupported range input: {byte_range}") + + semaphore = asyncio.Semaphore(config.get("async.concurrency")) + + futs: list[Coroutine[Any, Any, list[_Response]]] = [] + for path, bounded_ranges in per_file_bounded_requests.items(): + futs.append( + _make_bounded_requests(store, path, bounded_ranges, prototype, semaphore=semaphore) + ) + + for request in other_requests: + futs.append(_make_other_request(store, request, prototype, semaphore=semaphore)) # noqa: PERF401 + + for suffix_request in suffix_requests: + futs.append(_make_suffix_request(store, suffix_request, prototype, semaphore=semaphore)) # noqa: PERF401 + + buffers: list[Buffer | None] = [None] * len(key_ranges) + + for responses in await asyncio.gather(*futs): + for resp in responses: + buffers[resp["original_request_index"]] = resp["buffer"] + + return buffers diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_utils.py b/.venv/lib/python3.11/site-packages/zarr/storage/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..10ac395b360e1b0a0419a5bcdeaa5d383a5d81f7 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_utils.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +from zarr.abc.store import OffsetByteRequest, RangeByteRequest, SuffixByteRequest + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + + from zarr.abc.store import ByteRequest + from zarr.core.buffer import Buffer + + +def normalize_path(path: str | bytes | Path | None) -> str: + if path is None: + result = "" + elif isinstance(path, bytes): + result = str(path, "ascii") + + # handle pathlib.Path + elif isinstance(path, Path): + result = str(path) + + elif isinstance(path, str): + result = path + + else: + raise TypeError(f'Object {path} has an invalid type for "path": {type(path).__name__}') + + # convert backslash to forward slash + result = result.replace("\\", "/") + + # remove leading and trailing slashes + result = result.strip("/") + + # collapse any repeated slashes + pat = re.compile(r"//+") + result = pat.sub("/", result) + + # disallow path segments with just '.' or '..' + segments = result.split("/") + if any(s in {".", ".."} for s in segments): + raise ValueError( + f"The path {path!r} is invalid because its string representation contains '.' or '..' segments." + ) + + return result + + +def _normalize_byte_range_index(data: Buffer, byte_range: ByteRequest | None) -> tuple[int, int]: + """ + Convert a ByteRequest into an explicit start and stop + """ + if byte_range is None: + start = 0 + stop = len(data) + 1 + elif isinstance(byte_range, RangeByteRequest): + start = byte_range.start + stop = byte_range.end + elif isinstance(byte_range, OffsetByteRequest): + start = byte_range.offset + stop = len(data) + 1 + elif isinstance(byte_range, SuffixByteRequest): + start = len(data) - byte_range.suffix + stop = len(data) + 1 + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + return (start, stop) + + +def _join_paths(paths: Iterable[str]) -> str: + """ + Filter out instances of '' and join the remaining strings with '/'. + + Parameters + ---------- + paths : Iterable[str] + + Returns + ------- + str + + Examples + -------- + ```python + from zarr.storage._utils import _join_paths + _join_paths(["", "a", "b"]) + # 'a/b' + _join_paths(["a", "b", "c"]) + # 'a/b/c' + ``` + """ + return "/".join(filter(lambda v: v != "", paths)) + + +def _relativize_path(*, path: str, prefix: str) -> str: + """ + Make a "/"-delimited path relative to some prefix. If the prefix is '', then the path is + returned as-is. Otherwise, the prefix is removed from the path as well as the separator + string "/". + + If ``prefix`` is not the empty string and ``path`` does not start with ``prefix`` + followed by a "/" character, then an error is raised. + + This function assumes that the prefix does not end with "/". + + Parameters + ---------- + path : str + The path to make relative to the prefix. + prefix : str + The prefix to make the path relative to. + + Returns + ------- + str + + Examples + -------- + ```python + from zarr.storage._utils import _relativize_path + _relativize_path(path="a/b", prefix="") + # 'a/b' + _relativize_path(path="a/b/c", prefix="a/b") + # 'c' + ``` + """ + if prefix == "": + return path + else: + _prefix = prefix + "/" + if not path.startswith(_prefix): + raise ValueError(f"The first component of {path} does not start with {prefix}.") + return path.removeprefix(f"{prefix}/") + + +def _normalize_paths(paths: Iterable[str]) -> tuple[str, ...]: + """ + Normalize the input paths according to the normalization scheme used for zarr node paths. + If any two paths normalize to the same value, raise a ValueError. + """ + path_map: dict[str, str] = {} + for path in paths: + parsed = normalize_path(path) + if parsed in path_map: + msg = ( + f"After normalization, the value '{path}' collides with '{path_map[parsed]}'. " + f"Both '{path}' and '{path_map[parsed]}' normalize to the same value: '{parsed}'. " + f"You should use either '{path}' or '{path_map[parsed]}', but not both." + ) + raise ValueError(msg) + path_map[parsed] = path + return tuple(path_map.keys()) + + +T = TypeVar("T") + + +def _normalize_path_keys(data: Mapping[str, T]) -> dict[str, T]: + """ + Normalize the keys of the input dict according to the normalization scheme used for zarr node + paths. If any two keys in the input normalize to the same value, raise a ValueError. + Returns a dict where the keys are the elements of the input and the values are the + normalized form of each key. + """ + parsed_keys = _normalize_paths(data.keys()) + return dict(zip(parsed_keys, data.values(), strict=True)) diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_wrapper.py b/.venv/lib/python3.11/site-packages/zarr/storage/_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..e8a2859abcc7f4e8a9dfba5e99860a9dc522f5cb --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_wrapper.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Generic, TypeVar, cast + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Iterable + from types import TracebackType + from typing import Any, Self + + from zarr.abc.buffer import Buffer + from zarr.abc.store import ByteRequest + from zarr.core.buffer import BufferPrototype + +from zarr.abc.store import Store + +T_Store = TypeVar("T_Store", bound=Store) + + +class WrapperStore(Store, Generic[T_Store]): + """ + Store that wraps an existing Store. + + By default all of the store methods are delegated to the wrapped store instance, which is + accessible via the ``._store`` attribute of this class. + + Use this class to modify or extend the behavior of the other store classes. + """ + + _store: T_Store + + def __init__(self, store: T_Store) -> None: + self._store = store + + def _with_store(self, store: T_Store) -> Self: + """ + Constructs a new instance of the wrapper store with the same details but a new store. + """ + return type(self)(store=store) + + @classmethod + async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self: + store = store_cls(*args, **kwargs) + await store._open() + return cls(store=store) + + def with_read_only(self, read_only: bool = False) -> Self: + return self._with_store(cast(T_Store, self._store.with_read_only(read_only))) + + def __enter__(self) -> Self: + return self._with_store(self._store.__enter__()) + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + return self._store.__exit__(exc_type, exc_value, traceback) + + async def _open(self) -> None: + await self._store._open() + + async def _ensure_open(self) -> None: + await self._store._ensure_open() + + async def is_empty(self, prefix: str) -> bool: + return await self._store.is_empty(prefix) + + @property + def _is_open(self) -> bool: + return self._store._is_open + + @_is_open.setter + def _is_open(self, value: bool) -> None: + raise NotImplementedError("WrapperStore must be opened via the `_open` method") + + async def clear(self) -> None: + return await self._store.clear() + + @property + def read_only(self) -> bool: + return self._store.read_only + + def _check_writable(self) -> None: + return self._store._check_writable() + + def __eq__(self, value: object) -> bool: + return type(self) is type(value) and self._store.__eq__(value._store) # type: ignore[attr-defined] + + def __str__(self) -> str: + return f"wrapping-{self._store}" + + def __repr__(self) -> str: + return f"WrapperStore({self._store.__class__.__name__}, '{self._store}')" + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + return await self._store.get(key, prototype, byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + return await self._store.get_partial_values(prototype, key_ranges) + + async def exists(self, key: str) -> bool: + return await self._store.exists(key) + + async def set(self, key: str, value: Buffer) -> None: + await self._store.set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + return await self._store.set_if_not_exists(key, value) + + async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None: + await self._store._set_many(values) + + @property + def supports_writes(self) -> bool: + return self._store.supports_writes + + @property + def supports_deletes(self) -> bool: + return self._store.supports_deletes + + async def delete(self, key: str) -> None: + await self._store.delete(key) + + @property + def supports_listing(self) -> bool: + return self._store.supports_listing + + def list(self) -> AsyncIterator[str]: + return self._store.list() + + def list_prefix(self, prefix: str) -> AsyncIterator[str]: + return self._store.list_prefix(prefix) + + def list_dir(self, prefix: str) -> AsyncIterator[str]: + return self._store.list_dir(prefix) + + async def delete_dir(self, prefix: str) -> None: + return await self._store.delete_dir(prefix) + + def close(self) -> None: + self._store.close() + + async def _get_many( + self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]] + ) -> AsyncGenerator[tuple[str, Buffer | None], None]: + async for req in self._store._get_many(requests): + yield req diff --git a/.venv/lib/python3.11/site-packages/zarr/storage/_zip.py b/.venv/lib/python3.11/site-packages/zarr/storage/_zip.py new file mode 100644 index 0000000000000000000000000000000000000000..72bf9e335a96fdfc72c9d051ecdf9ddbe2ceed56 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/storage/_zip.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +import os +import shutil +import threading +import time +import zipfile +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr.core.buffer import Buffer, BufferPrototype + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable + +ZipStoreAccessModeLiteral = Literal["r", "w", "a"] + + +class ZipStore(Store): + """ + Store using a ZIP file. + + Parameters + ---------- + path : str + Location of file. + mode : str, optional + One of 'r' to read an existing file, 'w' to truncate and write a new + file, 'a' to append to an existing file, or 'x' to exclusively create + and write a new file. + compression : int, optional + Compression method to use when writing to the archive. + allowZip64 : bool, optional + If True (the default) will create ZIP files that use the ZIP64 + extensions when the zipfile is larger than 2 GiB. If False + will raise an exception when the ZIP file would require ZIP64 + extensions. + + Attributes + ---------- + allowed_exceptions + supports_writes + supports_deletes + supports_listing + path + compression + allowZip64 + """ + + supports_writes: bool = True + supports_deletes: bool = False + supports_listing: bool = True + + path: Path + compression: int + allowZip64: bool + + _zf: zipfile.ZipFile + _lock: threading.RLock + + def __init__( + self, + path: Path | str, + *, + mode: ZipStoreAccessModeLiteral = "r", + read_only: bool | None = None, + compression: int = zipfile.ZIP_STORED, + allowZip64: bool = True, + ) -> None: + if read_only is None: + read_only = mode == "r" + + super().__init__(read_only=read_only) + + if isinstance(path, str): + path = Path(path) + assert isinstance(path, Path) + self.path = path # root? + + self._zmode = mode + self.compression = compression + self.allowZip64 = allowZip64 + + def _sync_open(self) -> None: + if self._is_open: + raise ValueError("store is already open") + + self._lock = threading.RLock() + + self._zf = zipfile.ZipFile( + self.path, + mode=self._zmode, + compression=self.compression, + allowZip64=self.allowZip64, + ) + + self._is_open = True + + async def _open(self) -> None: + self._sync_open() + + def __getstate__(self) -> dict[str, Any]: + # We need a copy to not modify the state of the original store + state = self.__dict__.copy() + for attr in ["_zf", "_lock"]: + state.pop(attr, None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__ = state + self._is_open = False + self._sync_open() + + def close(self) -> None: + # docstring inherited + super().close() + with self._lock: + self._zf.close() + + async def clear(self) -> None: + # docstring inherited + with self._lock: + self._check_writable() + self._zf.close() + os.remove(self.path) + self._zf = zipfile.ZipFile( + self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64 + ) + + def __str__(self) -> str: + return f"zip://{self.path}" + + def __repr__(self) -> str: + return f"ZipStore('{self}')" + + def __eq__(self, other: object) -> bool: + return isinstance(other, type(self)) and self.path == other.path + + def _get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if not self._is_open: + self._sync_open() + # docstring inherited + try: + with self._zf.open(key) as f: # will raise KeyError + if byte_range is None: + return prototype.buffer.from_bytes(f.read()) + elif isinstance(byte_range, RangeByteRequest): + f.seek(byte_range.start) + return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell())) + size = f.seek(0, os.SEEK_END) + if isinstance(byte_range, OffsetByteRequest): + f.seek(byte_range.offset) + elif isinstance(byte_range, SuffixByteRequest): + f.seek(max(0, size - byte_range.suffix)) + else: + raise TypeError(f"Unexpected byte_range, got {byte_range}.") + return prototype.buffer.from_bytes(f.read()) + except KeyError: + return None + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + assert isinstance(key, str) + + with self._lock: + return self._get(key, prototype=prototype, byte_range=byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + out = [] + with self._lock: + for key, byte_range in key_ranges: + out.append(self._get(key, prototype=prototype, byte_range=byte_range)) + return out + + def _set(self, key: str, value: Buffer) -> None: + if not self._is_open: + self._sync_open() + # generally, this should be called inside a lock + keyinfo = zipfile.ZipInfo(filename=key, date_time=time.localtime(time.time())[:6]) + keyinfo.compress_type = self.compression + if keyinfo.filename[-1] == os.sep: + keyinfo.external_attr = 0o40775 << 16 # drwxrwxr-x + keyinfo.external_attr |= 0x10 # MS-DOS directory flag + else: + keyinfo.external_attr = 0o644 << 16 # ?rw-r--r-- + self._zf.writestr(keyinfo, value.to_bytes()) + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + self._check_writable() + if not self._is_open: + self._sync_open() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + with self._lock: + self._set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + self._check_writable() + with self._lock: + members = self._zf.namelist() + if key not in members: + self._set(key, value) + + async def delete_dir(self, prefix: str) -> None: + # only raise NotImplementedError if any keys are found + self._check_writable() + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + async for _ in self.list_prefix(prefix): + raise NotImplementedError + + async def delete(self, key: str) -> None: + # docstring inherited + # we choose to only raise NotImplementedError here if the key exists + # this allows the array/group APIs to avoid the overhead of existence checks + self._check_writable() + if await self.exists(key): + raise NotImplementedError + + async def exists(self, key: str) -> bool: + # docstring inherited + with self._lock: + try: + self._zf.getinfo(key) + except KeyError: + return False + else: + return True + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + with self._lock: + for key in self._zf.namelist(): + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + async for key in self.list(): + if key.startswith(prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + prefix = prefix.rstrip("/") + + keys = self._zf.namelist() + seen = set() + if prefix == "": + keys_unique = {k.split("/")[0] for k in keys} + for key in keys_unique: + if key not in seen: + seen.add(key) + yield key + else: + for key in keys: + if key.startswith(prefix + "/") and key.strip("/") != prefix: + k = key.removeprefix(prefix + "/").split("/")[0] + if k not in seen: + seen.add(k) + yield k + + async def move(self, path: Path | str) -> None: + """ + Move the store to another path. + """ + if isinstance(path, str): + path = Path(path) + self.close() + os.makedirs(path.parent, exist_ok=True) + shutil.move(self.path, path) + self.path = path + await self._open() diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/__init__.py b/.venv/lib/python3.11/site-packages/zarr/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..21a3572846b4f485ff4490eee569a54f24364ee2 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/__init__.py @@ -0,0 +1,17 @@ +import importlib.util +import warnings + +from zarr.errors import ZarrUserWarning + +if importlib.util.find_spec("pytest") is not None: + from zarr.testing.store import StoreTests +else: + warnings.warn( + "pytest not installed, skipping test suite", category=ZarrUserWarning, stacklevel=2 + ) + +from zarr.testing.utils import assert_bytes_equal + +# TODO: import public buffer tests? + +__all__ = ["StoreTests", "assert_bytes_equal"] diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/buffer.py b/.venv/lib/python3.11/site-packages/zarr/testing/buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..6096ece2f82bf58ba09bd107b946d0d566a69d63 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/buffer.py @@ -0,0 +1,86 @@ +# mypy: ignore-errors +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import numpy.typing as npt + +from zarr.core.buffer import Buffer, BufferPrototype, cpu +from zarr.storage import MemoryStore + +if TYPE_CHECKING: + from collections.abc import Iterable + from typing import Self + + +__all__ = [ + "NDBufferUsingTestNDArrayLike", + "StoreExpectingTestBuffer", + "TestBuffer", +] + + +class TestNDArrayLike(np.ndarray): + """An example of an ndarray-like class""" + + __test__ = False + + +class TestBuffer(cpu.Buffer): + """Example of a custom Buffer that handles ArrayLike""" + + __test__ = False + + +class NDBufferUsingTestNDArrayLike(cpu.NDBuffer): + """Example of a custom NDBuffer that handles MyNDArrayLike""" + + @classmethod + def create( + cls, + *, + shape: Iterable[int], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + fill_value: Any | None = None, + ) -> Self: + """Overwrite `NDBuffer.create` to create a TestNDArrayLike instance""" + ret = cls(TestNDArrayLike(shape=shape, dtype=dtype, order=order)) + if fill_value is not None: + ret.fill(fill_value) + return ret + + @classmethod + def empty( + cls, + shape: tuple[int, ...], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + ) -> Self: + return super(cpu.NDBuffer, cls).empty(shape=shape, dtype=dtype, order=order) + + +class StoreExpectingTestBuffer(MemoryStore): + """Example of a custom Store that expect MyBuffer for all its non-metadata + + We assume that keys containing "json" is metadata + """ + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + if "json" not in key: + assert isinstance(value, TestBuffer) + await super().set(key, value, byte_range) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: tuple[int, int | None] | None = None, + ) -> Buffer | None: + if "json" not in key: + assert prototype.buffer is TestBuffer + ret = await super().get(key=key, prototype=prototype, byte_range=byte_range) + if ret is not None: + assert isinstance(ret, prototype.buffer) + return ret diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/conftest.py b/.venv/lib/python3.11/site-packages/zarr/testing/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..59c148e0ecb82eb31ea99911771afd96b344fc60 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/conftest.py @@ -0,0 +1,9 @@ +import pytest + + +def pytest_configure(config: pytest.Config) -> None: + # The tests in zarr.testing are intended to be run by downstream projects. + # To allow those downstream projects to run with `--strict-markers`, we need + # to register an entry point with pytest11 and register our "plugin" with it, + # which just registers the markers used in zarr.testing + config.addinivalue_line("markers", "gpu: mark a test as requiring CuPy and GPU") diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/stateful.py b/.venv/lib/python3.11/site-packages/zarr/testing/stateful.py new file mode 100644 index 0000000000000000000000000000000000000000..382f1467da8f2979ecf1122317b5dc8276aecc1a --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/stateful.py @@ -0,0 +1,630 @@ +import builtins +import functools +from collections.abc import Callable +from typing import Any, TypeVar, cast + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +from hypothesis import assume, note +from hypothesis.stateful import ( + RuleBasedStateMachine, + initialize, + invariant, + precondition, + rule, +) +from hypothesis.strategies import DataObject + +import zarr +from zarr import Array +from zarr.abc.store import Store +from zarr.codecs.bytes import BytesCodec +from zarr.core.buffer import Buffer, BufferPrototype, cpu, default_buffer_prototype +from zarr.core.sync import SyncMixin +from zarr.storage import LocalStore, MemoryStore +from zarr.testing.strategies import ( + basic_indices, + chunk_paths, + dimension_names, + key_ranges, + node_names, + np_array_and_chunks, + orthogonal_indices, +) +from zarr.testing.strategies import keys as zarr_keys + +MAX_BINARY_SIZE = 100 + +F = TypeVar("F", bound=Callable[..., Any]) + + +def with_frequency(frequency: float) -> Callable[[F], F]: + """This needs to be deterministic for hypothesis replaying""" + + def decorator(func: F) -> F: + counter_attr = f"__{func.__name__}_counter" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + @precondition + def frequency_check(f: Any) -> Any: + if not hasattr(f, counter_attr): + setattr(f, counter_attr, 0) + + current_count = getattr(f, counter_attr) + 1 + setattr(f, counter_attr, current_count) + + return (current_count * frequency) % 1.0 >= (1.0 - frequency) + + return cast(F, frequency_check(wrapper)) + + return decorator + + +def split_prefix_name(path: str) -> tuple[str, str]: + split = path.rsplit("/", maxsplit=1) + if len(split) > 1: + prefix, name = split + else: + prefix = "" + (name,) = split + return prefix, name + + +class ZarrHierarchyStateMachine(SyncMixin, RuleBasedStateMachine): + """ + This state machine models operations that modify a zarr store's + hierarchy. That is, user actions that modify arrays/groups as well + as list operations. It is intended to be used by external stores, and + compares their results to a MemoryStore that is assumed to be perfect. + """ + + def __init__(self, store: Store) -> None: + super().__init__() + + self.store = store + + self.model = MemoryStore() + zarr.group(store=self.model) + + # Track state of the hierarchy, these should contain fully qualified paths + self.all_groups: set[str] = set() + self.all_arrays: set[str] = set() + + @initialize() + def init_store(self) -> None: + # This lets us reuse the fixture provided store. + self._sync(self.store.clear()) + zarr.group(store=self.store) + + def can_add(self, path: str) -> bool: + return path not in self.all_groups and path not in self.all_arrays + + # -------------------- store operations ----------------------- + @rule(name=node_names, data=st.data()) + def add_group(self, name: str, data: DataObject) -> None: + # Handle possible case-insensitive file systems (e.g. MacOS) + if isinstance(self.store, LocalStore): + name = name.lower() + if self.all_groups: + parent = data.draw(st.sampled_from(sorted(self.all_groups)), label="Group parent") + else: + parent = "" + path = f"{parent}/{name}".lstrip("/") + assume(self.can_add(path)) + note(f"Adding group: path='{path}'") + self.all_groups.add(path) + zarr.group(store=self.store, path=path) + zarr.group(store=self.model, path=path) + + @rule(data=st.data(), name=node_names, array_and_chunks=np_array_and_chunks()) + def add_array( + self, + data: DataObject, + name: str, + array_and_chunks: tuple[np.ndarray[Any, Any], tuple[int, ...]], + ) -> None: + # Handle possible case-insensitive file systems (e.g. MacOS) + if isinstance(self.store, LocalStore): + name = name.lower() + array, chunks = array_and_chunks + fill_value = data.draw(npst.from_dtype(array.dtype)) + if self.all_groups: + parent = data.draw(st.sampled_from(sorted(self.all_groups)), label="Array parent") + else: + parent = "" + # TODO: support creating deeper paths + # TODO: support overwriting potentially by just skipping `self.can_add` + path = f"{parent}/{name}".lstrip("/") + assume(self.can_add(path)) + note(f"Adding array: path='{path}' shape={array.shape} chunks={chunks}") + for store in [self.store, self.model]: + zarr.array( + array, + chunks=chunks, + path=path, + store=store, + fill_value=fill_value, + zarr_format=3, + dimension_names=data.draw( + dimension_names(ndim=array.ndim), label="dimension names" + ), + # Chose bytes codec to avoid wasting time compressing the data being written + codecs=[BytesCodec()], + ) + self.all_arrays.add(path) + + @rule() + @with_frequency(0.25) + def clear(self) -> None: + note("clearing") + import zarr + + self._sync(self.store.clear()) + self._sync(self.model.clear()) + + assert self._sync(self.store.is_empty("/")) + assert self._sync(self.model.is_empty("/")) + + self.all_groups.clear() + self.all_arrays.clear() + + zarr.group(store=self.store) + zarr.group(store=self.model) + + # TODO: MemoryStore is broken? + # assert not self._sync(self.store.is_empty("/")) + # assert not self._sync(self.model.is_empty("/")) + + def draw_directory(self, data: DataObject) -> str: + group_st = st.sampled_from(sorted(self.all_groups)) if self.all_groups else st.nothing() + array_st = st.sampled_from(sorted(self.all_arrays)) if self.all_arrays else st.nothing() + array_or_group = data.draw(st.one_of(group_st, array_st)) + if data.draw(st.booleans()) and array_or_group in self.all_arrays: + arr = zarr.open_array(path=array_or_group, store=self.model) + path = data.draw( + st.one_of( + st.sampled_from([array_or_group]), + chunk_paths(ndim=arr.ndim, numblocks=arr.cdata_shape).map( + lambda x: f"{array_or_group}/c/" + ), + ) + ) + else: + path = array_or_group + return path + + @precondition(lambda self: bool(self.all_groups)) + @rule(data=st.data()) + def check_list_dir(self, data: DataObject) -> None: + path = self.draw_directory(data) + note(f"list_dir for {path=!r}") + # Consider .list_dir("path/to/array") for an array with a single chunk. + # The MemoryStore model will return `"c", "zarr.json"` only if the chunk exists + # If that chunk was deleted, then `"c"` is not returned. + # LocalStore will not have this behaviour :/ + # There are similar consistency issues with delete_dir("/path/to/array/c/0/0") + assume(not isinstance(self.store, LocalStore)) + model_ls = sorted(self._sync_iter(self.model.list_dir(path))) + store_ls = sorted(self._sync_iter(self.store.list_dir(path))) + assert model_ls == store_ls, (model_ls, store_ls) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def delete_chunk(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + arr = zarr.open_array(path=array, store=self.model) + chunk_path = data.draw(chunk_paths(ndim=arr.ndim, numblocks=arr.cdata_shape, subset=False)) + path = f"{array}/c/{chunk_path}" + note(f"deleting chunk {path=!r}") + self._sync(self.model.delete(path)) + self._sync(self.store.delete(path)) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def check_array(self, data: DataObject) -> None: + path = data.draw(st.sampled_from(sorted(self.all_arrays))) + actual = zarr.open_array(self.store, path=path)[:] + expected = zarr.open_array(self.model, path=path)[:] + np.testing.assert_equal(actual, expected) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def overwrite_array_basic_indexing(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + slicer = data.draw(basic_indices(shape=model_array.shape)) + note(f"overwriting array with basic indexer: {slicer=}") + new_data = data.draw( + npst.arrays(shape=np.shape(model_array[slicer]), dtype=model_array.dtype) + ) + model_array[slicer] = new_data + store_array[slicer] = new_data + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def overwrite_array_orthogonal_indexing(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + indexer, _ = data.draw(orthogonal_indices(shape=model_array.shape)) + note(f"overwriting array orthogonal {indexer=}") + new_data = data.draw( + npst.arrays(shape=model_array.oindex[indexer].shape, dtype=model_array.dtype) # type: ignore[union-attr] + ) + model_array.oindex[indexer] = new_data + store_array.oindex[indexer] = new_data + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def resize_array(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + ndim = model_array.ndim + new_shape = tuple( + 0 if oldsize == 0 else newsize + for newsize, oldsize in zip( + data.draw(npst.array_shapes(max_dims=ndim, min_dims=ndim, min_side=0)), + model_array.shape, + strict=True, + ) + ) + + note(f"resizing array from {model_array.shape} to {new_shape}") + model_array.resize(new_shape) + store_array.resize(new_shape) + + @precondition(lambda self: bool(self.all_arrays) or bool(self.all_groups)) + @rule(data=st.data()) + def delete_dir(self, data: DataObject) -> None: + path = self.draw_directory(data) + note(f"delete_dir with {path=!r}") + self._sync(self.model.delete_dir(path)) + self._sync(self.store.delete_dir(path)) + + matches = set() + for node in self.all_groups | self.all_arrays: + if node.startswith(path): + matches.add(node) + self.all_groups = self.all_groups - matches + self.all_arrays = self.all_arrays - matches + + # @precondition(lambda self: bool(self.all_groups)) + # @precondition(lambda self: bool(self.all_arrays)) + # @rule(data=st.data()) + # def move_array(self, data): + # array_path = data.draw(st.sampled_from(self.all_arrays), label="Array move source") + # to_group = data.draw(st.sampled_from(self.all_groups), label="Array move destination") + + # # fixme renaming to self? + # array_name = os.path.basename(array_path) + # assume(self.model.can_add(to_group, array_name)) + # new_path = f"{to_group}/{array_name}".lstrip("/") + # note(f"moving array '{array_path}' -> '{new_path}'") + # self.model.rename(array_path, new_path) + # self.repo.store.rename(array_path, new_path) + + # @precondition(lambda self: len(self.all_groups) >= 2) + # @rule(data=st.data()) + # def move_group(self, data): + # from_group = data.draw(st.sampled_from(self.all_groups), label="Group move source") + # to_group = data.draw(st.sampled_from(self.all_groups), label="Group move destination") + # assume(not to_group.startswith(from_group)) + + # from_group_name = os.path.basename(from_group) + # assume(self.model.can_add(to_group, from_group_name)) + # # fixme renaming to self? + # new_path = f"{to_group}/{from_group_name}".lstrip("/") + # note(f"moving group '{from_group}' -> '{new_path}'") + # self.model.rename(from_group, new_path) + # self.repo.store.rename(from_group, new_path) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: len(self.all_arrays) >= 1) + @rule(data=st.data()) + def delete_array_using_del(self, data: DataObject) -> None: + array_path = data.draw( + st.sampled_from(sorted(self.all_arrays)), label="Array deletion target" + ) + prefix, array_name = split_prefix_name(array_path) + note(f"Deleting array '{array_path}' ({prefix=!r}, {array_name=!r}) using del") + for store in [self.model, self.store]: + group = zarr.open_group(path=prefix, store=store) + group[array_name] # check that it exists + del group[array_name] + self.all_arrays.remove(array_path) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: bool(self.all_groups)) + @rule(data=st.data()) + def delete_group_using_del(self, data: DataObject) -> None: + group_path = data.draw( + st.sampled_from(sorted(self.all_groups)), + label="Group deletion target", + ) + prefix, group_name = split_prefix_name(group_path) + note(f"Deleting group '{group_path=!r}', {prefix=!r}, {group_name=!r} using delete") + members = zarr.open_group(store=self.model, path=group_path).members(max_depth=None) + for _, obj in members: + if isinstance(obj, Array): + self.all_arrays.remove(obj.path) + else: + self.all_groups.remove(obj.path) + for store in [self.store, self.model]: + group = zarr.open_group(store=store, path=prefix) + group[group_name] # check that it exists + del group[group_name] + self.all_groups.remove(group_path) + + # # --------------- assertions ----------------- + # def check_group_arrays(self, group): + # # note(f"Checking arrays of '{group}'") + # g1 = self.model.get_group(group) + # g2 = zarr.open_group(path=group, mode="r", store=self.repo.store) + # model_arrays = sorted(g1.arrays(), key=itemgetter(0)) + # our_arrays = sorted(g2.arrays(), key=itemgetter(0)) + # for (n1, a1), (n2, a2) in zip_longest(model_arrays, our_arrays): + # assert n1 == n2 + # assert_array_equal(a1, a2) + + # def check_subgroups(self, group_path): + # g1 = self.model.get_group(group_path) + # g2 = zarr.open_group(path=group_path, mode="r", store=self.repo.store) + # g1_children = [name for (name, _) in g1.groups()] + # g2_children = [name for (name, _) in g2.groups()] + # # note(f"Checking {len(g1_children)} subgroups of group '{group_path}'") + # assert g1_children == g2_children + + # def check_list_prefix_from_group(self, group): + # prefix = f"meta/root/{group}" + # model_list = sorted(self.model.list_prefix(prefix)) + # al_list = sorted(self.repo.store.list_prefix(prefix)) + # # note(f"Checking {len(model_list)} keys under '{prefix}'") + # assert model_list == al_list + + # prefix = f"data/root/{group}" + # model_list = sorted(self.model.list_prefix(prefix)) + # al_list = sorted(self.repo.store.list_prefix(prefix)) + # # note(f"Checking {len(model_list)} keys under '{prefix}'") + # assert model_list == al_list + + # @precondition(lambda self: self.model.is_persistent_session()) + # @rule(data=st.data()) + # def check_group_path(self, data): + # t0 = time.time() + # group = data.draw(st.sampled_from(self.all_groups)) + # self.check_list_prefix_from_group(group) + # self.check_subgroups(group) + # self.check_group_arrays(group) + # t1 = time.time() + # note(f"Checks took {t1 - t0} sec.") + @invariant() + def check_list_prefix_from_root(self) -> None: + model_list = self._sync_iter(self.model.list_prefix("")) + store_list = self._sync_iter(self.store.list_prefix("")) + note(f"Checking {len(model_list)} expected keys vs {len(store_list)} actual keys") + assert sorted(model_list) == sorted(store_list), ( + sorted(model_list), + sorted(store_list), + ) + + # check that our internal state matches that of the store and model + assert all(f"{path}/zarr.json" in model_list for path in self.all_groups | self.all_arrays) + assert all(f"{path}/zarr.json" in store_list for path in self.all_groups | self.all_arrays) + + +class SyncStoreWrapper(zarr.core.sync.SyncMixin): + def __init__(self, store: Store) -> None: + """Synchronous Store wrapper + + This class holds synchronous methods that map to async methods of Store classes. + The synchronous wrapper is needed because hypothesis' stateful testing infra does + not support asyncio so we redefine sync versions of the Store API. + https://github.com/HypothesisWorks/hypothesis/issues/3712#issuecomment-1668999041 + """ + self.store = store + + @property + def read_only(self) -> bool: + return self.store.read_only + + def set(self, key: str, data_buffer: Buffer) -> None: + return self._sync(self.store.set(key, data_buffer)) + + def list(self) -> builtins.list[str]: + return self._sync_iter(self.store.list()) + + def get(self, key: str, prototype: BufferPrototype) -> Buffer | None: + return self._sync(self.store.get(key, prototype=prototype)) + + def get_partial_values( + self, key_ranges: builtins.list[Any], prototype: BufferPrototype + ) -> builtins.list[Buffer | None]: + return self._sync(self.store.get_partial_values(prototype=prototype, key_ranges=key_ranges)) + + def delete(self, path: str) -> None: + return self._sync(self.store.delete(path)) + + def is_empty(self, prefix: str) -> bool: + return self._sync(self.store.is_empty(prefix=prefix)) + + def clear(self) -> None: + return self._sync(self.store.clear()) + + def exists(self, key: str) -> bool: + return self._sync(self.store.exists(key)) + + def list_dir(self, prefix: str) -> None: + raise NotImplementedError + + def list_prefix(self, prefix: str) -> None: + raise NotImplementedError + + @property + def supports_listing(self) -> bool: + return self.store.supports_listing + + @property + def supports_writes(self) -> bool: + return self.store.supports_writes + + @property + def supports_deletes(self) -> bool: + return self.store.supports_deletes + + +class ZarrStoreStateMachine(RuleBasedStateMachine): + """ " + Zarr store state machine + + This is a subclass of a Hypothesis RuleBasedStateMachine. + It is testing a framework to ensure that the state of a Zarr store matches + an expected state after a set of random operations. It contains a store + (currently, a Zarr MemoryStore) and a model, a simplified version of a + zarr store (in this case, a dict). It also contains rules which represent + actions that can be applied to a zarr store. Rules apply an action to both + the store and the model, and invariants assert that the state of the model + is equal to the state of the store. Hypothesis then generates sequences of + rules, running invariants after each rule. It raises an error if a sequence + produces discontinuity between state of the model and state of the store + (ie. an invariant is violated). + https://hypothesis.readthedocs.io/en/latest/stateful.html + """ + + def __init__(self, store: Store) -> None: + super().__init__() + self.model: dict[str, Buffer] = {} + self.store = SyncStoreWrapper(store) + self.prototype = default_buffer_prototype() + + @initialize() + def init_store(self) -> None: + self.store.clear() + + @rule(key=zarr_keys(), data=st.binary(min_size=0, max_size=MAX_BINARY_SIZE)) + def set(self, key: str, data: bytes) -> None: + note(f"(set) Setting {key!r} with {data!r}") + assert not self.store.read_only + data_buf = cpu.Buffer.from_bytes(data) + self.store.set(key, data_buf) + self.model[key] = data_buf + + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(key=zarr_keys(), data=st.data()) + def get(self, key: str, data: DataObject) -> None: + key = data.draw( + st.sampled_from(sorted(self.model.keys())) + ) # hypothesis wants to sample from sorted list + note("(get)") + store_value = self.store.get(key, self.prototype) + # to bytes here necessary because data_buf set to model in set() + assert self.model[key] == store_value + + @rule(key=zarr_keys(), data=st.data()) + def get_invalid_zarr_keys(self, key: str, data: DataObject) -> None: + note("(get_invalid)") + assume(key not in self.model) + assert self.store.get(key, self.prototype) is None + + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(data=st.data()) + def get_partial_values(self, data: DataObject) -> None: + key_range = data.draw( + key_ranges(keys=st.sampled_from(sorted(self.model.keys())), max_size=MAX_BINARY_SIZE) + ) + note(f"(get partial) {key_range=}") + obs_maybe = self.store.get_partial_values(key_range, self.prototype) + observed = [] + + for obs in obs_maybe: + assert obs is not None + observed.append(obs.to_bytes()) + + model_vals_ls = [] + + for key, byte_range in key_range: + start = byte_range.start + stop = byte_range.end + model_vals_ls.append(self.model[key][start:stop]) + + assert all( + obs == exp.to_bytes() for obs, exp in zip(observed, model_vals_ls, strict=True) + ), ( + observed, + model_vals_ls, + ) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(data=st.data()) + def delete(self, data: DataObject) -> None: + key = data.draw(st.sampled_from(sorted(self.model.keys()))) + note(f"(delete) Deleting {key=}") + + self.store.delete(key) + del self.model[key] + + @rule() + def clear(self) -> None: + assert not self.store.read_only + note("(clear)") + self.store.clear() + self.model.clear() + + assert self.store.is_empty("") + + assert len(self.model.keys()) == len(list(self.store.list())) == 0 + + @rule() + # Local store can be non-empty when there are subdirectories but no files + @precondition(lambda self: not isinstance(self.store.store, LocalStore)) + def is_empty(self) -> None: + note("(is_empty)") + + # make sure they either both are or both aren't empty (same state) + assert self.store.is_empty("") == (not self.model) + + @rule(key=zarr_keys()) + def exists(self, key: str) -> None: + note("(exists)") + + assert self.store.exists(key) == (key in self.model) + + @invariant() + def check_paths_equal(self) -> None: + note("Checking that paths are equal") + paths = sorted(self.store.list()) + + assert sorted(self.model.keys()) == paths + + @invariant() + def check_vals_equal(self) -> None: + note("Checking values equal") + for key, val in self.model.items(): + store_item = self.store.get(key, self.prototype) + assert val == store_item + + @invariant() + def check_num_zarr_keys_equal(self) -> None: + note("check num zarr_keys equal") + + assert len(self.model) == len(list(self.store.list())) + + @invariant() + def check_zarr_keys(self) -> None: + keys = list(self.store.list()) + + if not keys: + assert self.store.is_empty("") is True + + else: + assert self.store.is_empty("") is False + + for key in keys: + assert self.store.exists(key) is True + note("checking keys / exists / empty") diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/store.py b/.venv/lib/python3.11/site-packages/zarr/testing/store.py new file mode 100644 index 0000000000000000000000000000000000000000..ce83715b8649bbd2d04b85bfc90f42f9dcabecc8 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/store.py @@ -0,0 +1,712 @@ +from __future__ import annotations + +import asyncio +import json +import pickle +from abc import abstractmethod +from typing import TYPE_CHECKING, Generic, Self, TypeVar + +from zarr.storage import WrapperStore + +if TYPE_CHECKING: + from typing import Any + + from zarr.core.buffer.core import BufferPrototype + +import pytest + +from zarr.abc.store import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, +) +from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.sync import _collect_aiterator, sync +from zarr.storage._utils import _normalize_byte_range_index +from zarr.testing.utils import assert_bytes_equal + +__all__ = ["StoreTests"] + + +S = TypeVar("S", bound=Store) +B = TypeVar("B", bound=Buffer) + + +class StoreTests(Generic[S, B]): + store_cls: type[S] + buffer_cls: type[B] + + @staticmethod + def _require_get_sync(store: S) -> SupportsGetSync: + """Skip unless *store* implements :class:`SupportsGetSync`.""" + if not isinstance(store, SupportsGetSync): + pytest.skip("store does not implement SupportsGetSync") + return store # type: ignore[unreachable] + + @staticmethod + def _require_set_sync(store: S) -> SupportsSetSync: + """Skip unless *store* implements :class:`SupportsSetSync`.""" + if not isinstance(store, SupportsSetSync): + pytest.skip("store does not implement SupportsSetSync") + return store # type: ignore[unreachable] + + @staticmethod + def _require_delete_sync(store: S) -> SupportsDeleteSync: + """Skip unless *store* implements :class:`SupportsDeleteSync`.""" + if not isinstance(store, SupportsDeleteSync): + pytest.skip("store does not implement SupportsDeleteSync") + return store # type: ignore[unreachable] + + @abstractmethod + async def set(self, store: S, key: str, value: Buffer) -> None: + """ + Insert a value into a storage backend, with a specific key. + This should not use any store methods. Bypassing the store methods allows them to be + tested. + """ + ... + + @abstractmethod + async def get(self, store: S, key: str) -> Buffer: + """ + Retrieve a value from a storage backend, by key. + This should not use any store methods. Bypassing the store methods allows them to be + tested. + """ + ... + + @abstractmethod + @pytest.fixture + def store_kwargs(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Kwargs for instantiating a store""" + ... + + @abstractmethod + def test_store_repr(self, store: S) -> None: ... + + @abstractmethod + def test_store_supports_writes(self, store: S) -> None: ... + + def test_store_supports_partial_writes(self, store: S) -> None: + assert not store.supports_partial_writes + + @abstractmethod + def test_store_supports_listing(self, store: S) -> None: ... + + @pytest.fixture + def open_kwargs(self, store_kwargs: dict[str, Any]) -> dict[str, Any]: + return store_kwargs + + @pytest.fixture + async def store(self, open_kwargs: dict[str, Any]) -> Store: + return await self.store_cls.open(**open_kwargs) + + @pytest.fixture + async def store_not_open(self, store_kwargs: dict[str, Any]) -> Store: + return self.store_cls(**store_kwargs) + + def test_store_type(self, store: S) -> None: + assert isinstance(store, Store) + assert isinstance(store, self.store_cls) + + def test_store_eq(self, store: S, store_kwargs: dict[str, Any]) -> None: + # check self equality + assert store == store + + # check store equality with same inputs + # asserting this is important for being able to compare (de)serialized stores + store2 = self.store_cls(**store_kwargs) + assert store == store2 + + async def test_serializable_store(self, store: S) -> None: + new_store: S = pickle.loads(pickle.dumps(store)) + assert new_store == store + assert new_store.read_only == store.read_only + # quickly roundtrip data to a key to test that new store works + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "foo" + await store.set(key, data_buf) + observed = await store.get(key, prototype=default_buffer_prototype()) + assert_bytes_equal(observed, data_buf) + + def test_store_read_only(self, store: S) -> None: + assert not store.read_only + + with pytest.raises(AttributeError): + store.read_only = False # type: ignore[misc] + + @pytest.mark.parametrize("read_only", [True, False]) + async def test_store_open_read_only(self, open_kwargs: dict[str, Any], read_only: bool) -> None: + open_kwargs["read_only"] = read_only + store = await self.store_cls.open(**open_kwargs) + assert store._is_open + assert store.read_only == read_only + + async def test_store_context_manager(self, open_kwargs: dict[str, Any]) -> None: + # Test that the context manager closes the store + with await self.store_cls.open(**open_kwargs) as store: + assert store._is_open + # Test trying to open an already open store + with pytest.raises(ValueError, match="store is already open"): + await store._open() + assert not store._is_open + + async def test_read_only_store_raises(self, open_kwargs: dict[str, Any]) -> None: + kwargs = {**open_kwargs, "read_only": True} + store = await self.store_cls.open(**kwargs) + assert store.read_only + + # set + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + + # delete + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.delete("foo") + + async def test_with_read_only_store(self, open_kwargs: dict[str, Any]) -> None: + kwargs = {**open_kwargs, "read_only": True} + store = await self.store_cls.open(**kwargs) + assert store.read_only + + # Test that you cannot write to a read-only store + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + + # Check if the store implements with_read_only + try: + writer = store.with_read_only(read_only=False) + except NotImplementedError: + # Test that stores that do not implement with_read_only raise NotImplementedError with the correct message + with pytest.raises( + NotImplementedError, + match=f"with_read_only is not implemented for the {type(store)} store type.", + ): + store.with_read_only(read_only=False) + return + + # Test that you can write to a new store copy + assert not writer._is_open + assert not writer.read_only + await writer.set("foo", self.buffer_cls.from_bytes(b"bar")) + await writer.delete("foo") + + # Test that you cannot write to the original store + assert store.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.delete("foo") + + # Test that you cannot write to a read-only store copy + reader = store.with_read_only(read_only=True) + assert reader.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.set("foo", self.buffer_cls.from_bytes(b"bar")) + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.delete("foo") + + @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize( + ("data", "byte_range"), + [ + (b"\x01\x02\x03\x04", None), + (b"\x01\x02\x03\x04", RangeByteRequest(1, 4)), + (b"\x01\x02\x03\x04", OffsetByteRequest(1)), + (b"\x01\x02\x03\x04", SuffixByteRequest(1)), + (b"", None), + ], + ) + async def test_get(self, store: S, key: str, data: bytes, byte_range: ByteRequest) -> None: + """ + Ensure that data can be read from the store using the store.get method. + """ + data_buf = self.buffer_cls.from_bytes(data) + await self.set(store, key, data_buf) + observed = await store.get(key, prototype=default_buffer_prototype(), byte_range=byte_range) + start, stop = _normalize_byte_range_index(data_buf, byte_range=byte_range) + expected = data_buf[start:stop] + assert_bytes_equal(observed, expected) + + async def test_get_not_open(self, store_not_open: S) -> None: + """ + Ensure that data can be read from the store that isn't yet open using the store.get method. + """ + assert not store_not_open._is_open + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "c/0" + await self.set(store_not_open, key, data_buf) + observed = await store_not_open.get(key, prototype=default_buffer_prototype()) + assert_bytes_equal(observed, data_buf) + + async def test_get_raises(self, store: S) -> None: + """ + Ensure that a ValueError is raise for invalid byte range syntax + """ + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + await self.set(store, "c/0", data_buf) + with pytest.raises((ValueError, TypeError), match=r"Unexpected byte_range, got.*"): + await store.get("c/0", prototype=default_buffer_prototype(), byte_range=(0, 2)) # type: ignore[arg-type] + + async def test_get_many(self, store: S) -> None: + """ + Ensure that multiple keys can be retrieved at once with the _get_many method. + """ + keys = tuple(map(str, range(10))) + values = tuple(f"{k}".encode() for k in keys) + for k, v in zip(keys, values, strict=False): + await self.set(store, k, self.buffer_cls.from_bytes(v)) + observed_buffers = await _collect_aiterator( + store._get_many( + zip( + keys, + (default_buffer_prototype(),) * len(keys), + (None,) * len(keys), + strict=False, + ) + ) + ) + observed_kvs = sorted(((k, b.to_bytes()) for k, b in observed_buffers)) # type: ignore[union-attr] + expected_kvs = sorted(((k, b) for k, b in zip(keys, values, strict=False))) + assert observed_kvs == expected_kvs + + @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""]) + async def test_getsize(self, store: S, key: str, data: bytes) -> None: + """ + Test the result of store.getsize(). + """ + data_buf = self.buffer_cls.from_bytes(data) + expected = len(data_buf) + await self.set(store, key, data_buf) + observed = await store.getsize(key) + assert observed == expected + + async def test_getsize_prefix(self, store: S) -> None: + """ + Test the result of store.getsize_prefix(). + """ + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + keys = ["c/0/0", "c/0/1", "c/1/0", "c/1/1"] + keys_values = [(k, data_buf) for k in keys] + await store._set_many(keys_values) + expected = len(data_buf) * len(keys) + observed = await store.getsize_prefix("c") + assert observed == expected + + async def test_getsize_raises(self, store: S) -> None: + """ + Test that getsize() raise a FileNotFoundError if the key doesn't exist. + """ + with pytest.raises(FileNotFoundError): + await store.getsize("c/1000") + + @pytest.mark.parametrize("key", ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""]) + async def test_set(self, store: S, key: str, data: bytes) -> None: + """ + Ensure that data can be written to the store using the store.set method. + """ + assert not store.read_only + data_buf = self.buffer_cls.from_bytes(data) + await store.set(key, data_buf) + observed = await self.get(store, key) + assert_bytes_equal(observed, data_buf) + + async def test_set_not_open(self, store_not_open: S) -> None: + """ + Ensure that data can be written to the store that's not yet open using the store.set method. + """ + assert not store_not_open._is_open + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "c/0" + await store_not_open.set(key, data_buf) + observed = await self.get(store_not_open, key) + assert_bytes_equal(observed, data_buf) + + async def test_set_many(self, store: S) -> None: + """ + Test that a dict of key : value pairs can be inserted into the store via the + `_set_many` method. + """ + keys = ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"] + data_buf = [self.buffer_cls.from_bytes(k.encode()) for k in keys] + store_dict = dict(zip(keys, data_buf, strict=True)) + await store._set_many(store_dict.items()) + for k, v in store_dict.items(): + assert (await self.get(store, k)).to_bytes() == v.to_bytes() + + @pytest.mark.parametrize( + "key_ranges", + [ + [], + [("zarr.json", RangeByteRequest(0, 2))], + [("c/0", RangeByteRequest(0, 2)), ("zarr.json", None)], + [ + ("c/0/0", RangeByteRequest(0, 2)), + ("c/0/1", SuffixByteRequest(2)), + ("c/0/2", OffsetByteRequest(2)), + ], + ], + ) + async def test_get_partial_values( + self, store: S, key_ranges: list[tuple[str, ByteRequest]] + ) -> None: + # put all of the data + for key, _ in key_ranges: + await self.set(store, key, self.buffer_cls.from_bytes(bytes(key, encoding="utf-8"))) + + # read back just part of it + observed_maybe = await store.get_partial_values( + prototype=default_buffer_prototype(), key_ranges=key_ranges + ) + + observed: list[Buffer] = [] + expected: list[Buffer] = [] + + for obs in observed_maybe: + assert obs is not None + observed.append(obs) + + for idx in range(len(observed)): + key, byte_range = key_ranges[idx] + result = await store.get( + key, prototype=default_buffer_prototype(), byte_range=byte_range + ) + assert result is not None + expected.append(result) + + assert all( + obs.to_bytes() == exp.to_bytes() for obs, exp in zip(observed, expected, strict=True) + ) + + async def test_exists(self, store: S) -> None: + assert not await store.exists("foo") + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + assert await store.exists("foo/zarr.json") + + async def test_delete(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + assert await store.exists("foo/zarr.json") + await store.delete("foo/zarr.json") + assert not await store.exists("foo/zarr.json") + + async def test_delete_dir(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.set("zarr.json", self.buffer_cls.from_bytes(b"root")) + await store.set("foo-bar/zarr.json", self.buffer_cls.from_bytes(b"root")) + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + await store.set("foo/c/0", self.buffer_cls.from_bytes(b"chunk")) + await store.delete_dir("foo") + assert await store.exists("zarr.json") + assert await store.exists("foo-bar/zarr.json") + assert not await store.exists("foo/zarr.json") + assert not await store.exists("foo/c/0") + + async def test_delete_nonexistent_key_does_not_raise(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.delete("nonexistent_key") + + async def test_is_empty(self, store: S) -> None: + assert await store.is_empty("") + await self.set( + store, "foo/bar", self.buffer_cls.from_bytes(bytes("something", encoding="utf-8")) + ) + assert not await store.is_empty("") + assert await store.is_empty("fo") + assert not await store.is_empty("foo/") + assert not await store.is_empty("foo") + assert await store.is_empty("spam/") + + async def test_clear(self, store: S) -> None: + await self.set( + store, "key", self.buffer_cls.from_bytes(bytes("something", encoding="utf-8")) + ) + await store.clear() + assert await store.is_empty("") + + async def test_list(self, store: S) -> None: + assert await _collect_aiterator(store.list()) == () + prefix = "foo" + data = self.buffer_cls.from_bytes(b"") + store_dict = { + prefix + "/zarr.json": data, + **{prefix + f"/c/{idx}": data for idx in range(10)}, + } + await store._set_many(store_dict.items()) + expected_sorted = sorted(store_dict.keys()) + observed = await _collect_aiterator(store.list()) + observed_sorted = sorted(observed) + assert observed_sorted == expected_sorted + + async def test_list_prefix(self, store: S) -> None: + """ + Test that the `list_prefix` method works as intended. Given a prefix, it should return + all the keys in storage that start with this prefix. + """ + prefixes = ("", "a/", "a/b/", "a/b/c/") + data = self.buffer_cls.from_bytes(b"") + fname = "zarr.json" + store_dict = {p + fname: data for p in prefixes} + + await store._set_many(store_dict.items()) + + for prefix in prefixes: + observed = tuple(sorted(await _collect_aiterator(store.list_prefix(prefix)))) + expected: tuple[str, ...] = () + for key in store_dict: + if key.startswith(prefix): + expected += (key,) + expected = tuple(sorted(expected)) + assert observed == expected + + async def test_list_empty_path(self, store: S) -> None: + """ + Verify that list and list_prefix work correctly when path is an empty string, + i.e. no unwanted replacement occurs. + """ + data = self.buffer_cls.from_bytes(b"") + store_dict = { + "foo/bar/zarr.json": data, + "foo/bar/c/1": data, + "foo/baz/c/0": data, + } + await store._set_many(store_dict.items()) + + # Test list() + observed_list = await _collect_aiterator(store.list()) + observed_list_sorted = sorted(observed_list) + expected_list_sorted = sorted(store_dict.keys()) + assert observed_list_sorted == expected_list_sorted + + # Test list_prefix() with an empty prefix + observed_prefix_empty = await _collect_aiterator(store.list_prefix("")) + observed_prefix_empty_sorted = sorted(observed_prefix_empty) + expected_prefix_empty_sorted = sorted(store_dict.keys()) + assert observed_prefix_empty_sorted == expected_prefix_empty_sorted + + # Test list_prefix() with a non-empty prefix + observed_prefix = await _collect_aiterator(store.list_prefix("foo/bar/")) + observed_prefix_sorted = sorted(observed_prefix) + expected_prefix_sorted = sorted(k for k in store_dict if k.startswith("foo/bar/")) + assert observed_prefix_sorted == expected_prefix_sorted + + async def test_list_dir(self, store: S) -> None: + roots_and_keys: list[tuple[str, dict[str, Buffer]]] = [ + ( + "foo", + { + "foo/zarr.json": self.buffer_cls.from_bytes(b"bar"), + "foo/c/1": self.buffer_cls.from_bytes(b"\x01"), + }, + ), + ( + "foo/bar", + { + "foo/bar/foobar_first_child": self.buffer_cls.from_bytes(b"1"), + "foo/bar/foobar_second_child/zarr.json": self.buffer_cls.from_bytes(b"2"), + }, + ), + ] + + assert await _collect_aiterator(store.list_dir("")) == () + + for root, store_dict in roots_and_keys: + assert await _collect_aiterator(store.list_dir(root)) == () + + await store._set_many(store_dict.items()) + + keys_observed = await _collect_aiterator(store.list_dir(root)) + keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict} + assert sorted(keys_observed) == sorted(keys_expected) + + keys_observed = await _collect_aiterator(store.list_dir(root + "/")) + assert sorted(keys_expected) == sorted(keys_observed) + + async def test_set_if_not_exists(self, store: S) -> None: + key = "k" + data_buf = self.buffer_cls.from_bytes(b"0000") + await self.set(store, key, data_buf) + + new = self.buffer_cls.from_bytes(b"1111") + await store.set_if_not_exists("k", new) # no error + + result = await store.get(key, default_buffer_prototype()) + assert result == data_buf + + await store.set_if_not_exists("k2", new) # no error + + result = await store.get("k2", default_buffer_prototype()) + assert result == new + + async def test_get_bytes(self, store: S) -> None: + """ + Test that the get_bytes method reads bytes. + """ + data = b"hello world" + key = "zarr.json" + await self.set(store, key, self.buffer_cls.from_bytes(data)) + assert await store._get_bytes(key, prototype=default_buffer_prototype()) == data + with pytest.raises(FileNotFoundError): + await store._get_bytes("nonexistent_key", prototype=default_buffer_prototype()) + + def test_get_bytes_sync(self, store: S) -> None: + """ + Test that the get_bytes_sync method reads bytes. + """ + data = b"hello world" + key = "zarr.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(data))) + assert store._get_bytes_sync(key, prototype=default_buffer_prototype()) == data + + async def test_get_json(self, store: S) -> None: + """ + Test that the get_json method reads json. + """ + data = {"foo": "bar"} + data_bytes = json.dumps(data).encode("utf-8") + key = "zarr.json" + await self.set(store, key, self.buffer_cls.from_bytes(data_bytes)) + assert await store._get_json(key, prototype=default_buffer_prototype()) == data + + def test_get_json_sync(self, store: S) -> None: + """ + Test that the get_json method reads json. + """ + data = {"foo": "bar"} + data_bytes = json.dumps(data).encode("utf-8") + key = "zarr.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(data_bytes))) + assert store._get_json_sync(key, prototype=default_buffer_prototype()) == data + + # ------------------------------------------------------------------- + # Synchronous store methods (SupportsSyncStore protocol) + # ------------------------------------------------------------------- + + def test_get_sync(self, store: S) -> None: + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_get" + sync(self.set(store, key, data_buf)) + result = getter.get_sync(key) + assert result is not None + assert_bytes_equal(result, data_buf) + + def test_get_sync_missing(self, store: S) -> None: + getter = self._require_get_sync(store) + result = getter.get_sync("nonexistent") + assert result is None + + def test_set_sync(self, store: S) -> None: + setter = self._require_set_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_set" + setter.set_sync(key, data_buf) + result = sync(self.get(store, key)) + assert_bytes_equal(result, data_buf) + + def test_delete_sync(self, store: S) -> None: + setter = self._require_set_sync(store) + deleter = self._require_delete_sync(store) + getter = self._require_get_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_delete" + setter.set_sync(key, data_buf) + deleter.delete_sync(key) + result = getter.get_sync(key) + assert result is None + + def test_delete_sync_missing(self, store: S) -> None: + deleter = self._require_delete_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + # should not raise + deleter.delete_sync("nonexistent_sync") + + +class LatencyStore(WrapperStore[Store]): + """ + A wrapper class that takes any store class in its constructor and + adds latency to the `set` and `get` methods. This can be used for + performance testing. + """ + + get_latency: float + set_latency: float + + def __init__(self, store: Store, *, get_latency: float = 0, set_latency: float = 0) -> None: + self.get_latency = float(get_latency) + self.set_latency = float(set_latency) + self._store = store + + def _with_store(self, store: Store) -> Self: + return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency) + + async def set(self, key: str, value: Buffer) -> None: + """ + Add latency to the ``set`` method. + + Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. + + Parameters + ---------- + key : str + The key to set + value : Buffer + The value to set + + Returns + ------- + None + """ + await asyncio.sleep(self.set_latency) + await self._store.set(key, value) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """ + Add latency to the ``get`` method. + + Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. + + Parameters + ---------- + key : str + The key to get + prototype : BufferPrototype + The BufferPrototype to use. + byte_range : ByteRequest, optional + An optional byte range. + + Returns + ------- + buffer : Buffer or None + """ + await asyncio.sleep(self.get_latency) + return await self._store.get(key, prototype=prototype, byte_range=byte_range) diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/strategies.py b/.venv/lib/python3.11/site-packages/zarr/testing/strategies.py new file mode 100644 index 0000000000000000000000000000000000000000..330f220b560d54eb59cb8e4237f636be0ccb3384 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/strategies.py @@ -0,0 +1,448 @@ +import math +import sys +from collections.abc import Callable, Mapping +from typing import Any, Literal + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +import numpy.typing as npt +from hypothesis import event +from hypothesis.strategies import SearchStrategy + +import zarr +from zarr.abc.store import RangeByteRequest, Store +from zarr.codecs.bytes import BytesCodec +from zarr.core.array import Array +from zarr.core.chunk_grids import RegularChunkGrid +from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding +from zarr.core.common import JSON, ZarrFormat +from zarr.core.dtype import get_data_type_from_native_dtype +from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.sync import sync +from zarr.storage import MemoryStore, StoreLike +from zarr.storage._common import _dereference_path +from zarr.storage._utils import normalize_path +from zarr.types import AnyArray + +TrueOrFalse = Literal[True, False] + +# Copied from Xarray +_attr_keys = st.text(st.characters(), min_size=1) +_attr_values = st.recursive( + st.none() | st.booleans() | st.text(st.characters(), max_size=5), + lambda children: st.lists(children) | st.dictionaries(_attr_keys, children), + max_leaves=3, +) + + +@st.composite +def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str: + return draw(st.lists(node_names, min_size=1, max_size=max_num_nodes).map("/".join)) + + +@st.composite +def paths(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str: + return draw(st.just("/") | keys(max_num_nodes=max_num_nodes)) + + +def dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return ( + npst.boolean_dtypes() + | npst.integer_dtypes(endianness="=") + | npst.unsigned_integer_dtypes(endianness="=") + | npst.floating_dtypes(endianness="=") + | npst.complex_number_dtypes(endianness="=") + | npst.byte_string_dtypes(endianness="=") + | npst.unicode_string_dtypes(endianness="=") + | npst.datetime64_dtypes(endianness="=") + | npst.timedelta64_dtypes(endianness="=") + ) + + +def v3_dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return dtypes() + + +def v2_dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return dtypes() + + +def safe_unicode_for_dtype(dtype: np.dtype[np.str_]) -> st.SearchStrategy[str]: + """Generate UTF-8-safe text constrained to max_len of dtype.""" + # account for utf-32 encoding (i.e. 4 bytes/character) + max_len = max(1, dtype.itemsize // 4) + + return st.text( + alphabet=st.characters( + exclude_categories=["Cs"], # Avoid *technically allowed* surrogates + min_codepoint=32, + ), + min_size=1, + max_size=max_len, + ) + + +def clear_store(x: Store) -> Store: + sync(x.clear()) + return x + + +# From https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#node-names +# 1. must not be the empty string ("") +# 2. must not include the character "/" +# 3. must not be a string composed only of period characters, e.g. "." or ".." +# 4. must not start with the reserved prefix "__" +zarr_key_chars = st.sampled_from( + ".-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz" +) +node_names = ( + st.text(zarr_key_chars, min_size=1) + .filter(lambda t: t not in (".", "..") and not t.startswith("__")) + .filter(lambda name: name.lower() != "zarr.json") +) +short_node_names = ( + st.text(zarr_key_chars, max_size=3, min_size=1) + .filter(lambda t: t not in (".", "..") and not t.startswith("__")) + .filter(lambda name: name.lower() != "zarr.json") +) +array_names = node_names +attrs: st.SearchStrategy[Mapping[str, JSON] | None] = st.none() | st.dictionaries( + _attr_keys, _attr_values +) +# st.builds will only call a new store constructor for different keyword arguments +# i.e. stores.examples() will always return the same object per Store class. +# So we map a clear to reset the store. +stores = st.builds(MemoryStore, st.just({})).map(clear_store) +compressors = st.sampled_from([None, "default"]) +zarr_formats: st.SearchStrategy[ZarrFormat] = st.sampled_from([3, 2]) +# We de-prioritize arrays having dim sizes 0, 1, 2 +array_shapes = npst.array_shapes(max_dims=4, min_side=3, max_side=5) | npst.array_shapes( + max_dims=4, min_side=0 +) + + +@st.composite +def dimension_names(draw: st.DrawFn, *, ndim: int | None = None) -> list[None | str] | None: + simple_text = st.text(zarr_key_chars, min_size=0) + return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim)) # type: ignore[arg-type] + + +@st.composite +def array_metadata( + draw: st.DrawFn, + *, + array_shapes: Callable[..., st.SearchStrategy[tuple[int, ...]]] = npst.array_shapes, + zarr_formats: st.SearchStrategy[ZarrFormat] = zarr_formats, + attributes: SearchStrategy[Mapping[str, JSON] | None] = attrs, +) -> ArrayV2Metadata | ArrayV3Metadata: + zarr_format = draw(zarr_formats) + # separator = draw(st.sampled_from(['/', '\\'])) + shape = draw(array_shapes()) + ndim = len(shape) + chunk_shape = draw(array_shapes(min_dims=ndim, max_dims=ndim)) + np_dtype = draw(dtypes()) + dtype = get_data_type_from_native_dtype(np_dtype) + fill_value = draw(npst.from_dtype(np_dtype)) + if zarr_format == 2: + return ArrayV2Metadata( + shape=shape, + chunks=chunk_shape, + dtype=dtype, + fill_value=fill_value, + order=draw(st.sampled_from(["C", "F"])), + attributes=draw(attributes), # type: ignore[arg-type] + dimension_separator=draw(st.sampled_from([".", "/"])), + filters=None, + compressor=None, + ) + else: + return ArrayV3Metadata( + shape=shape, + data_type=dtype, + chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape), + fill_value=fill_value, + attributes=draw(attributes), # type: ignore[arg-type] + dimension_names=draw(dimension_names(ndim=ndim)), + chunk_key_encoding=DefaultChunkKeyEncoding(separator="/"), # FIXME + codecs=[BytesCodec()], + storage_transformers=(), + ) + + +@st.composite +def numpy_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, + dtype: np.dtype[Any] | None = None, +) -> npt.NDArray[Any]: + """ + Generate numpy arrays that can be saved in the provided Zarr format. + """ + if dtype is None: + dtype = draw(dtypes()) + if np.issubdtype(dtype, np.str_): + safe_unicode_strings = safe_unicode_for_dtype(dtype) + return draw(npst.arrays(dtype=dtype, shape=shapes, elements=safe_unicode_strings)) + + return draw(npst.arrays(dtype=dtype, shape=shapes)) + + +@st.composite +def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]: + # We want this strategy to shrink towards arrays with smaller number of chunks + # 1. st.integers() shrinks towards smaller values. So we use that to generate number of chunks + numchunks = draw( + st.tuples(*[st.integers(min_value=0 if size == 0 else 1, max_value=size) for size in shape]) + ) + # 2. and now generate the chunks tuple + chunks = tuple( + size // nchunks if nchunks > 0 else 0 + for size, nchunks in zip(shape, numchunks, strict=True) + ) + + for c in chunks: + event("chunk size", c) + + if any((c != 0 and s % c != 0) for s, c in zip(shape, chunks, strict=True)): + event("smaller last chunk") + + return chunks + + +@st.composite +def shard_shapes( + draw: st.DrawFn, *, shape: tuple[int, ...], chunk_shape: tuple[int, ...] +) -> tuple[int, ...]: + # We want this strategy to shrink towards arrays with smaller number of shards + # shards must be an integral number of chunks + assert all(c != 0 for c in chunk_shape) + numchunks = tuple(s // c for s, c in zip(shape, chunk_shape, strict=True)) + multiples = tuple(draw(st.integers(min_value=1, max_value=nc)) for nc in numchunks) + return tuple(m * c for m, c in zip(multiples, chunk_shape, strict=True)) + + +@st.composite +def np_array_and_chunks( + draw: st.DrawFn, + *, + arrays: st.SearchStrategy[npt.NDArray[Any]] = numpy_arrays(), # noqa: B008 +) -> tuple[np.ndarray, tuple[int, ...]]: # type: ignore[type-arg] + """A hypothesis strategy to generate small sized random arrays. + + Returns: a tuple of the array and a suitable random chunking for it. + """ + array = draw(arrays) + return (array, draw(chunk_shapes(shape=array.shape))) + + +@st.composite +def arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, + compressors: st.SearchStrategy = compressors, + stores: st.SearchStrategy[StoreLike] = stores, + paths: st.SearchStrategy[str] = paths(), # noqa: B008 + array_names: st.SearchStrategy = array_names, + arrays: st.SearchStrategy | None = None, + attrs: st.SearchStrategy = attrs, + zarr_formats: st.SearchStrategy = zarr_formats, +) -> AnyArray: + store = draw(stores, label="store") + path = draw(paths, label="array parent") + name = draw(array_names, label="array name") + attributes = draw(attrs, label="attributes") + zarr_format = draw(zarr_formats, label="zarr format") + if arrays is None: + arrays = numpy_arrays(shapes=shapes) + nparray = draw(arrays, label="array data") + chunk_shape = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") + dim_names: None | list[str | None] = None + if zarr_format == 3 and all(c > 0 for c in chunk_shape): + shard_shape = draw( + st.none() | shard_shapes(shape=nparray.shape, chunk_shape=chunk_shape), + label="shard shape", + ) + dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names") + else: + shard_shape = None + # test that None works too. + fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)])) + # compressor = draw(compressors) + + expected_attrs = {} if attributes is None else attributes + + array_path = _dereference_path(path, name) + root = zarr.open_group(store, mode="w", zarr_format=zarr_format) + + a = root.create_array( + array_path, + shape=nparray.shape, + chunks=chunk_shape, + shards=shard_shape, + dtype=nparray.dtype, + attributes=attributes, + # compressor=compressor, # FIXME + fill_value=fill_value, + dimension_names=dim_names, + ) + + assert isinstance(a, Array) + if a.metadata.zarr_format == 3: + assert a.fill_value is not None + assert a.name is not None + assert a.path == normalize_path(array_path) + assert a.name == "/" + a.path + assert isinstance(root[array_path], Array) + assert nparray.shape == a.shape + assert chunk_shape == a.chunks + assert shard_shape == a.shards + assert a.basename == name, (a.basename, name) + assert dict(a.attrs) == expected_attrs + + a[:] = nparray + + return a + + +@st.composite +def simple_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, +) -> Any: + return draw( + arrays( + shapes=shapes, + paths=paths(max_num_nodes=2), + array_names=short_node_names, + attrs=st.none(), + compressors=st.sampled_from([None, "default"]), + ) + ) + + +def is_negative_slice(idx: Any) -> bool: + return isinstance(idx, slice) and idx.step is not None and idx.step < 0 + + +@st.composite +def end_slices(draw: st.DrawFn, *, shape: tuple[int, ...]) -> Any: + """ + A strategy that slices ranges that include the last chunk. + This is intended to stress-test handling of a possibly smaller last chunk. + """ + slicers = [] + for size in shape: + start = draw(st.integers(min_value=size // 2, max_value=size - 1)) + length = draw(st.integers(min_value=0, max_value=size - start)) + slicers.append(slice(start, start + length)) + event("drawing end slice") + return tuple(slicers) + + +@st.composite +def basic_indices( + draw: st.DrawFn, + *, + shape: tuple[int, ...], + min_dims: int = 0, + max_dims: int | None = None, + allow_newaxis: TrueOrFalse = False, + allow_ellipsis: TrueOrFalse = True, +) -> Any: + """Basic indices without unsupported negative slices.""" + strategy = npst.basic_indices( + shape=shape, + min_dims=min_dims, + max_dims=max_dims, + allow_newaxis=allow_newaxis, + allow_ellipsis=allow_ellipsis, + ).filter( + lambda idxr: ( + not ( + is_negative_slice(idxr) + or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr)) + ) + ) + ) + if math.prod(shape) >= 3: + strategy = end_slices(shape=shape) | strategy + return draw(strategy) + + +@st.composite +def orthogonal_indices( + draw: st.DrawFn, *, shape: tuple[int, ...] +) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: + """ + Strategy that returns + (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays. + (2) a tuple of integer arrays that can be used for equivalent indexing of numpy arrays + """ + zindexer = [] + npindexer = [] + ndim = len(shape) + for axis, size in enumerate(shape): + if size != 0: + strategy = npst.integer_array_indices( + shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1) + ) | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + else: + strategy = basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + + val = draw( + strategy + # bare ints, slices + .map(lambda x: (x,) if not isinstance(x, tuple) else x) + # skip empty tuple + .filter(bool) + ) + (idxr,) = val + if isinstance(idxr, int): + idxr = np.array([idxr]) + zindexer.append(idxr) + if isinstance(idxr, slice): + idxr = np.arange(*idxr.indices(size)) + elif isinstance(idxr, (tuple, int)): + idxr = np.array(idxr) + newshape = [1] * ndim + newshape[axis] = idxr.size + npindexer.append(idxr.reshape(newshape)) + + # casting the output of broadcast_arrays is needed for numpy < 2 + return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) + + +def key_ranges( + keys: SearchStrategy[str] = node_names, max_size: int = sys.maxsize +) -> SearchStrategy[list[tuple[str, RangeByteRequest]]]: + """ + Function to generate key_ranges strategy for get_partial_values() + returns list strategy w/ form:: + + [(key, (range_start, range_end)), + (key, (range_start, range_end)),...] + """ + + def make_request(start: int, length: int) -> RangeByteRequest: + return RangeByteRequest(start, end=min(start + length, max_size)) + + byte_ranges = st.builds( + make_request, + start=st.integers(min_value=0, max_value=max_size), + length=st.integers(min_value=0, max_value=max_size), + ) + key_tuple = st.tuples(keys, byte_ranges) + return st.lists(key_tuple, min_size=1, max_size=10) + + +@st.composite +def chunk_paths(draw: st.DrawFn, ndim: int, numblocks: tuple[int, ...], subset: bool = True) -> str: + blockidx = draw( + st.tuples(*tuple(st.integers(min_value=0, max_value=max(0, b - 1)) for b in numblocks)) + ) + subset_slicer = slice(draw(st.integers(min_value=0, max_value=ndim))) if subset else slice(None) + return "/".join(map(str, blockidx[subset_slicer])) diff --git a/.venv/lib/python3.11/site-packages/zarr/testing/utils.py b/.venv/lib/python3.11/site-packages/zarr/testing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4c3e45c54fc67209ac35e059ee639b81cd7f23 --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/testing/utils.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar, cast + +import pytest + +from zarr.core.buffer import Buffer + +if TYPE_CHECKING: + from zarr.core.common import BytesLike + +__all__ = ["assert_bytes_equal"] + + +def assert_bytes_equal(b1: Buffer | BytesLike | None, b2: Buffer | BytesLike | None) -> None: + """Help function to assert if two bytes-like or Buffers are equal + + Warnings + -------- + Always copies data, only use for testing and debugging + """ + if isinstance(b1, Buffer): + b1 = b1.to_bytes() + if isinstance(b2, Buffer): + b2 = b2.to_bytes() + assert b1 == b2 + + +def has_cupy() -> bool: + try: + import cupy + + return cast("bool", cupy.cuda.runtime.getDeviceCount() > 0) + except ImportError: + return False + except cupy.cuda.runtime.CUDARuntimeError: + return False + + +T = TypeVar("T") + + +gpu_mark = pytest.mark.gpu +skip_if_no_gpu = pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available") + + +# Decorator for GPU tests +def gpu_test(func: T) -> T: + return cast(T, gpu_mark(skip_if_no_gpu(func))) diff --git a/.venv/lib/python3.11/site-packages/zarr/types.py b/.venv/lib/python3.11/site-packages/zarr/types.py new file mode 100644 index 0000000000000000000000000000000000000000..38990982f91dc7d8e409fc8bfad5bc2cf0eecd6d --- /dev/null +++ b/.venv/lib/python3.11/site-packages/zarr/types.py @@ -0,0 +1,23 @@ +from typing import Any, TypeAlias + +from zarr.core.array import Array, AsyncArray +from zarr.core.metadata.v2 import ArrayV2Metadata +from zarr.core.metadata.v3 import ArrayV3Metadata + +AnyAsyncArray: TypeAlias = AsyncArray[Any] +"""A Zarr format 2 or 3 `AsyncArray`""" + +AsyncArrayV2: TypeAlias = AsyncArray[ArrayV2Metadata] +"""A Zarr format 2 `AsyncArray`""" + +AsyncArrayV3: TypeAlias = AsyncArray[ArrayV3Metadata] +"""A Zarr format 3 `AsyncArray`""" + +AnyArray: TypeAlias = Array[Any] +"""A Zarr format 2 or 3 `Array`""" + +ArrayV2: TypeAlias = Array[ArrayV2Metadata] +"""A Zarr format 2 `Array`""" + +ArrayV3: TypeAlias = Array[ArrayV3Metadata] +"""A Zarr format 3 `Array`"""