From 8fbc267158a823e54c9b0f54dd2c702fe9941644 Mon Sep 17 00:00:00 2001
From: jonschz <jonschz@users.noreply.github.com>
Date: Sat, 29 Jun 2024 08:04:18 +0200
Subject: [PATCH] Implement static import for `ISequentialStream` (#474)

---
 comtypes/client/_generate.py  |  1 +
 comtypes/objidl.py            | 59 +++++++++++++++++++++++++++++++++++
 comtypes/test/test_istream.py | 39 +++++++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 comtypes/objidl.py
 create mode 100644 comtypes/test/test_istream.py

diff --git a/comtypes/client/_generate.py b/comtypes/client/_generate.py
index f2e65941..f825a71b 100644
--- a/comtypes/client/_generate.py
+++ b/comtypes/client/_generate.py
@@ -272,6 +272,7 @@ def _get_known_namespaces() -> Tuple[
         "comtypes.persist",
         "comtypes.typeinfo",
         "comtypes.automation",
+        "comtypes.objidl",
         "comtypes",
         "ctypes.wintypes",
         "ctypes",
diff --git a/comtypes/objidl.py b/comtypes/objidl.py
new file mode 100644
index 00000000..bbad9737
--- /dev/null
+++ b/comtypes/objidl.py
@@ -0,0 +1,59 @@
+from ctypes import Array, c_ubyte, c_ulong, HRESULT, POINTER, pointer
+from typing import Tuple, TYPE_CHECKING
+
+from comtypes import COMMETHOD, GUID, IUnknown
+
+
+class ISequentialStream(IUnknown):
+    """Defines methods for the stream objects in sequence."""
+
+    _iid_ = GUID("{0C733A30-2A1C-11CE-ADE5-00AA0044773D}")
+    _idlflags_ = []
+
+    _methods_ = [
+        COMMETHOD(
+            [],
+            HRESULT,
+            "RemoteRead",
+            # This call only works if `pv` is pre-allocated with `cb` bytes,
+            # which cannot be done by the high level function generated by metaclasses.
+            # Therefore, we override the high level function to implement this behaviour
+            # and then delegate the call the raw COM method.
+            (["out"], POINTER(c_ubyte), "pv"),
+            (["in"], c_ulong, "cb"),
+            (["out"], POINTER(c_ulong), "pcbRead"),
+        ),
+        COMMETHOD(
+            [],
+            HRESULT,
+            "RemoteWrite",
+            (["in"], POINTER(c_ubyte), "pv"),
+            (["in"], c_ulong, "cb"),
+            (["out"], POINTER(c_ulong), "pcbWritten"),
+        ),
+    ]
+
+    def RemoteRead(self, cb: int) -> Tuple["Array[c_ubyte]", int]:
+        """Reads a specified number of bytes from the stream object into memory
+        starting at the current seek pointer.
+        """
+        # Behaves as if `pv` is pre-allocated with `cb` bytes by the high level func.
+        pv = (c_ubyte * cb)()
+        pcb_read = pointer(c_ulong(0))
+        self.__com_RemoteRead(pv, c_ulong(cb), pcb_read)  # type: ignore
+        return pv, pcb_read.contents.value
+
+    if TYPE_CHECKING:
+
+        def RemoteWrite(self, pv: "Array[c_ubyte]", cb: int) -> int:
+            """Writes a specified number of bytes into the stream object starting at
+            the current seek pointer.
+            """
+            ...
+
+
+# fmt: off
+__known_symbols__ = [
+    'ISequentialStream',
+]
+# fmt: on
diff --git a/comtypes/test/test_istream.py b/comtypes/test/test_istream.py
new file mode 100644
index 00000000..4ccebe43
--- /dev/null
+++ b/comtypes/test/test_istream.py
@@ -0,0 +1,39 @@
+import unittest as ut
+
+from ctypes import POINTER, byref, c_bool, c_ubyte
+import comtypes
+import comtypes.client
+
+comtypes.client.GetModule("portabledeviceapi.dll")
+from comtypes.gen.PortableDeviceApiLib import IStream
+
+
+class Test_IStream(ut.TestCase):
+    def test_istream(self):
+        # Create an IStream
+        stream: IStream = POINTER(IStream)()  # type: ignore
+        comtypes._ole32.CreateStreamOnHGlobal(None, c_bool(True), byref(stream))
+
+        test_data = "Some data".encode("utf-8")
+        pv = (c_ubyte * len(test_data)).from_buffer(bytearray(test_data))
+
+        written = stream.RemoteWrite(pv, len(test_data))
+        self.assertEqual(written, len(test_data))
+        # make sure the data actually gets written before trying to read back
+        stream.Commit(0)
+
+        # Move the stream back to the beginning
+        STREAM_SEEK_SET = 0
+        stream.RemoteSeek(0, STREAM_SEEK_SET)
+
+        read_buffer_size = 1024
+
+        read_buffer, data_read = stream.RemoteRead(read_buffer_size)
+
+        # Verification
+        self.assertEqual(data_read, len(test_data))
+        self.assertEqual(bytearray(read_buffer)[0:data_read], test_data)
+
+
+if __name__ == "__main__":
+    ut.main()