From 88ba0cd6089ddc6d98a5bbfbe8e7395944757f88 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <deepak@cherian.net>
Date: Thu, 18 Apr 2024 15:39:43 -0600
Subject: [PATCH] Fix `is_total_slice` for size-1 dimensions

Closes #1730

Co-authored-by: Ryan Abernathey <ryan.abernathey@gmail.com>
---
 docs/release.rst        |  2 ++
 zarr/tests/test_util.py |  9 +++++++++
 zarr/util.py            | 13 +++++++++++--
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/docs/release.rst b/docs/release.rst
index 75193bc3e3..e8ec447673 100644
--- a/docs/release.rst
+++ b/docs/release.rst
@@ -20,6 +20,8 @@ Unreleased
 
 Enhancements
 ~~~~~~~~~~~~
+* Performance improvement for reading and writing chunks if any of the dimensions is size 1. :issue:`1730`
+  By :user:`Deepak Cherian <dcherian>`.
 
 
 Docs
diff --git a/zarr/tests/test_util.py b/zarr/tests/test_util.py
index 1f7efc9214..d908c7b2d7 100644
--- a/zarr/tests/test_util.py
+++ b/zarr/tests/test_util.py
@@ -89,6 +89,15 @@ def test_is_total_slice():
     assert not is_total_slice((slice(0, 50), slice(0, 50)), (100, 100))
     assert not is_total_slice((slice(0, 100, 2), slice(0, 100)), (100, 100))
 
+    # size-1 dimension edge-case
+    # https://github.com/zarr-developers/zarr-python/issues/1730
+    assert is_total_slice((slice(0, 1),), (1,))
+    # this is an equivalent selection (without a slice)
+    assert is_total_slice((0,), (1,))
+    # same for multidimensional selection
+    assert is_total_slice((slice(0, 1), slice(0, 10)), (1, 10))
+    assert is_total_slice((0, slice(0, 10)), (1, 10))
+
     with pytest.raises(TypeError):
         is_total_slice("foo", (100,))
 
diff --git a/zarr/util.py b/zarr/util.py
index 848f1ed114..e58aed80ab 100644
--- a/zarr/util.py
+++ b/zarr/util.py
@@ -234,8 +234,17 @@ def is_total_slice(item, shape: Tuple[int]) -> bool:
     if isinstance(item, tuple):
         return all(
             (
-                isinstance(it, slice)
-                and ((it == slice(None)) or ((it.stop - it.start == sh) and (it.step in [1, None])))
+                (
+                    isinstance(it, slice)
+                    and (
+                        (it == slice(None))
+                        or ((it.stop - it.start == sh) and (it.step in [1, None]))
+                    )
+                )
+                # The only scalar edge case, indexing with int 0 along a size-1 dimension
+                # is identical to a total slice
+                # https://github.com/zarr-developers/zarr-python/issues/1730
+                or (isinstance(it, int) and it == 0 and sh == 1)
             )
             for it, sh in zip(item, shape)
         )