X-Git-Url: https://git.cworth.org/git?a=blobdiff_plain;f=trace.py;h=368462a917579a90c16989fef5a84dff0c6aa159;hb=bb8760bc379ec3c37db0d18c9972b3c4c3df1179;hp=9f115a32e92687086280fee60d415419d44405a9;hpb=8bec8a0c309cffea31148ccbf0d5cf76e3e93619;p=apitrace

diff --git a/trace.py b/trace.py
index 9f115a3..368462a 100644
--- a/trace.py
+++ b/trace.py
@@ -271,7 +271,12 @@ unwrap_instance = Unwrapper().visit
 
 class Tracer:
 
+    def __init__(self):
+        self.api = None
+
     def trace_api(self, api):
+        self.api = api
+
         self.header(api)
 
         # Includes
@@ -427,8 +432,15 @@ class Tracer:
             wrap_instance(method.type, '__result')
         print '    Trace::EndLeave();'
         if method.name == 'QueryInterface':
-            print '    if (*ppvObj == m_pInstance)'
-            print '        *ppvObj = this;'
+            print '    if (ppvObj && *ppvObj) {'
+            print '        if (*ppvObj == m_pInstance) {'
+            print '            *ppvObj = this;'
+            print '        }'
+            for iface in self.api.interfaces:
+                print '        else if (riid == IID_%s) {' % iface.name
+                print '            *ppvObj = new Wrap%s((%s *) *ppvObj);' % (iface.name, iface.name)
+                print '        }'
+            print '    }'
         if method.name == 'Release':
             assert method.type is not stdapi.Void
             print '    if (!__result)'
@@ -444,9 +456,6 @@ class DllTracer(Tracer):
     def __init__(self, dllname):
         self.dllname = dllname
     
-    def get_function_address(self, function):
-        return '__%s' % (function.name,)
-
     def header(self, api):
         print '''
 static HINSTANCE g_hDll = NULL;