Add tests for importing generic legacy and generic modern containers

Author: oleksandr-pavlyk

dpctl/tests/test_usm_ndarray_dlpack.py

@@ -733,3 +733,76 @@ def test_copy_via_host_gh_1789():
         dpt.from_dlpack(x_np)
     with pytest.raises(BufferError):
         dpt.from_dlpack(x_np, device=(14, 0))
+
+
+class LegacyContainer:
+    "Helper class implementing legacy `__dlpack__` protocol"
+
+    def __init__(self, array):
+        self._array = array
+
+    def __dlpack__(self, stream=None):
+        return self._array.__dlpack__(stream=stream)
+
+    def __dlpack_device__(self):
+        return self._array.__dlpack_device__()
+
+
+class Container:
+    "Helper class implementing legacy `__dlpack__` protocol"
+
+    def __init__(self, array):
+        self._array = array
+
+    def __dlpack__(
+        self, max_version=None, dl_device=None, copy=None, stream=None
+    ):
+        return self._array.__dlpack__(
+            max_version=max_version,
+            dl_device=dl_device,
+            copy=copy,
+            stream=stream,
+        )
+
+    def __dlpack_device__(self):
+        return self._array.__dlpack_device__()
+
+
+def test_generic_container_legacy():
+    get_queue_or_skip()
+    C = LegacyContainer(dpt.linspace(0, 100, num=20, dtype="int16"))
+
+    X = dpt.from_dlpack(C)
+    assert isinstance(X, dpt.usm_ndarray)
+    assert X._pointer == C._array._pointer
+    assert X.sycl_device == C._array.sycl_device
+    assert X.dtype == C._array.dtype
+
+    Y = dpt.from_dlpack(C, device=(dpt.DLDeviceType.kDLCPU, 0))
+    assert isinstance(Y, np.ndarray)
+    assert Y.dtype == X.dtype
+
+    Z = dpt.from_dlpack(C, device=X.device)
+    assert isinstance(Z, dpt.usm_ndarray)
+    assert Z._pointer == X._pointer
+    assert Z.device == X.device
+
+
+def test_generic_container():
+    get_queue_or_skip()
+    C = Container(dpt.linspace(0, 100, num=20, dtype="int16"))
+
+    X = dpt.from_dlpack(C)
+    assert isinstance(X, dpt.usm_ndarray)
+    assert X._pointer == C._array._pointer
+    assert X.sycl_device == C._array.sycl_device
+    assert X.dtype == C._array.dtype
+
+    Y = dpt.from_dlpack(C, device=(dpt.DLDeviceType.kDLCPU, 0))
+    assert isinstance(Y, np.ndarray)
+    assert Y.dtype == X.dtype
+
+    Z = dpt.from_dlpack(C, device=X.device)
+    assert isinstance(Z, dpt.usm_ndarray)
+    assert Z._pointer == X._pointer
+    assert Z.device == X.device