4 files changed, 25 insertions, 39 deletions

diff --git a/cppgenerator.cpp b/cppgenerator.cpp
index 70e974899..a75542baa 100644
--- a/cppgenerator.cpp
+++ b/cppgenerator.cpp
@@ -533,7 +533,7 @@ void CppGenerator::writeVirtualMethodNative(QTextStream &s, const AbstractMetaFu
             argConversions << argConv;
         }
 
-        s << "Py_BuildValue(\"(" << getFormatUnitString(func) << ")\"," << endl;
+        s << "Py_BuildValue(\"(" << getFormatUnitString(func, false) << ")\"," << endl;
         s << argConversions.join(",\n") << endl;
         s << INDENT << "));" << endl;
     }
@@ -801,8 +801,6 @@ void CppGenerator::writeConstructorWrapper(QTextStream& s, const AbstractMetaFun
     // Python owns it and C++ wrapper is false.
     if (shouldGenerateCppWrapper(overloads.first()->ownerClass()))
         s << INDENT << "sbkSelf->containsCppWrapper = 1;" << endl;
-    if (needsReferenceCountControl(metaClass))
-        s << INDENT << "sbkSelf->referredObjects = new Shiboken::RefCountMap;" << endl;
     s << INDENT << "BindingManager::instance().registerWrapper(sbkSelf, cptr);" << endl;
 
     // Create metaObject and register signal/slot
@@ -2040,12 +2038,20 @@ void CppGenerator::writeMethodCall(QTextStream& s, const AbstractMetaFunction* f
                 continue;
             const AbstractMetaClass* wrappedClass = 0;
             QString pyArgName = argumentNameFromIndex(func, arg_mod.index, &wrappedClass);
-            if (!wrappedClass) {
+            if (pyArgName.isEmpty()) {
                 s << "#error Invalid reference count modification for argument " << arg_mod.index << endl << endl;
                 break;
             }
+
             s << INDENT << "Shiboken::keepReference(reinterpret_cast<SbkBaseWrapper*>(self), \"";
-            s << func->minimalSignature() << arg_mod.index << "\", " << pyArgName << ");" << endl;
+            QString varName = arg_mod.referenceCounts.first().varName;
+            if (varName.isEmpty())
+                varName = func->minimalSignature() + QString().number(arg_mod.index);
+
+            s << varName << "\", " << pyArgName << ");" << endl;
+
+            if (arg_mod.index == 0)
+                hasReturnPolicy = true;
         }
     }
     writeParentChildManagement(s, func, !hasReturnPolicy);
diff --git a/libshiboken/basewrapper.cpp b/libshiboken/basewrapper.cpp
index eaa6a2aae..21e429535 100644
--- a/libshiboken/basewrapper.cpp
+++ b/libshiboken/basewrapper.cpp
@@ -376,25 +376,35 @@ void deallocWrapperWithPrivateDtor(PyObject* self)
 
 void keepReference(SbkBaseWrapper* self, const char* key, PyObject* referredObject)
 {
+    bool isNone = (!referredObject || (referredObject == Py_None));
     if (!self->referredObjects)
-        return;
+        self->referredObjects = new Shiboken::RefCountMap;
+
     RefCountMap& refCountMap = *(self->referredObjects);
-    Py_INCREF(referredObject);
+    if (!isNone)
+        Py_INCREF(referredObject);
+
     RefCountMap::iterator iter = refCountMap.find(key);
-    if (iter != refCountMap.end())
+    if (iter != refCountMap.end()){
         Py_DECREF(iter->second);
-    refCountMap[key] = referredObject;
+        refCountMap.erase(iter);
+    }
+
+    if (!isNone)
+        refCountMap[key] = referredObject;
 }
 
 void clearReferences(SbkBaseWrapper* self)
 {
     if (!self->referredObjects)
         return;
+
     RefCountMap& refCountMap = *(self->referredObjects);
     RefCountMap::iterator iter;
     for (iter = refCountMap.begin(); iter != refCountMap.end(); ++iter)
         Py_DECREF(iter->second);
     delete self->referredObjects;
+    self->referredObjects = 0;
 }
 
 bool importModule(const char* moduleName, PyTypeObject*** cppApiPtr)
diff --git a/shibokengenerator.cpp b/shibokengenerator.cpp
index df6093195..4c310804d 100644
--- a/shibokengenerator.cpp
+++ b/shibokengenerator.cpp
@@ -1359,30 +1359,6 @@ bool ShibokenGenerator::hasMultipleInheritanceInAncestry(const AbstractMetaClass
     return hasMultipleInheritanceInAncestry(metaClass->baseClass());
 }
 
-bool ShibokenGenerator::needsReferenceCountControl(const AbstractMetaClass* metaClass)
-{
-    if (!metaClass->fields().isEmpty() || hasMethodsWithReferenceCountModifications(metaClass))
-        return true;
-    foreach (const AbstractMetaClass* parent, getBaseClasses(metaClass)) {
-        if (needsReferenceCountControl(parent))
-            return true;
-    }
-    return false;
-}
-
-bool ShibokenGenerator::hasMethodsWithReferenceCountModifications(const AbstractMetaClass* metaClass)
-{
-    foreach (const AbstractMetaFunction* func, metaClass->functions()) {
-        foreach (FunctionModification func_mod, func->modifications()) {
-            foreach (ArgumentModification arg_mod, func_mod.argument_mods) {
-                if (!arg_mod.referenceCounts.isEmpty())
-                    return true;
-            }
-        }
-    }
-    return false;
-}
-
 bool ShibokenGenerator::classNeedsGetattroFunction(const AbstractMetaClass* metaClass)
 {
     if (!metaClass)
diff --git a/shibokengenerator.h b/shibokengenerator.h
index 17002caac..9f2feda3e 100644
--- a/shibokengenerator.h
+++ b/shibokengenerator.h
@@ -178,12 +178,6 @@ public:
     /// Returns true if there are cases of multiple inheritance in any of its ancestors.
     bool hasMultipleInheritanceInAncestry(const AbstractMetaClass* metaClass);
 
-    /// Returns true if the class needs reference counting control.
-    bool needsReferenceCountControl(const AbstractMetaClass* metaClass);
-
-    /// Returns true if the class has any method that modifies the reference counting of any of its arguments.
-    bool hasMethodsWithReferenceCountModifications(const AbstractMetaClass* metaClass);
-
     /// Returns true if the class needs to have a getattro function.
     bool classNeedsGetattroFunction(const AbstractMetaClass* metaClass);