ENH: test bitwise_{left,right}_shift with scalars

Author: ev-br

array_api_tests/hypothesis_helpers.py

@@ -457,7 +457,11 @@ def scalars(draw, dtypes, finite=False, **kwds):
     """
     dtype = draw(dtypes)
     if dh.is_int_dtype(dtype):
-        m, M = dh.dtype_ranges[dtype]
+        mM = kwds.pop('mM', None)
+        if mM is None:
+            m, M = dh.dtype_ranges[dtype]
+        else:
+            m, M = mM
         return draw(integers(m, M))
     elif dtype == bool_dtype:
         return draw(booleans())
@@ -588,18 +592,20 @@ def two_mutual_arrays(
 
 
 @composite
-def array_and_py_scalar(draw, dtypes):
+def array_and_py_scalar(draw, dtypes, mM=None, positive=False):
     """Draw a pair: (array, scalar) or (scalar, array)."""
     dtype = draw(sampled_from(dtypes))
 
-    scalar_var = draw(scalars(just(dtype), finite=True,
-        **{'min_value': 1/ (2<<5), 'max_value': 2<<5}
-    ))
+    scalar_var = draw(scalars(just(dtype), finite=True, mM=mM))
+    if positive:
+        assume (scalar_var > 0)
 
     elements={}
     if dtype in dh.real_float_dtypes:
         elements = {'allow_nan': False, 'allow_infinity': False,
                     'min_value': 1.0 / (2<<5), 'max_value': 2<<5}
+    if positive:
+        elements = {'min_value': 0}
     array_var = draw(arrays(dtype, shape=shapes(min_dims=1), elements=elements))
 
     if draw(booleans()):

array_api_tests/test_operators_and_elementwise_functions.py

@@ -1881,3 +1881,20 @@ def test_binary_with_scalars_bitwise(func_data, x1x2):
     refimpl_ = lambda l, r: mock_int_dtype(refimpl(l, r), xp.int32 )
     _check_binary_with_scalars((func_name, refimpl_, kwargs, expected), x1x2)
 
+
+@pytest.mark.min_version("2024.12")
+@pytest.mark.parametrize('func_data',
+    # func_name, refimpl, kwargs, expected_dtype
+    [
+        ("bitwise_left_shift", operator.lshift, {}, None),
+        ("bitwise_right_shift", operator.rshift, {}, None),
+    ],
+    ids=lambda func_data: func_data[0]  # use names for test IDs
+)
+@given(x1x2=hh.array_and_py_scalar([xp.int32], positive=True, mM=(1, 3)))
+def test_binary_with_scalars_bitwise_shifts(func_data, x1x2):
+    func_name, refimpl, kwargs, expected = func_data
+    # repack the refimpl
+    refimpl_ = lambda l, r: mock_int_dtype(refimpl(l, r), xp.int32 )
+    _check_binary_with_scalars((func_name, refimpl_, kwargs, expected), x1x2)
+