Merge pull request #89 from hakril/comimplem_multiinterface

COMImplementation can now implement multiple interface at once

Author: hakril

tests/test_com.py

@@ -34,15 +34,15 @@ class ICallFrameEventsImplemSimple(windows.com.COMImplementation):
 
     def __init__(self):
         super(ICallFrameEventsImplemSimple, self).__init__()
-        self.called = 0
+        self.oncall = 0
 
     def OnCall(self, This, pFrame):
-        self.called += 1
-        return self.called
+        self.oncall += 1
+        return self.oncall
 
 def test_com_implementation_simple():
     obj = ICallFrameEventsImplemSimple()
-    assert obj.com_refcount == 1
+    assert obj._com_refcount == 1
     assert obj.AddRef() == 2
     assert obj.Release() == 1
     assert id(obj) in obj._get_keepalive_registry()
@@ -55,7 +55,7 @@ def test_com_implementation_query_interface():
     # Emulate a call from C code
     obj.QueryInterface(obj._as_parameter_, ctypes.pointer(gdef.IUnknown.IID), ctypes.pointer(iunk))
     assert obj._as_parameter_ == iunk.value
-    assert obj.com_refcount == 2 # +1 on QueryInterface
+    assert obj._com_refcount == 2 # +1 on QueryInterface
     assert iunk.Release() == 1 # We are now working via ctypes/vtable call
 
 def test_com_implementation_keep_alive():
@@ -75,3 +75,129 @@ def test_com_implementation_keep_alive():
     assert objcom.Release() == 0 # trigger Revoke / free !
     gc.collect()
     assert wref() is None # Check object is dead in python world
+
+
+# Test multiple interface code
+
+class SimpleMultipleInterfaceImplem(windows.com.COMImplementation):
+    IMPLEMENT = [
+        gdef.ICallFrameEvents, # OnCall
+        gdef.IPersist # GetClassID
+    ]
+
+    def __init__(self):
+        super(SimpleMultipleInterfaceImplem, self).__init__()
+        self.oncall = 0
+        self.getclassid = 0
+        self.global_call = 0 # Proof this is the same object
+
+    def OnCall(self, This, pFrame):
+        self.oncall += 1
+        self.global_call += 1
+        return self.oncall
+
+    def GetClassID(self, This, pClassID):
+        self.getclassid += 1
+        self.global_call += 1
+        return self.getclassid
+
+def test_com_multiple_implementation_no_collision():
+    obj = SimpleMultipleInterfaceImplem()
+    with pytest.raises(ValueError):
+        x = obj._as_parameter_ # _as_parameter_ cannot be ask directly on object implementing multiple interfaces
+
+    # Create a C-like pointer from the interface
+    objcom = gdef.ICallFrameEvents(obj.as_interface(gdef.ICallFrameEvents))
+    assert objcom.OnCall(None) == 1 # Count the number of call : proof of correct python-side execution
+    assert objcom.OnCall(None) == 2
+    assert obj.global_call == 2
+
+    objcom2 = gdef.IPersist(obj.as_interface(gdef.IPersist))
+    assert objcom2.GetClassID(None) == 1 # Count the number of call : proof of correct python-side execution
+    assert objcom2.GetClassID(None) == 2
+    assert obj.global_call == 4
+
+
+
+def test_com_multiple_implementation_query_interface():
+    obj = SimpleMultipleInterfaceImplem()
+    iunk = gdef.IUnknown()
+    # Emulate a call from C code
+    obj.QueryInterface(obj.as_interface(gdef.ICallFrameEvents), ctypes.pointer(gdef.IUnknown.IID), ctypes.pointer(iunk))
+    assert obj._com_refcount == 2
+    assert iunk.value
+    assert iunk.value in obj._vtables_ptr.values()
+    # Use the standard C/ctypes Iunknown pointer to query other types and verify pointer/refcount
+    assert iunk.query(gdef.ICallFrameEvents).value == obj._vtables_ptr[str(gdef.ICallFrameEvents.IID)]
+    assert obj._com_refcount == 3
+    assert iunk.query(gdef.IPersist).value == obj._vtables_ptr[str(gdef.IPersist.IID)]
+    assert obj._com_refcount == 4
+
+
+class CollisionMultipleInterfaceImplem(windows.com.COMImplementation):
+    # Both share the same root & have "Invoke"
+    IMPLEMENT = [
+        gdef.IIdleSettings, # Invoke
+        gdef.ITaskFolderCollection, # Invoke
+    ]
+
+    def __init__(self):
+        super(CollisionMultipleInterfaceImplem, self).__init__()
+        self.iidlesettings_invoke = 0
+        self.itaskfoldercollection_invoke = 0
+
+    def IIdleSettings_Invoke(self, This, dispIdMember, riid, lcid, wFlags, pDispParams, pVarResult, pExcepInfo, puArgErr):
+        print('IIdleSettings.Invoke')
+        self.iidlesettings_invoke += 1
+        return self.iidlesettings_invoke
+
+    def ITaskFolderCollection_Invoke(self, This, dispIdMember, riid, lcid, wFlags, pDispParams, pVarResult, pExcepInfo, puArgErr):
+        print('ITaskFolderCollection.Invoke')
+        self.itaskfoldercollection_invoke += 1
+        return self.itaskfoldercollection_invoke
+
+    def not_implemented(self, *args):
+        return gdef.E_NOTIMPL
+
+    # IIdleSettings
+
+    IIdleSettings_GetTypeInfoCount = not_implemented
+    IIdleSettings_GetTypeInfo = not_implemented
+    IIdleSettings_GetIDsOfNames = not_implemented
+    get_IdleDuration = not_implemented
+    put_IdleDuration = not_implemented
+    get_WaitTimeout = not_implemented
+    put_WaitTimeout = not_implemented
+    get_StopOnIdleEnd = not_implemented
+    put_StopOnIdleEnd = not_implemented
+    get_RestartOnIdle = not_implemented
+    put_RestartOnIdle = not_implemented
+
+    # ITaskFolderCollection
+
+    ITaskFolderCollection_GetTypeInfoCount = not_implemented
+    ITaskFolderCollection_GetTypeInfo = not_implemented
+    ITaskFolderCollection_GetIDsOfNames = not_implemented
+    get_Count = not_implemented
+    get_Item = not_implemented
+    get__NewEnum = not_implemented
+
+
+def test_com_multiple_implementation_with_collision():
+    obj = CollisionMultipleInterfaceImplem()
+    assert obj
+
+    # Emulate a call from C code
+    iunk = gdef.IUnknown(obj.as_interface(gdef.IUnknown))
+
+    idlesetting = iunk.query(gdef.IIdleSettings)
+    taskfoldercollection = iunk.query(gdef.ITaskFolderCollection)
+    taskfoldercollection.Invoke(0, None, 0, 0, None, None, None, None)
+    assert obj.itaskfoldercollection_invoke == 1
+    assert obj.iidlesettings_invoke == 0
+    taskfoldercollection.Invoke(0, None, 0, 0, None, None, None, None)
+    assert obj.itaskfoldercollection_invoke == 2
+    assert obj.iidlesettings_invoke == 0
+    idlesetting.Invoke(0, None, 0, 0, None, None, None, None)
+    assert obj.itaskfoldercollection_invoke == 2
+    assert obj.iidlesettings_invoke == 1

windows/com.py

@@ -3,6 +3,7 @@
 import ctypes
 import functools
 from ctypes import HRESULT, byref, cast
+from collections import defaultdict
 
 import windows
 from windows import winproxy
@@ -364,36 +365,98 @@ def extract_methods_order(self, interface):
         index_and_method = sorted((self.get_index_of_method(m),name, m) for name, m in interface._functions_.items())
         return index_and_method
 
-    def verify_implem(self, interface):
-        for func_name in interface._functions_:
-            implem = getattr(self, func_name, None)
-            if implem is None:
-                raise ValueError("<{0}> implementing <{1}> has no method <{2}>".format(type(self).__name__, self.IMPLEMENT.__name__, func_name))
-            if not callable(implem):
-                raise ValueError("{0} implementing <{1}>: <{2}> is not callable".format(type(self).__name__, self.IMPLEMENT.__name__, func_name))
+
+    def _get_implemented_interfaces(self):
+        implist = self.IMPLEMENT
+        if isinstance(implist, type) and issubclass(implist, windows.generated_def.interfaces.COMInterface):
+            # Only one class :
+            implist = (implist,)
+        return tuple(implist)
+
+    def _verify_interfaces_collision(self):
+        function_names = defaultdict(int)
+        for interface in self._get_implemented_interfaces():
+            for function in interface._functions_:
+                function_names[function] += 1
+        return set(function for function, count in function_names.items() if count > 1) - {"QueryInterface", "AddRef", "Release"}
+
+    def _verify_interfaces_implem(self):
+        for interface in self._get_implemented_interfaces():
+            for func_name in interface._functions_:
+                explicit_func_name = "{}_{}".format(interface.__name__, func_name)
+                explicit_implem = getattr(self, explicit_func_name, None)
+                if explicit_implem is not None:
+                    if not callable(explicit_implem):
+                        raise ValueError("{0} implementing <{1}>: <{2}> is not callable".format(type(self).__name__, interface.__name__, explicit_func_name))
+                    continue # Ok for explicit implem
+
+                # No explicit implem : go for implicit
+                implem = getattr(self, func_name, None)
+                if implem is None:
+                    raise ValueError("<{0}> implementing <{1}> has no method <{2}>".format(type(self).__name__, interface.__name__, func_name))
+                if not callable(implem):
+                    raise ValueError("{0} implementing <{1}>: <{2}> is not callable".format(type(self).__name__, interface.__name__, func_name))
+
+                # Prevent multiple usage of the same name if multiple implemented interface have a function with the same name: force explicit name
+                if func_name in self._colliding_functions:
+                    raise ValueError("{0} method <{1}> is used by more than 1 implemented interface : use explicit method name 'Interface_Method' (Ex: IPersist_GetClassID) ".format(type(self).__name__, func_name))
+
         return True
 
     def _create_vtable(self, interface):
         implems = []
         names = []
         for index, name, method in self.extract_methods_order(interface):
-            func_implem = getattr(self, name)
+            explicit_func_name = "{}_{}".format(interface.__name__, name)
+            # Explicit name (IPersist_GetClassID)
+            func_implem = getattr(self, explicit_func_name, None)
+            if func_implem is None:
+                # Implicit name (GetClassID) # Collision have been checked in verify_interfaces_implem()
+                func_implem = getattr(self, name, None)
+
             #'this' is a COM-interface of the type we are implementing
             types = [method.restype, interface] + list(method.argtypes)
             implems.append(ctypes.WINFUNCTYPE(*types)(func_implem))
             names.append(name)
         class Vtable(ctypes.Structure):
             _fields_ = [(name, ctypes.c_void_p) for name in names]
-        return Vtable(*[ctypes.cast(x, ctypes.c_void_p) for x in implems]), implems
+        vtable = Vtable(*[ctypes.cast(x, ctypes.c_void_p) for x in implems])
+        vtable._implems = implems # Keep a reference to prevent early garbage-collection
+        return ctypes.pointer(vtable)
+
+    def _create_vtables(self):
+        vtables = {}
+        for interface in self._get_implemented_interfaces():
+            vtables[str(interface.IID).upper()] = self._create_vtable(interface)
+        return vtables
+
+    @property
+    def _as_parameter_(self):
+        implist = self._get_implemented_interfaces()
+        if len(implist) > 1:
+            raise ValueError("Cannot use _as_parameter_ on {} implementing multiple COM interface : use .as_interface() for ctypes parameter passing".format(self))
+        return self._vtables_ptr[str(implist[0].IID)]
+
+    def as_interface(self, interface):
+        try:
+            iid = interface.IID
+        except AttributeError:
+            iid = interface
+
+        if iid == gdef.IUnknown.IID: # Any interface pointer will respect IUnknown
+            return list(self._vtables_ptr.values())[0]
+
+        try:
+            return self._vtables_ptr[str(iid)]
+        except KeyError as e:
+            raise ValueError("{} does not implement IID {} (interface={})".format(self, iid, interface))
 
     def __init__(self):
-        self.verify_implem(self.IMPLEMENT)
-        vtable, implems = self._create_vtable(self.IMPLEMENT)
-        self.vtable = vtable
-        self.implems = implems
-        self.vtable_pointer = ctypes.pointer(self.vtable)
-        self._as_parameter_ = ctypes.addressof(self.vtable_pointer)
-        self.com_refcount = 1
+        self._colliding_functions = self._verify_interfaces_collision()
+        self._verify_interfaces_implem()
+        self._vtables = self._create_vtables()
+        self._vtables_ptr = {interface: ctypes.addressof(vtable) for interface, vtable in self._vtables.items()}
+        self._com_refcount = 1
         self.register()
 
     # PFW COMImplementation API
@@ -417,23 +480,23 @@ def revoke(self):
 
     def QueryInterface(self, this, piid, result):
         """Default ``QueryInterface`` implementation that returns ``self`` if piid is the implemented interface"""
-        if piid[0] in (gdef.IUnknown.IID, self.IMPLEMENT.IID):
-            result[0] = this
+        if piid[0] == (gdef.IUnknown.IID) or piid[0] in (i.IID for i in self._get_implemented_interfaces()):
+            result[0] = self.as_interface(piid[0])
             self.AddRef()
             return 1
         return E_NOINTERFACE
 
     def AddRef(self, *args):
         """Default ``AddRef`` implementation that returns ``1``"""
-        self.com_refcount += 1
-        return self.com_refcount
+        self._com_refcount += 1
+        return self._com_refcount
 
     def Release(self, *args):
         """Default ``Release`` implementation that returns ``1``"""
-        self.com_refcount -= 1
-        if self.com_refcount == 0:
+        self._com_refcount -= 1
+        if self._com_refcount == 0:
             self.revoke()
-        return self.com_refcount
+        return self._com_refcount
 
     def __repr__(self):
-        return "<{0} refcount={1} at {2:#08x}>".format(type(self).__name__, self.com_refcount, id(self))
+        return "<{0} refcount={1} at {2:#08x}>".format(type(self).__name__, self._com_refcount, id(self))