Skip to content

Commit c49f799

Browse files
committed
✨: add HasArrayRShift protocol to support right shift operator for array classes
1 parent cbe6dbb commit c49f799

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

src/array_api_typing/_array.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,22 @@ def __lshift__(self, other: Self | int, /) -> Self:
268268
...
269269

270270

271+
class HasArrayRShift(Protocol):
272+
"""Protocol for array classes that support the right shift operator."""
273+
274+
def __rshift__(self, other: Self | int, /) -> Self:
275+
"""Calculates the right shift for each element of an array instance with the respective element of the array other.
276+
277+
Args:
278+
other: operand array. Must be compatible with self (see Broadcasting). Should have an integer data type.
279+
280+
Returns:
281+
an array containing the element-wise right shifts. The returned array must have a data type determined by Type Promotion Rules.
282+
283+
""" # noqa: E501
284+
...
285+
286+
271287
class Array(
272288
HasArrayNamespace[NS_co],
273289
HasArrayPos,
@@ -285,6 +301,7 @@ class Array(
285301
HasArrayOr,
286302
HasArrayXOr,
287303
HasArrayLShift,
304+
HasArrayRShift,
288305
Protocol,
289306
):
290307
"""Array API specification for array object attributes and methods."""

0 commit comments

Comments
 (0)