Add more typing to zarr.group (#1870)

Co-authored-by: Joe Hamman <joe@earthmover.io>

Author: dstansby

src/zarr/group.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterator
 from dataclasses import asdict, dataclass, field, replace
 
 import asyncio
@@ -256,7 +256,7 @@ async def _save_metadata(self) -> None:
         await asyncio.gather(*awaitables)
 
     @property
-    def attrs(self):
+    def attrs(self) -> dict[str, Any]:
         return self.metadata.attributes
 
     @property
@@ -301,7 +301,7 @@ async def create_array(
             zarr_format=self.metadata.zarr_format,
         )
 
-    async def update_attributes(self, new_attributes: dict[str, Any]):
+    async def update_attributes(self, new_attributes: dict[str, Any]) -> "AsyncGroup":
         # metadata.attributes is "frozen" so we simply clear and update the dict
         self.metadata.attributes.clear()
         self.metadata.attributes.update(new_attributes)
@@ -319,7 +319,7 @@ async def update_attributes(self, new_attributes: dict[str, Any]):
 
         return self
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<AsyncGroup {self.store_path}>"
 
     async def nmembers(self) -> int:
@@ -394,31 +394,31 @@ async def arrays(self) -> AsyncGenerator[AsyncArray, None]:
             if isinstance(value, AsyncArray):
                 yield value
 
-    async def tree(self, expand=False, level=None) -> Any:
+    async def tree(self, expand: bool = False, level: int | None = None) -> Any:
         raise NotImplementedError
 
-    async def empty(self, **kwargs) -> AsyncArray:
+    async def empty(self, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def zeros(self, **kwargs) -> AsyncArray:
+    async def zeros(self, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def ones(self, **kwargs) -> AsyncArray:
+    async def ones(self, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def full(self, **kwargs) -> AsyncArray:
+    async def full(self, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def empty_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
+    async def empty_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def zeros_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
+    async def zeros_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def ones_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
+    async def ones_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
-    async def full_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
+    async def full_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
         raise NotImplementedError
 
     async def move(self, source: str, dest: str) -> None:
@@ -462,16 +462,16 @@ def __getitem__(self, path: str) -> Array | Group:
         else:
             return Group(obj)
 
-    def __delitem__(self, key) -> None:
+    def __delitem__(self, key: str) -> None:
         self._sync(self._async_group.delitem(key))
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[str]:
         raise NotImplementedError
 
-    def __len__(self):
+    def __len__(self) -> int:
         raise NotImplementedError
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: Any) -> None:
         """__setitem__ is not supported in v3"""
         raise NotImplementedError
 
@@ -502,7 +502,7 @@ def attrs(self) -> Attributes:
     def info(self):
         return self._async_group.info
 
-    def update_attributes(self, new_attributes: dict[str, Any]):
+    def update_attributes(self, new_attributes: dict[str, Any]) -> "Group":
         self._sync(self._async_group.update_attributes(new_attributes))
         return self
 
@@ -521,7 +521,7 @@ def members(self) -> tuple[tuple[str, Array | Group], ...]:
         result = tuple(map(lambda kv: (kv[0], _parse_async_node(kv[1])), _members))
         return result
 
-    def __contains__(self, member) -> bool:
+    def __contains__(self, member: str) -> bool:
         return self._sync(self._async_group.contains(member))
 
     def group_keys(self) -> tuple[str, ...]:
@@ -537,37 +537,37 @@ def array_keys(self) -> tuple[str, ...]:
     def arrays(self) -> tuple[Array, ...]:
         return tuple(Array(obj) for obj in self._sync_iter(self._async_group.arrays()))
 
-    def tree(self, expand=False, level=None) -> Any:
+    def tree(self, expand: bool = False, level: int | None = None) -> Any:
         return self._sync(self._async_group.tree(expand=expand, level=level))
 
-    def create_group(self, name: str, **kwargs) -> Group:
+    def create_group(self, name: str, **kwargs: Any) -> Group:
         return Group(self._sync(self._async_group.create_group(name, **kwargs)))
 
-    def create_array(self, name: str, **kwargs) -> Array:
+    def create_array(self, name: str, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.create_array(name, **kwargs)))
 
-    def empty(self, **kwargs) -> Array:
+    def empty(self, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.empty(**kwargs)))
 
-    def zeros(self, **kwargs) -> Array:
+    def zeros(self, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.zeros(**kwargs)))
 
-    def ones(self, **kwargs) -> Array:
+    def ones(self, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.ones(**kwargs)))
 
-    def full(self, **kwargs) -> Array:
+    def full(self, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.full(**kwargs)))
 
-    def empty_like(self, prototype: AsyncArray, **kwargs) -> Array:
+    def empty_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.empty_like(prototype, **kwargs)))
 
-    def zeros_like(self, prototype: AsyncArray, **kwargs) -> Array:
+    def zeros_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.zeros_like(prototype, **kwargs)))
 
-    def ones_like(self, prototype: AsyncArray, **kwargs) -> Array:
+    def ones_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.ones_like(prototype, **kwargs)))
 
-    def full_like(self, prototype: AsyncArray, **kwargs) -> Array:
+    def full_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
         return Array(self._sync(self._async_group.full_like(prototype, **kwargs)))
 
     def move(self, source: str, dest: str) -> None: