From 35cba428a3daabf2163d534f20a27a1b7c1da1e7 Mon Sep 17 00:00:00 2001
From: Marta Iborra <martaiborra24@gmail.com>
Date: Fri, 13 Sep 2024 11:04:01 +0200
Subject: [PATCH] Add type annotations in proxy.py

---
 src/blosc2/proxy.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/blosc2/proxy.py b/src/blosc2/proxy.py
index 66d35d4b..2ad263d2 100644
--- a/src/blosc2/proxy.py
+++ b/src/blosc2/proxy.py
@@ -8,6 +8,7 @@
 from abc import ABC, abstractmethod
 
 import blosc2
+import numpy as np
 
 
 class ProxySource(ABC):
@@ -24,7 +25,7 @@ class ProxySource(ABC):
     """
 
     @abstractmethod
-    def get_chunk(self, nchunk):
+    def get_chunk(self, nchunk: int) -> bytes:
         """
         Return the compressed chunk in :paramref:`self`.
 
@@ -48,7 +49,7 @@ class Proxy(blosc2.Operand):
     which follows the :ref:`ProxySource` interface in an urlpath.
     """
 
-    def __init__(self, src, urlpath=None, **kwargs):
+    def __init__(self, src: ProxySource, urlpath: str = None, **kwargs: dict):
         """
         Create a new :ref:`Proxy` to serve like a cache to save accessed
         chunks locally.
@@ -56,7 +57,7 @@ def __init__(self, src, urlpath=None, **kwargs):
         Parameters
         ----------
         src: :ref:`ProxySource`
-            The original container
+            The original container.
         urlpath: str, optional
             The urlpath where to save the container that will work as a cache.
 
@@ -115,7 +116,7 @@ def __init__(self, src, urlpath=None, **kwargs):
             for key in vlmeta:
                 self._schunk_cache.vlmeta[key] = vlmeta[key]
 
-    def fetch(self, item=None):
+    def fetch(self, item: slice | list[slice] = None) -> blosc2.NDArray | blosc2.schunk.SChunk:
         """
         Get the container used as cache with the requested data updated.
 
@@ -146,7 +147,7 @@ def fetch(self, item=None):
 
         return self._cache
 
-    async def afetch(self, item=None):
+    async def afetch(self, item: slice | list[slice] = None) -> blosc2.NDArray | blosc2.schunk.SChunk:
         """
         Get the container used as cache with the requested data updated
         in an asynchronous way.
@@ -185,7 +186,7 @@ async def afetch(self, item=None):
 
         return self._cache
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: slice | list[slice]) -> np.ndarray:
         """
         Get a slice as a numpy.ndarray using the :ref:`Proxy`.
 
@@ -204,12 +205,12 @@ def __getitem__(self, item):
         return self._cache[item]
 
     @property
-    def dtype(self):
+    def dtype(self) -> np.dtype:
         """The dtype of :paramref:`self` or None if the data is unidimensional"""
         return self._cache.dtype if isinstance(self._cache, blosc2.NDArray) else None
 
     @property
-    def shape(self):
+    def shape(self) -> tuple[int]:
         """The shape of :paramref:`self`"""
         return self._cache.shape if isinstance(self._cache, blosc2.NDArray) else len(self._cache)
 
@@ -217,7 +218,7 @@ def __str__(self):
         return f"Proxy({self.src}, urlpath={self.urlpath})"
 
     @property
-    def vlmeta(self):
+    def vlmeta(self) -> blosc2.schunk.vlmeta:
         """
         Get the vlmeta of the cache.
 
@@ -228,7 +229,7 @@ def vlmeta(self):
         return self._schunk_cache.vlmeta
 
     @property
-    def fields(self):
+    def fields(self) -> dict:
         """
         Dictionary with the fields of :paramref:`self`.
 
@@ -254,7 +255,7 @@ def __init__(self, proxy: Proxy, field: str):
         self.shape = proxy.shape
         self.dtype = proxy.dtype
 
-    def __getitem__(self, item: slice):
+    def __getitem__(self, item: slice | list[slice]) -> np.ndarray:
         """
         Get a slice as a numpy.ndarray using the `field` in `proxy`.