add test for python scalar second argument for searchsorted

Author: ndgrigorian

dpctl/tests/test_usm_ndarray_searchsorted.py

@@ -1,3 +1,21 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2025 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ctypes
+
 import numpy as np
 import pytest
 
@@ -7,6 +25,30 @@
 
 from .helper import get_queue_or_skip, skip_if_dtype_not_supported
 
+_integer_dtypes = [
+    "i1",
+    "u1",
+    "i2",
+    "u2",
+    "i4",
+    "u4",
+    "i8",
+    "u8",
+]
+
+_floating_dtypes = [
+    "f2",
+    "f4",
+    "f8",
+]
+
+_complex_dtypes = [
+    "c8",
+    "c16",
+]
+
+_all_dtypes = ["?"] + _integer_dtypes + _floating_dtypes + _complex_dtypes
+
 
 def _check(hay_stack, needles, needles_np):
     assert hay_stack.dtype == needles.dtype
@@ -73,19 +115,7 @@ def test_searchsorted_strided_bool():
     )
 
 
-@pytest.mark.parametrize(
-    "idt",
-    [
-        dpt.int8,
-        dpt.uint8,
-        dpt.int16,
-        dpt.uint16,
-        dpt.int32,
-        dpt.uint32,
-        dpt.int64,
-        dpt.uint64,
-    ],
-)
+@pytest.mark.parametrize("idt", _integer_dtypes)
 def test_searchsorted_contig_int(idt):
     q = get_queue_or_skip()
     skip_if_dtype_not_supported(idt, q)
@@ -105,19 +135,7 @@ def test_searchsorted_contig_int(idt):
     )
 
 
-@pytest.mark.parametrize(
-    "idt",
-    [
-        dpt.int8,
-        dpt.uint8,
-        dpt.int16,
-        dpt.uint16,
-        dpt.int32,
-        dpt.uint32,
-        dpt.int64,
-        dpt.uint64,
-    ],
-)
+@pytest.mark.parametrize("idt", _integer_dtypes)
 def test_searchsorted_strided_int(idt):
     q = get_queue_or_skip()
     skip_if_dtype_not_supported(idt, q)
@@ -144,12 +162,12 @@ def _add_extended_fp(array):
     array[-1] = dpt.nan
 
 
-@pytest.mark.parametrize("idt", [dpt.float16, dpt.float32, dpt.float64])
-def test_searchsorted_contig_fp(idt):
+@pytest.mark.parametrize("fdt", _floating_dtypes)
+def test_searchsorted_contig_fp(fdt):
     q = get_queue_or_skip()
-    skip_if_dtype_not_supported(idt, q)
+    skip_if_dtype_not_supported(fdt, q)
 
-    dt = dpt.dtype(idt)
+    dt = dpt.dtype(fdt)
 
     hay_stack = dpt.linspace(0, 1, num=255, dtype=dt, endpoint=True)
     _add_extended_fp(hay_stack)
@@ -165,12 +183,12 @@ def test_searchsorted_contig_fp(idt):
     )
 
 
-@pytest.mark.parametrize("idt", [dpt.float16, dpt.float32, dpt.float64])
-def test_searchsorted_strided_fp(idt):
+@pytest.mark.parametrize("fdt", _floating_dtypes)
+def test_searchsorted_strided_fp(fdt):
     q = get_queue_or_skip()
-    skip_if_dtype_not_supported(idt, q)
+    skip_if_dtype_not_supported(fdt, q)
 
-    dt = dpt.dtype(idt)
+    dt = dpt.dtype(fdt)
 
     hay_stack = dpt.repeat(
         dpt.linspace(0, 1, num=255, dtype=dt, endpoint=True), 4
@@ -213,12 +231,12 @@ def _add_extended_cfp(array):
     return dpt.sort(dpt.concat((ev, array)))
 
 
-@pytest.mark.parametrize("idt", [dpt.complex64, dpt.complex128])
-def test_searchsorted_contig_cfp(idt):
+@pytest.mark.parametrize("cdt", _complex_dtypes)
+def test_searchsorted_contig_cfp(cdt):
     q = get_queue_or_skip()
-    skip_if_dtype_not_supported(idt, q)
+    skip_if_dtype_not_supported(cdt, q)
 
-    dt = dpt.dtype(idt)
+    dt = dpt.dtype(cdt)
 
     hay_stack = dpt.linspace(0, 1, num=255, dtype=dt, endpoint=True)
     hay_stack = _add_extended_cfp(hay_stack)
@@ -233,12 +251,12 @@ def test_searchsorted_contig_cfp(idt):
     )
 
 
-@pytest.mark.parametrize("idt", [dpt.complex64, dpt.complex128])
-def test_searchsorted_strided_cfp(idt):
+@pytest.mark.parametrize("cdt", _complex_dtypes)
+def test_searchsorted_strided_cfp(cdt):
     q = get_queue_or_skip()
-    skip_if_dtype_not_supported(idt, q)
+    skip_if_dtype_not_supported(cdt, q)
 
-    dt = dpt.dtype(idt)
+    dt = dpt.dtype(cdt)
 
     hay_stack = dpt.repeat(
         dpt.linspace(0, 1, num=255, dtype=dt, endpoint=True), 4
@@ -375,3 +393,23 @@ def test_searchsorted_strided_scalar_needle():
     needles = dpt.asarray(needles_np)
 
     _check(hay_stack, needles, needles_np)
+
+
+@pytest.mark.parametrize("dt", _all_dtypes)
+def test_searchsorted_py_scalars(dt):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+
+    x = dpt.zeros(10, dtype=dt, sycl_queue=q)
+    py_zeros = (
+        bool(0),
+        int(0),
+        float(0),
+        complex(0),
+        np.float32(0),
+        ctypes.c_int(0),
+    )
+    for sc in py_zeros:
+        r1 = dpt.searchsorted(x, sc)
+        assert isinstance(r1, dpt.usm_ndarray)
+        assert r1.shape == ()