2 files changed, 14 insertions, 7 deletions

diff --git a/source4/scripting/python/pytalloc.c b/source4/scripting/python/pytalloc.c
index 4032ff75a4..d0b8cb83f2 100644
--- a/source4/scripting/python/pytalloc.c
+++ b/source4/scripting/python/pytalloc.c
@@ -22,14 +22,16 @@
 void py_talloc_dealloc(PyObject* self)
 {
 	py_talloc_Object *obj = (py_talloc_Object *)self;
-	talloc_free(obj->talloc_ptr);
+	talloc_free(obj->talloc_ctx);
 	PyObject_Del(self);
 }
 
-PyObject *py_talloc_import(PyTypeObject *py_type, void *ptr)
+PyObject *py_talloc_import(PyTypeObject *py_type, TALLOC_CTX *mem_ctx, 
+						   void *ptr)
 {
 	PyObject *ret = PyObject_New(py_talloc_Object, &py_type);
-	ret->talloc_ptr = talloc_reference(NULL, ptr); 
+	ret->talloc_ctx = talloc_reference(mem_ctx, ptr); 
+	ret->ptr = ptr;
 	return ret;
 }
 
@@ -38,5 +40,5 @@ PyObject *py_talloc_default_repr(PyObject *py_obj)
 	py_talloc_Object *obj = (py_talloc_Object *)py_obj;
 
 	return PyString_FromFormat("<talloc: %s>", 
-							   talloc_get_name(obj->talloc_ptr));
+							   talloc_get_name(obj->talloc_ctx));
 }
diff --git a/source4/scripting/python/pytalloc.h b/source4/scripting/python/pytalloc.h
index 735829bfcb..aad5840a67 100644
--- a/source4/scripting/python/pytalloc.h
+++ b/source4/scripting/python/pytalloc.h
@@ -24,7 +24,8 @@
 
 typedef struct {
 	PyObject_HEAD
-	void *talloc_ptr;
+	TALLOC_CTX *talloc_ctx;
+	void *ptr;
 } py_talloc_Object;
 
 /* Deallocate a py_talloc_Object */
@@ -36,9 +37,13 @@ void py_talloc_dealloc(PyObject* self);
 /* FIXME: Call PyErr_SetString(PyExc_TypeError, "expected " __STR(type) ") 
  * when talloc_get_type() returns NULL. */
 #define py_talloc_get_type(py_obj, type) \
-	talloc_get_type(((py_talloc_Object *)py_obj)->talloc_ptr, type)
+	talloc_get_type(py_talloc_get_ptr(py_obj), type)
 
-PyObject *py_talloc_import(PyTypeObject *py_type, void *ptr);
+#define py_talloc_get_ptr(py_obj) ((py_talloc_Object *)py_obj)->ptr
+#define py_talloc_get_mem_ctx(py_obj)  ((py_talloc_Object *)py_obj)->talloc_ctx
+
+PyObject *py_talloc_import_ex(PyTypeObject *py_type, TALLOC_CTX *mem_ctx, void *ptr);
+#define py_talloc_import(py_type, talloc_ptr) py_talloc_import_ex(py_type, talloc_ptr, talloc_ptr)
 
 /* Sane default implementation of reprfunc. */
 PyObject *py_talloc_default_repr(PyObject *py_obj);