From dbeb43a53b9cce45bd39aa87b4f6d7b832e50132 Mon Sep 17 00:00:00 2001
From: Marta Iborra <martaiborra24@gmail.com>
Date: Fri, 13 Sep 2024 10:44:54 +0200
Subject: [PATCH] Add type annotations to lazyexpr.py

---
 src/blosc2/lazyexpr.py | 20 +++++++++++---------
 src/blosc2/ndarray.py  |  2 +-
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py
index 643aa15c..10af03d6 100644
--- a/src/blosc2/lazyexpr.py
+++ b/src/blosc2/lazyexpr.py
@@ -13,7 +13,7 @@
 import pathlib
 import threading
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from enum import Enum
 from pathlib import Path
 from queue import Empty, Queue
@@ -61,7 +61,7 @@ class LazyArrayEnum(Enum):
 
 class LazyArray(ABC):
     @abstractmethod
-    def eval(self, item, **kwargs):
+    def eval(self, item: slice | list[slice] = None, **kwargs: dict) -> blosc2.NDArray:
         """
         Return a :ref:`NDArray` containing the evaluation of the :ref:`LazyArray`.
 
@@ -90,7 +90,7 @@ def eval(self, item, **kwargs):
         pass
 
     @abstractmethod
-    def __getitem__(self, item):
+    def __getitem__(self, item: int | slice | Sequence[slice]) -> blosc2.NDArray:
         """
         Return a NumPy.ndarray containing the evaluation of the :ref:`LazyArray`.
 
@@ -107,7 +107,7 @@ def __getitem__(self, item):
         pass
 
     @abstractmethod
-    def save(self, **kwargs):
+    def save(self, **kwargs: dict) -> None:
         """
         Save the :ref:`LazyArray` on disk.
 
@@ -133,7 +133,7 @@ def save(self, **kwargs):
 
     @property
     @abstractmethod
-    def dtype(self):
+    def dtype(self) -> np.dtype:
         """
         Get the data type of the :ref:`LazyArray`.
 
@@ -146,7 +146,7 @@ def dtype(self):
 
     @property
     @abstractmethod
-    def shape(self):
+    def shape(self) -> tuple[int]:
         """
         Get the shape of the :ref:`LazyArray`.
 
@@ -159,7 +159,7 @@ def shape(self):
 
     @property
     @abstractmethod
-    def info(self):
+    def info(self) -> InfoReporter:
         """
         Get information about the :ref:`LazyArray`.
 
@@ -1873,7 +1873,8 @@ def _open_lazyarray(array):
     return expr
 
 
-def lazyudf(func, inputs, dtype, chunked_eval=True, **kwargs):
+def lazyudf(func: Callable[[tuple, np.ndarray, tuple[int]], None], inputs: tuple | list,
+            dtype: np.dtype, chunked_eval: bool = True, **kwargs: dict) -> LazyUDF:
     """
     Get a LazyUDF from a python user-defined function.
 
@@ -1908,7 +1909,8 @@ def lazyudf(func, inputs, dtype, chunked_eval=True, **kwargs):
     return LazyUDF(func, inputs, dtype, chunked_eval, **kwargs)
 
 
-def lazyexpr(expression, operands=None, out=None, where=None):
+def lazyexpr(expression: str | bytes | LazyExpr, operands: dict = None,
+             out: blosc2.NDArray | np.ndarray = None, where: tuple | list = None) -> LazyExpr:
     """
     Get a LazyExpr from an expression.
 
diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py
index fa0deac4..73246043 100644
--- a/src/blosc2/ndarray.py
+++ b/src/blosc2/ndarray.py
@@ -798,7 +798,7 @@ def keep_last_read(self, value: bool) -> None:
         self._keep_last_read = value
 
     @property
-    def info(self):
+    def info(self) -> InfoReporter:
         """
         Print information about this array.