Merge pull request #88 from hakril/COMImplementation_keepalive

COMImplementation class are now kept alive based on com refcount & gl…

Author: hakril

tests/test_com.py

@@ -0,0 +1,77 @@
+import ctypes
+import gc
+import weakref
+
+import pytest
+
+import windows
+import windows.com
+
+import windows.generated_def as gdef
+
+# ICallFrameEvents is a good candidate for testing :
+# - only 1 fonction (not counting IUnkown())
+# - Only 1 param for the function
+
+# class ICallFrameEventsImplem(windows.com.COMImplementation):
+#     IMPLEMENT = ICallFrameEvents
+#
+#     def OnCall(self, This, pFrame):
+#         print('ICallFrameEvents.OnCall')
+#         return E_NOTIMPL
+
+class ICallFrameEventsImplemIncomplete(windows.com.COMImplementation):
+    IMPLEMENT = gdef.ICallFrameEvents
+
+
+def test_com_implementation_incomplete():
+    with pytest.raises(ValueError):
+        obj = ICallFrameEventsImplemIncomplete()
+
+
+class ICallFrameEventsImplemSimple(windows.com.COMImplementation):
+    IMPLEMENT = gdef.ICallFrameEvents
+
+    def __init__(self):
+        super(ICallFrameEventsImplemSimple, self).__init__()
+        self.called = 0
+
+    def OnCall(self, This, pFrame):
+        self.called += 1
+        return self.called
+
+def test_com_implementation_simple():
+    obj = ICallFrameEventsImplemSimple()
+    assert obj.com_refcount == 1
+    assert obj.AddRef() == 2
+    assert obj.Release() == 1
+    assert id(obj) in obj._get_keepalive_registry()
+    assert obj.Release() == 0
+    assert id(obj) not in obj._get_keepalive_registry()
+
+def test_com_implementation_query_interface():
+    obj = ICallFrameEventsImplemSimple()
+    iunk = gdef.IUnknown()
+    # 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 iunk.Release() == 1 # We are now working via ctypes/vtable call
+
+def test_com_implementation_keep_alive():
+    obj = ICallFrameEventsImplemSimple()
+    # Create a C-like pointer from the interface
+    objcom = gdef.ICallFrameEvents(obj._as_parameter_)
+    assert objcom.AddRef() == 2
+    wref = weakref.ref(obj)
+    assert wref() is obj # Check object is still alive in python world
+    del obj # No more python-side direct reference
+    gc.collect() # Force gc : code below would crash without keep-alive logic
+    assert objcom.Release() == 1
+    assert objcom.OnCall(None) == 1 # Count the number of call : proof of correct python-side execution
+    assert objcom.OnCall(None) == 2
+    gc.collect()
+    assert objcom.OnCall(None) == 3
+    assert objcom.Release() == 0 # trigger Revoke / free !
+    gc.collect()
+    assert wref() is None # Check object is dead in python world

windows/com.py

@@ -349,6 +349,7 @@ def __repr__(self):
 class COMImplementation(object):
     """The base class to implements COM object respecting a given interface"""
     IMPLEMENT = None
+    COM_KEEPALIVE_REGISTRY = {} # Will keep on COM-alive object to prevent early garbage collection
 
     def get_index_of_method(self, method):
         # This code is horrible but not totally my fault
@@ -392,18 +393,47 @@ def __init__(self):
         self.implems = implems
         self.vtable_pointer = ctypes.pointer(self.vtable)
         self._as_parameter_ = ctypes.addressof(self.vtable_pointer)
+        self.com_refcount = 1
+        self.register()
+
+    # PFW COMImplementation API
+    def _get_keepalive_registry(self):
+        """Allow subclass to chose where are store alive COMImplementation object"""
+        return self.COM_KEEPALIVE_REGISTRY
+
+    def register(self):
+        """Called once at instanciation : should be used to register it in any pointer-keeper logic to prevent early garbage-collection
+        The method : revoke() will be called when RefCount go to 0.
+        """
+        self._get_keepalive_registry()[id(self)] = self # Add object to global table to prevent early garbage collection
+        return True
+
+    def revoke(self):
+        """Called once when Refcount is lowered to 0 : use it to remove from any global-state object for refcounting.
+        This indicate the object can be safely garbage collected from the COM point of view
+        """
+        del self._get_keepalive_registry()[id(self)]
+        return True
 
     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
+            self.AddRef()
             return 1
         return E_NOINTERFACE
 
     def AddRef(self, *args):
         """Default ``AddRef`` implementation that returns ``1``"""
-        return 1
+        self.com_refcount += 1
+        return self.com_refcount
 
     def Release(self, *args):
         """Default ``Release`` implementation that returns ``1``"""
-        return 0
+        self.com_refcount -= 1
+        if self.com_refcount == 0:
+            self.revoke()
+        return self.com_refcount
+
+    def __repr__(self):
+        return "<{0} refcount={1} at {2:#08x}>".format(type(self).__name__, self.com_refcount, id(self))