Skip to content

Commit

Permalink
Add some return types
Browse files Browse the repository at this point in the history
  • Loading branch information
@martaiborra
martaiborra committed Sep 13, 2024
1 parent 5428211 commit dc410f9
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 36 deletions.
1 change: 0 additions & 1 deletion src/blosc2/blosc2_ext.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2486,7 +2486,6 @@ cdef class NDArray:
new_shape_[i] = s
_check_rc(b2nd_resize(self.array, new_shape_, NULL),
"Error while resizing the array")
return self

def squeeze(self):
_check_rc(b2nd_squeeze(self.array), "Error while performing the squeeze")
Expand Down
78 changes: 43 additions & 35 deletions src/blosc2/ndarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -601,74 +601,74 @@ def all(ndarr: NDArray | NDField | blosc2.C2Array | blosc2.LazyExpr, axis: int |
class Operand:
"""Base class for all operands in expressions."""

def __neg__(self):
def __neg__(self) -> blosc2.LazyExpr:
return blosc2.LazyExpr(new_op=(0, "-", self))

def __and__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __and__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "&", value))

def __add__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __add__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "+", value))

def __iadd__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __iadd__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "+", value))

def __radd__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __radd__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(value, "+", self))

def __sub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __sub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "-", value))

def __isub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __isub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "-", value))

def __rsub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __rsub__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(value, "-", self))

def __mul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __mul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "*", value))

def __imul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __imul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "*", value))

def __rmul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __rmul__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(value, "*", self))

def __truediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __truediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "/", value))

def __itruediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __itruediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "/", value))

def __rtruediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __rtruediv__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(value, "/", self))

def __lt__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __lt__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "<", value))

def __le__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __le__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "<=", value))

def __gt__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __gt__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, ">", value))

def __ge__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __ge__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, ">=", value))

Expand All @@ -678,19 +678,19 @@ def __eq__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
return self is value
return blosc2.LazyExpr(new_op=(self, "==", value))

def __ne__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __ne__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "!=", value))

def __pow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __pow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "**", value))

def __ipow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __ipow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(self, "**", value))

def __rpow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /):
def __rpow__(self, value: int | float | NDArray | NDField | blosc2.C2Array, /) -> blosc2.LazyExpr:
_check_allowed_dtypes(value)
return blosc2.LazyExpr(new_op=(value, "**", self))

Expand All @@ -715,7 +715,7 @@ def var(self, axis=None, dtype=None, ddof=0, keepdims=False, **kwargs):
return expr.var(axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, **kwargs)

@is_documented_by(prod)
def prod(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs):
def prod(self, axis=None, dtype=None, keepdims=False, **kwargs):
expr = blosc2.LazyExpr(new_op=(self, None, None))
return expr.prod(axis=axis, dtype=dtype, keepdims=keepdims, **kwargs)

Expand Down Expand Up @@ -754,7 +754,7 @@ def __init__(self, **kwargs):
self._fields[field] = NDField(self, field)

@property
def fields(self):
def fields(self) -> dict:
"""
Dictionary with the fields of the structured array.
Expand All @@ -781,12 +781,12 @@ def fields(self):
return self._fields

@property
def keep_last_read(self):
def keep_last_read(self) -> bool:
"""Indicates whether the last read data should be kept in memory."""
return self._keep_last_read

@keep_last_read.setter
def keep_last_read(self, value: bool):
def keep_last_read(self, value: bool) -> None:
"""Set whether the last read data should be kept in memory.
This always clear the last read data (if any).
Expand Down Expand Up @@ -835,7 +835,7 @@ def info(self):
return InfoReporter(self)

@property
def info_items(self):
def info_items(self) -> list:
items = []
items += [("type", f"{self.__class__.__name__}")]
items += [("shape", self.shape)]
Expand All @@ -848,7 +848,7 @@ def info_items(self):
return items

@property
def schunk(self):
def schunk(self) -> blosc2.SChunk:
"""
The :ref:`SChunk <SChunk>` reference of the :ref:`NDArray`.
All the attributes from the :ref:`SChunk <SChunk>` can be accessed through
Expand All @@ -861,7 +861,7 @@ def schunk(self):
return self._schunk

@property
def blocksize(self):
def blocksize(self) -> int:
"""The block size (in bytes) for this container.
This is a shortcut to
Expand Down Expand Up @@ -1194,7 +1194,7 @@ def copy(self, dtype: np.dtype = None, **kwargs: dict) -> NDArray:

return super().copy(dtype, **kwargs)

def resize(self, newshape: tuple | list):
def resize(self, newshape: tuple | list) -> None:
"""Change the shape of the array by growing or shrinking one or more dimensions.
Parameters
Expand All @@ -1203,6 +1203,10 @@ def resize(self, newshape: tuple | list):
The new shape of the array. It should have the same dimensions
as :paramref:`self`.
Returns
-------
out: None
Notes
-----
The array values corresponding to the added positions are not initialized.
Expand All @@ -1219,12 +1223,12 @@ def resize(self, newshape: tuple | list):
>>> b = blosc2.asarray(a)
>>> newshape = [50, 10]
>>> # Extend first dimension, shrink second dimension
>>> _ = b.resize(newshape)
>>> b.resize(newshape)
>>> b.shape
(50, 10)
"""
blosc2_ext.check_access_mode(self.schunk.urlpath, self.schunk.mode)
return super().resize(newshape)
super().resize(newshape)

def slice(self, key: int | slice | Sequence[slice], **kwargs: dict) -> NDArray:
"""Get a (multidimensional) slice as a new :ref:`NDArray`.
Expand Down Expand Up @@ -1280,9 +1284,13 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: dict) -> NDArray:

return ndslice

def squeeze(self):
def squeeze(self) -> None:
"""Remove the 1's in array's shape.
Returns
-------
out: None
Examples
--------
>>> import blosc2
Expand Down Expand Up @@ -2643,12 +2651,12 @@ def __repr__(self):
return f"NDField({self.ndarr}, {self.field})"

@property
def shape(self):
def shape(self) -> tuple[int]:
"""The shape of the associated :ref:`NDArray`."""
return self.ndarr.shape

@property
def schunk(self):
def schunk(self) -> blosc2.SChunk:
"""The associated :ref:`SChunk <SChunk>`."""
return self.ndarr.schunk

Expand Down

0 comments on commit dc410f9

Please sign in to comment.