Fix ref count error issue #906

traits/ctraits.c

@@ -2921,7 +2921,7 @@ trait_new(PyTypeObject *trait_type, PyObject *args, PyObject *kw)
     }
 
     PyErr_Format(
-        TraitError, 
+        TraitError,
         "Invalid argument to trait constructor. The argument `kind` "
         "must be an integer between 0 and 8 but a value of %d was provided.",
         kind);
@@ -3674,37 +3674,45 @@ validate_trait_function(
 |  Verifies a Python value is a callable (or None):
 +----------------------------------------------------------------------------*/
 
+/* Internal function for validation:
+
+   Return 1 if the value is valid, 0 if it is invalid, and return -1
+   and set an error condition if anything goes wrong.
+*/
+
+static int
+_validate_trait_callable(PyObject *type_info, PyObject *value)
+{
+    if (value == Py_None) {
+        if (PyTuple_GET_SIZE(type_info) < 2) {
+            /* Backwards compatibility with old Callable */
+            return 1;
+        }
+        else {
+            /* 2nd element of tuple determines whether None is allowed */
+            return PyObject_IsTrue(PyTuple_GET_ITEM(type_info, 1));
+        }
+    }
+    else {
+        return PyCallable_Check(value);
+    }
+}
+
+
 static PyObject *
 validate_trait_callable(
     trait_object *trait, has_traits_object *obj, PyObject *name,
     PyObject *value)
 {
-    if (PyCallable_Check(value)) {
+    int valid = _validate_trait_callable(trait->py_validate, value);
+
+    if (valid == -1) {
+        return NULL;
+    }
+    if (valid == 1) {
         Py_INCREF(value);
         return value;
     }
-    else if (value == Py_None) {
-        int allow_none;
-        int tuple_size = PyTuple_GET_SIZE(trait->py_validate);
-
-        //Handle callables without allow_none, default to allow None
-        if (tuple_size < 2) {
-            Py_INCREF(value);
-            return value;
-        }
-
-        allow_none = PyObject_IsTrue(PyTuple_GET_ITEM(trait->py_validate, 1));
-
-        if (allow_none == -1) {
-            return NULL;
-        }
-
-        else if (allow_none) {
-            Py_INCREF(value);
-            return value;
-        }
-    }
-
     return raise_trait_error(trait, obj, name, value);
 }
 
@@ -4093,28 +4101,13 @@ validate_trait_complex(
 
             case 22: /* Callable check: */
                 {
-                    if (PyCallable_Check(value)) {
-                        return value;
-                    }
-                    else if (value == Py_None) {
-                        int allow_none;
-                        int tuple_size = PyTuple_GET_SIZE(trait->py_validate);
-
-                        //Handle callables without allow_none, default to allow None
-                        if (tuple_size < 2) {
-                            Py_INCREF(value);
-                            return value;
-                        }
-
-                        allow_none = PyObject_IsTrue(PyTuple_GET_ITEM(trait->py_validate, 1));
-
-                        if (allow_none == -1) {
-                            return NULL;
-                        }
+                    int valid = _validate_trait_callable(type_info, value);
 
-                        else if (allow_none) {
-                            return value;
-                        }
+                    if (valid == -1) {
+                        return NULL;
+                    }
+                    if (valid == 1) {
+                        goto done;
                     }
                     break;
                 }

traits/tests/test_callable.py

@@ -34,10 +34,37 @@ def instance_method(self):
         pass
 
 
+class OldCallable(BaseCallable):
+    """
+    Old-style Callable, whose validation tuple doesn't include
+    the allow_none field.
+
+    We only care about this case because it's possible that old pickles
+    could include Callable instances whose validation tuple has length 1.
+    """
+    def __init__(self, value=None, **metadata):
+        self.fast_validate = (ValidateTrait.callable,)
+        super().__init__(value, **metadata)
+
+
+class Unbool:
+    """
+    Object that can't be interpreted as a bool, for testing purposes.
+    """
+    def __bool__(self):
+        raise ZeroDivisionError()
+
+
 class MyCallable(HasTraits):
     value = Callable()
 
-    callable_or_str = Either(Callable(), Str)
+    callable_or_str = Either(Callable, Str)
+
+    old_callable_or_str = Either(OldCallable, Str)
+
+    bad_allow_none = Either(Callable(allow_none=Unbool()), Str)
+
+    non_none_callable_or_str = Either(Callable(allow_none=False), Str)
 
 
 class MyBaseCallable(HasTraits):
@@ -91,6 +118,25 @@ def test_callable_in_complex_trait(self):
                 a.callable_or_str = value
             self.assertEqual(a.callable_or_str, old_value)
 
+    def test_compound_callable_refcount(self):
+        # Regression test for enthought/traits#906.
+
+        def my_function():
+            return 72
+
+        a = MyCallable()
+
+        string_value = "some string"
+        callable_value = my_function
+
+        # We seem to need at least 3 repetitions to reliably produce
+        # a crash. Let's hoick it up to 5 to be safe.
+        for _ in range(5):
+            a.callable_or_str = string_value
+            a.callable_or_str = callable_value
+
+        self.assertEqual(a.callable_or_str(), 72)
+
     def test_disallow_none(self):
 
         class MyNewCallable(HasTraits):
@@ -129,12 +175,24 @@ class MyNewCallable2(HasTraits):
             obj.a_non_none_union = None
         obj.a_allow_none_union = None
 
-    def test_old_style_callable(self):
-        class OldCallable(Callable):
-            def __init__(self, value=None, **metadata):
-                self.fast_validate = (ValidateTrait.callable,)
-                super(BaseCallable, self).__init__(value, **metadata)
+    def test_implicitly_allowed_none_in_compound(self):
+        obj = MyCallable()
+        obj.old_callable_or_str = "bob"
+        obj.old_callable_or_str = None
+        self.assertIsNone(obj.old_callable_or_str)
 
+    def test_non_bool_allow_none(self):
+        obj = MyCallable()
+        obj.bad_allow_none = "a string"
+        with self.assertRaises(ZeroDivisionError):
+            obj.bad_allow_none = None
+
+    def test_none_not_allowed_in_compound(self):
+        obj = MyCallable()
+        with self.assertRaises(TraitError):
+            obj.non_none_callable_or_str = None
+
+    def test_old_style_callable(self):
         class MyCallable(HasTraits):
             # allow_none flag should be ineffective
             value = OldCallable()