pytalloc: Refactor the pytalloc_reference and pytalloc_steal to use a common method.

Signed-off-by: Kristján Valur <kristjan@rvx.is>
Reviewed-by: Noel Power <npower@samba.org>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

diff --git a/lib/talloc/pytalloc.h b/lib/talloc/pytalloc.h
index 37416fcb46fdb65d3d215b0ab2197e7656d9bfba..7db6c33cf014a147429b78dc996586b389983454 100644 (file)
--- a/lib/talloc/pytalloc.h
+++ b/lib/talloc/pytalloc.h
@@ -29,10 +29,10 @@ typedef struct {
        void *ptr; /* eg the array element */
 } pytalloc_Object;
 
-/* Return the PyTypeObject for pytalloc_Object. Returns a new reference. */
+/* Return the PyTypeObject for pytalloc_Object. Returns a borrowed reference. */
 PyTypeObject *pytalloc_GetObjectType(void);
 
-/* Return the PyTypeObject for pytalloc_BaseObject. Returns a new reference. */
+/* Return the PyTypeObject for pytalloc_BaseObject. Returns a borrowed reference. */
 PyTypeObject *pytalloc_GetBaseObjectType(void);
 
 /* Check whether a specific object is a talloc Object. */

diff --git a/lib/talloc/pytalloc_util.c b/lib/talloc/pytalloc_util.c
index 265bdcd7e8feeb531a04ed18a9b2738e1f5234ee..af835e437c51ecb6e2862c383832cd914c5ff7f4 100644 (file)
--- a/lib/talloc/pytalloc_util.c
+++ b/lib/talloc/pytalloc_util.c
@@ -24,6 +24,9 @@
 #include <assert.h>
 #include "pytalloc_private.h"
 
+static PyObject *pytalloc_steal_or_reference(PyTypeObject *py_type,
+                                        TALLOC_CTX *mem_ctx, void *ptr, bool steal);
+
 _PUBLIC_ PyTypeObject *pytalloc_GetObjectType(void)
 {
        static PyTypeObject *type = NULL;
@@ -90,56 +93,7 @@ static PyTypeObject *pytalloc_GetGenericObjectType(void)
 _PUBLIC_ PyObject *pytalloc_steal_ex(PyTypeObject *py_type, TALLOC_CTX *mem_ctx,
                                     void *ptr)
 {
-       PyTypeObject *BaseObjectType = pytalloc_GetBaseObjectType();
-       PyTypeObject *ObjectType = pytalloc_GetObjectType();
-
-       if (mem_ctx == NULL) {
-               return PyErr_NoMemory();
-       }
-
-       if (PyType_IsSubtype(py_type, BaseObjectType)) {
-               pytalloc_BaseObject *ret
-                       = (pytalloc_BaseObject *)py_type->tp_alloc(py_type, 0);
-
-               ret->talloc_ctx = talloc_new(NULL);
-               if (ret->talloc_ctx == NULL) {
-                       return NULL;
-               }
-
-               /*
-                * This allows us to keep multiple references to this object -
-                * we only reference this context, which is per ptr, not the
-                * talloc_ctx, which is per pytalloc_Object
-                */
-               if (talloc_steal(ret->talloc_ctx, mem_ctx) == NULL) {
-                       return NULL;
-               }
-               ret->talloc_ptr_ctx = mem_ctx;
-               talloc_set_name_const(ret->talloc_ctx, py_type->tp_name);
-               ret->ptr = ptr;
-               return (PyObject *)ret;
-
-       } else if (PyType_IsSubtype(py_type, ObjectType)) {
-               pytalloc_Object *ret
-                       = (pytalloc_Object *)py_type->tp_alloc(py_type, 0);
-
-               ret->talloc_ctx = talloc_new(NULL);
-               if (ret->talloc_ctx == NULL) {
-                       return NULL;
-               }
-
-               if (talloc_steal(ret->talloc_ctx, mem_ctx) == NULL) {
-                       return NULL;
-               }
-               talloc_set_name_const(ret->talloc_ctx, py_type->tp_name);
-               ret->ptr = ptr;
-               return (PyObject *)ret;
-       } else {
-               PyErr_SetString(PyExc_RuntimeError,
-                               "pytalloc_steal_ex() called for object type "
-                               "not based on talloc");
-               return NULL;
-       }
+       return pytalloc_steal_or_reference(py_type, mem_ctx, ptr, true);
 }
 
 /**
@@ -147,7 +101,7 @@ _PUBLIC_ PyObject *pytalloc_steal_ex(PyTypeObject *py_type, TALLOC_CTX *mem_ctx,
  */
 _PUBLIC_ PyObject *pytalloc_steal(PyTypeObject *py_type, void *ptr)
 {
-       return pytalloc_steal_ex(py_type, ptr, ptr);
+       return pytalloc_steal_or_reference(py_type, ptr, ptr, true);
 }
 
 
@@ -165,6 +119,18 @@ _PUBLIC_ PyObject *pytalloc_steal(PyTypeObject *py_type, void *ptr)
 _PUBLIC_ PyObject *pytalloc_reference_ex(PyTypeObject *py_type,
                                         TALLOC_CTX *mem_ctx, void *ptr)
 {
+       return pytalloc_steal_or_reference(py_type, mem_ctx, ptr, false);
+}
+
+
+/**
+ * Internal function that either steals or referecences the talloc
+ * pointer into a new talloc context.
+ */
+static PyObject *pytalloc_steal_or_reference(PyTypeObject *py_type,
+                                        TALLOC_CTX *mem_ctx, void *ptr, bool steal)
+{
+       bool ok = false;
        PyTypeObject *BaseObjectType = pytalloc_GetBaseObjectType();
        PyTypeObject *ObjectType = pytalloc_GetObjectType();
 
@@ -179,7 +145,12 @@ _PUBLIC_ PyObject *pytalloc_reference_ex(PyTypeObject *py_type,
                if (ret->talloc_ctx == NULL) {
                        return NULL;
                }
-               if (talloc_reference(ret->talloc_ctx, mem_ctx) == NULL) {
+               if (steal) {
+                       ok = (talloc_steal(ret->talloc_ctx, mem_ctx) != NULL);
+               } else {
+                       ok = (talloc_reference(ret->talloc_ctx, mem_ctx) != NULL);
+               }
+               if (!ok) {
                        return NULL;
                }
                talloc_set_name_const(ret->talloc_ctx, py_type->tp_name);
@@ -193,7 +164,12 @@ _PUBLIC_ PyObject *pytalloc_reference_ex(PyTypeObject *py_type,
                if (ret->talloc_ctx == NULL) {
                        return NULL;
                }
-               if (talloc_reference(ret->talloc_ctx, mem_ctx) == NULL) {
+               if (steal) {
+                       ok = (talloc_steal(ret->talloc_ctx, mem_ctx) != NULL);
+               } else {
+                       ok = (talloc_reference(ret->talloc_ctx, mem_ctx) != NULL);
+               }
+               if (!ok) {
                        return NULL;
                }
                talloc_set_name_const(ret->talloc_ctx, py_type->tp_name);
@@ -201,8 +177,7 @@ _PUBLIC_ PyObject *pytalloc_reference_ex(PyTypeObject *py_type,
                return (PyObject *)ret;
        } else {
                PyErr_SetString(PyExc_RuntimeError,
-                               "pytalloc_reference_ex() called for object type "
-                               "not based on talloc");
+                               "Expected type based on talloc");
                return NULL;
        }
 }