Enforce ruff/flake8-type-checking rules (TCH) (#2110)

* Apply ruff/flake8-type-checking rule TCH001

TCH001 Move application import into a type-checking block

* Apply ruff/flake8-type-checking rule TCH002

TCH002 Move third-party import into a type-checking block

* Apply ruff/flake8-type-checking rule TCH003

TCH003 Move standard library import into a type-checking block

* Apply ruff/flake8-type-checking rule TCH005

TCH005 Found empty type-checking block

* Enforce ruff/flake8-type-checking rules (TCH)

pyproject.toml

@@ -212,6 +212,7 @@ extend-select = [
     "UP",  # pyupgrade
     "RSE",
     "RUF",
+    "TCH", # flake8-type-checking
     "TRY", # tryceratops
 ]
 ignore = [

src/zarr/abc/codec.py

@@ -1,22 +1,22 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Awaitable, Callable, Iterable
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-import numpy as np
-
 from zarr.abc.metadata import Metadata
-from zarr.abc.store import ByteGetter, ByteSetter
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.chunk_grids import ChunkGrid
 from zarr.core.common import ChunkCoords, concurrent_map
 from zarr.core.config import config
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable, Iterable
+
+    import numpy as np
     from typing_extensions import Self
 
+    from zarr.abc.store import ByteGetter, ByteSetter
     from zarr.core.array_spec import ArraySpec
+    from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.common import JSON
     from zarr.core.indexing import SelectorTuple
 

src/zarr/abc/metadata.py

@@ -6,9 +6,9 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-from dataclasses import dataclass, fields
+    from zarr.core.common import JSON
 
-from zarr.core.common import JSON
+from dataclasses import dataclass, fields
 
 __all__ = ["Metadata"]
 

src/zarr/api/asynchronous.py

@@ -2,16 +2,12 @@
 
 import asyncio
 import warnings
-from collections.abc import Iterable
-from typing import Any, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Literal, Union, cast
 
 import numpy as np
 import numpy.typing as npt
 
-from zarr.abc.codec import Codec
 from zarr.core.array import Array, AsyncArray
-from zarr.core.buffer import NDArrayLike
-from zarr.core.chunk_key_encodings import ChunkKeyEncoding
 from zarr.core.common import JSON, AccessModeLiteral, ChunkCoords, MemoryOrder, ZarrFormat
 from zarr.core.group import AsyncGroup
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
@@ -20,6 +16,13 @@
     make_store_path,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from zarr.abc.codec import Codec
+    from zarr.core.buffer import NDArrayLike
+    from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+
 __all__ = [
     "consolidate_metadata",
     "copy",

src/zarr/api/synchronous.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import zarr.api.asynchronous as async_api
 from zarr.core.array import Array, AsyncArray
-from zarr.core.buffer import NDArrayLike
-from zarr.core.common import JSON, AccessModeLiteral, ChunkCoords, ZarrFormat
 from zarr.core.group import Group
 from zarr.core.sync import sync
-from zarr.store import StoreLike
+
+if TYPE_CHECKING:
+    from zarr.core.buffer import NDArrayLike
+    from zarr.core.common import JSON, AccessModeLiteral, ChunkCoords, ZarrFormat
+    from zarr.store import StoreLike
 
 __all__ = [
     "consolidate_metadata",

src/zarr/codecs/_v2.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import numcodecs
 from numcodecs.compat import ensure_bytes, ensure_ndarray
 
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec
-from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import Buffer, NDBuffer, default_buffer_prototype
 from zarr.core.common import JSON, to_thread
 from zarr.registry import get_ndbuffer_class
 
+if TYPE_CHECKING:
+    from zarr.core.array_spec import ArraySpec
+
 
 @dataclass(frozen=True)
 class V2Compressor(ArrayBytesCodec):

src/zarr/codecs/blosc.py

@@ -9,15 +9,16 @@
 from numcodecs.blosc import Blosc
 
 from zarr.abc.codec import BytesBytesCodec
-from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_enum, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer
+
 
 class BloscShuffle(Enum):
     noshuffle = "noshuffle"

src/zarr/codecs/bytes.py

@@ -8,14 +8,15 @@
 import numpy as np
 
 from zarr.abc.codec import ArrayBytesCodec
-from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from zarr.core.array_spec import ArraySpec
+
 
 class Endian(Enum):
     big = "big"

src/zarr/codecs/crc32c_.py

@@ -8,14 +8,15 @@
 from crc32c import crc32c
 
 from zarr.abc.codec import BytesBytesCodec
-from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer
 from zarr.core.common import JSON, parse_named_configuration
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer
+
 
 @dataclass(frozen=True)
 class Crc32cCodec(BytesBytesCodec):

src/zarr/codecs/gzip.py

@@ -6,15 +6,16 @@
 from numcodecs.gzip import GZip
 
 from zarr.abc.codec import BytesBytesCodec
-from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer
+
 
 def parse_gzip_level(data: JSON) -> int:
     if not isinstance(data, (int)):

src/zarr/codecs/pipeline.py

@@ -6,8 +6,6 @@
 from typing import TYPE_CHECKING, Any, TypeVar
 from warnings import warn
 
-import numpy as np
-
 from zarr.abc.codec import (
     ArrayArrayCodec,
     ArrayBytesCodec,
@@ -17,18 +15,19 @@
     Codec,
     CodecPipeline,
 )
-from zarr.abc.store import ByteGetter, ByteSetter
-from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer
-from zarr.core.chunk_grids import ChunkGrid
 from zarr.core.common import JSON, ChunkCoords, concurrent_map, parse_named_configuration
 from zarr.core.config import config
 from zarr.core.indexing import SelectorTuple, is_scalar, is_total_slice
 from zarr.registry import get_codec_class, register_pipeline
 
 if TYPE_CHECKING:
+    import numpy as np
     from typing_extensions 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
 
 T = TypeVar("T")
 U = TypeVar("U")

src/zarr/codecs/transpose.py

@@ -8,8 +8,6 @@
 
 from zarr.abc.codec import ArrayArrayCodec
 from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import NDBuffer
-from zarr.core.chunk_grids import ChunkGrid
 from zarr.core.common import JSON, ChunkCoordsLike, parse_named_configuration
 from zarr.registry import register_codec
 
@@ -18,6 +16,9 @@
 
     from typing_extensions import Self
 
+    from zarr.core.buffer import NDBuffer
+    from zarr.core.chunk_grids import ChunkGrid
+
 
 def parse_transpose_order(data: JSON | Iterable[int]) -> tuple[int, ...]:
     if not isinstance(data, Iterable):

src/zarr/codecs/zstd.py

@@ -8,15 +8,16 @@
 from numcodecs.zstd import Zstd
 
 from zarr.abc.codec import BytesBytesCodec
-from zarr.core.array_spec import ArraySpec
-from zarr.core.buffer import Buffer
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from zarr.core.array_spec import ArraySpec
+    from zarr.core.buffer import Buffer
+
 
 def parse_zstd_level(data: JSON) -> int:
     if isinstance(data, int):

src/zarr/core/array.py

@@ -2,14 +2,12 @@
 
 import json
 from asyncio import gather
-from collections.abc import Iterable
 from dataclasses import dataclass, field, replace
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import numpy.typing as npt
 
-from zarr.abc.codec import Codec, CodecPipeline
 from zarr.abc.store import set_or_delete
 from zarr.codecs import BytesCodec
 from zarr.codecs._v2 import V2Compressor, V2Filters
@@ -63,6 +61,11 @@
     ensure_no_existing_node,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from zarr.abc.codec import Codec, CodecPipeline
+
 
 def parse_array_metadata(data: Any) -> ArrayV2Metadata | ArrayV3Metadata:
     if isinstance(data, ArrayV2Metadata | ArrayV3Metadata):

src/zarr/core/array_spec.py

@@ -3,11 +3,11 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal
 
-import numpy as np
-
 from zarr.core.common import parse_dtype, parse_fill_value, parse_order, parse_shapelike
 
 if TYPE_CHECKING:
+    import numpy as np
+
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import ChunkCoords
 

src/zarr/core/buffer/core.py

@@ -2,7 +2,6 @@
 
 import sys
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -17,17 +16,18 @@
 import numpy as np
 import numpy.typing as npt
 
-from zarr.core.common import ChunkCoords
 from zarr.registry import (
     get_buffer_class,
     get_ndbuffer_class,
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+
     from typing_extensions import Self
 
     from zarr.codecs.bytes import Endian
-    from zarr.core.common import BytesLike
+    from zarr.core.common import BytesLike, ChunkCoords
 
 
 @runtime_checkable

src/zarr/core/buffer/cpu.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -11,15 +10,17 @@
 import numpy.typing as npt
 
 from zarr.core.buffer import core
-from zarr.core.buffer.core import ArrayLike, NDArrayLike
 from zarr.registry import (
     register_buffer,
     register_ndbuffer,
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
     from typing_extensions import Self
 
+    from zarr.core.buffer.core import ArrayLike, NDArrayLike
     from zarr.core.common import BytesLike
 
 

src/zarr/core/buffer/gpu.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Iterable
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -16,6 +15,8 @@
 from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
+
     from typing_extensions import Self
 
     from zarr.core.common import BytesLike

src/zarr/core/chunk_grids.py

@@ -4,7 +4,6 @@
 import math
 import operator
 from abc import abstractmethod
-from collections.abc import Iterator
 from dataclasses import dataclass
 from functools import reduce
 from typing import TYPE_CHECKING
@@ -22,6 +21,8 @@
 from zarr.core.indexing import ceildiv
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from typing_extensions import Self