#include "Python.h"

#include "pycore_capsule.h" // export _PyCapsule_SetTraverse()

#include "pycore_gc.h" // _PyObject_GC_IS_TRACKED()

#include "pycore_object.h" // _PyObject_GC_TRACK()

#include "pycore_pymem.h" // _PyMem_Strdup()

typedef struct {

PyObject_HEAD

void *pointer;

const char *name;

void *context;

PyCapsule_Destructor destructor;

traverseproc traverse_func;

inquiry clear_func;

} PyCapsule;

#define _PyCapsule_CAST(op) ((PyCapsule *)(op))

static int

_is_legal_capsule(PyObject *op, const char *invalid_capsule)

{

if (!op || !PyCapsule_CheckExact(op)) {

goto error;

}

PyCapsule *capsule = (PyCapsule *)op;

if (capsule->pointer == NULL) {

goto error;

}

return 1;

error:

PyErr_SetString(PyExc_ValueError, invalid_capsule);

return 0;

}

#define is_legal_capsule(capsule, name) \

(_is_legal_capsule(capsule, \

name " called with invalid PyCapsule object"))

static int

name_matches(const char *name1, const char *name2) {

/* if either is NULL, */

if (!name1 || !name2) {

/* they're only the same if they're both NULL. */

return name1 == name2;

}

return !strcmp(name1, name2);

}

PyObject *

PyCapsule_New(void *pointer, const char *name, PyCapsule_Destructor destructor)

{

PyCapsule *capsule;

if (!pointer) {

PyErr_SetString(PyExc_ValueError, "PyCapsule_New called with null pointer");

return NULL;

}

capsule = PyObject_GC_New(PyCapsule, &PyCapsule_Type);

if (capsule == NULL) {

return NULL;

}

capsule->pointer = pointer;

capsule->name = name;

capsule->context = NULL;

capsule->destructor = destructor;

capsule->traverse_func = NULL;

capsule->clear_func = NULL;

// Only track the object by the GC when _PyCapsule_SetTraverse() is called

return (PyObject *)capsule;

}

PyCapsule_IsValid(PyObject *op, const char *name)

{

PyCapsule *capsule = (PyCapsule *)op;

return (capsule != NULL &&

PyCapsule_CheckExact(capsule) &&

capsule->pointer != NULL &&

name_matches(capsule->name, name));

}

void *

PyCapsule_GetPointer(PyObject *op, const char *name)

{

if (!is_legal_capsule(op, "PyCapsule_GetPointer")) {

return NULL;

}

PyCapsule *capsule = (PyCapsule *)op;

if (!name_matches(name, capsule->name)) {

PyErr_SetString(PyExc_ValueError, "PyCapsule_GetPointer called with incorrect name");

return NULL;

}

return capsule->pointer;

}

const char *

PyCapsule_GetName(PyObject *op)

{

if (!is_legal_capsule(op, "PyCapsule_GetName")) {

return NULL;

}

PyCapsule *capsule = (PyCapsule *)op;

return capsule->name;

}

PyCapsule_GetDestructor(PyObject *op)

{

if (!is_legal_capsule(op, "PyCapsule_GetDestructor")) {

return NULL;

}

PyCapsule *capsule = (PyCapsule *)op;

return capsule->destructor;

}

void *

PyCapsule_GetContext(PyObject *op)

{

if (!is_legal_capsule(op, "PyCapsule_GetContext")) {

return NULL;

}

PyCapsule *capsule = (PyCapsule *)op;

return capsule->context;

}

PyCapsule_SetPointer(PyObject *op, void *pointer)

{

if (!is_legal_capsule(op, "PyCapsule_SetPointer")) {

return -1;

}

PyCapsule *capsule = (PyCapsule *)op;

if (!pointer) {

PyErr_SetString(PyExc_ValueError, "PyCapsule_SetPointer called with null pointer");

return -1;

}

capsule->pointer = pointer;

return 0;

}

PyCapsule_SetName(PyObject *op, const char *name)

{

if (!is_legal_capsule(op, "PyCapsule_SetName")) {

return -1;

}

PyCapsule *capsule = (PyCapsule *)op;

capsule->name = name;

return 0;

}

PyCapsule_SetDestructor(PyObject *op, PyCapsule_Destructor destructor)

{

if (!is_legal_capsule(op, "PyCapsule_SetDestructor")) {

return -1;

}

PyCapsule *capsule = (PyCapsule *)op;

capsule->destructor = destructor;

return 0;

}

PyCapsule_SetContext(PyObject *op, void *context)

{

if (!is_legal_capsule(op, "PyCapsule_SetContext")) {

return -1;

}

PyCapsule *capsule = (PyCapsule *)op;

capsule->context = context;

return 0;

}

_PyCapsule_SetTraverse(PyObject *op, traverseproc traverse_func, inquiry clear_func)

{

if (!is_legal_capsule(op, "_PyCapsule_SetTraverse")) {

return -1;

}

PyCapsule *capsule = (PyCapsule *)op;

if (traverse_func == NULL || clear_func == NULL) {

PyErr_SetString(PyExc_ValueError,

"_PyCapsule_SetTraverse() called with NULL callback");

return -1;

}

if (!_PyObject_GC_IS_TRACKED(op)) {

_PyObject_GC_TRACK(op);

}

capsule->traverse_func = traverse_func;

capsule->clear_func = clear_func;

return 0;

}

void *

PyCapsule_Import(const char *name, int Py_UNUSED(no_block))

{

PyObject *object = NULL;

void *return_value = NULL;

char *name_dup = _PyMem_Strdup(name);

if (!name_dup) {

return PyErr_NoMemory();

}

char *trace = name_dup;

while (1) {

char *dot = strchr(trace, '.');

if (dot) {

*dot = '\0';

}

if (object) {

PyObject *attr;

if (PyObject_GetOptionalAttrString(object, trace, &attr) < 0) {

Py_CLEAR(object);

break;

}

Py_SETREF(object, attr);

}

if (!dot) {

// We are done

break;

}

if (!object) {

object = PyImport_ImportModule(name_dup);

if (!object) {

break;

}

}

*dot = '.';

trace = dot + 1;

}

/* compare attribute name to module.name by hand */

if (PyCapsule_IsValid(object, name)) {

PyCapsule *capsule = (PyCapsule *)object;

return_value = capsule->pointer;

}

else if (!PyErr_Occurred()) {

PyErr_Format(PyExc_AttributeError,

"PyCapsule_Import \"%s\" is not valid",

name);

}

Py_XDECREF(object);

PyMem_Free(name_dup);

return return_value;

}

static void

capsule_dealloc(PyObject *op)

{

PyCapsule *capsule = _PyCapsule_CAST(op);

PyObject_GC_UnTrack(op);

if (capsule->destructor) {

capsule->destructor(op);

}

PyObject_GC_Del(op);

}

static PyObject *

capsule_repr(PyObject *o)

{

PyCapsule *capsule = _PyCapsule_CAST(o);

const char *name;

const char *quote;

if (capsule->name) {

quote = "\"";

name = capsule->name;

} else {

quote = "";

name = "NULL";

}

return PyUnicode_FromFormat("<capsule object %s%s%s at %p>",

quote, name, quote, capsule);

}

static int

capsule_traverse(PyObject *self, visitproc visit, void *arg)

{

// Capsule object is only tracked by the GC

// if _PyCapsule_SetTraverse() is called, but

// this can still be manually triggered by gc.get_referents()

PyCapsule *capsule = _PyCapsule_CAST(self);

if (capsule->traverse_func != NULL) {

return capsule->traverse_func(self, visit, arg);

}

return 0;

}

static int

capsule_clear(PyObject *self)

{

// Capsule object is only tracked by the GC

// if _PyCapsule_SetTraverse() is called

PyCapsule *capsule = _PyCapsule_CAST(self);

assert(capsule->clear_func != NULL);

return capsule->clear_func(self);

}

PyDoc_STRVAR(PyCapsule_Type__doc__,

");

PyTypeObject PyCapsule_Type = {

PyVarObject_HEAD_INIT(&PyType_Type, 0)

.tp_name = "PyCapsule",

.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,

.tp_basicsize = sizeof(PyCapsule),

.tp_dealloc = capsule_dealloc,

.tp_repr = capsule_repr,

.tp_doc = PyCapsule_Type__doc__,

.tp_traverse = capsule_traverse,

.tp_clear = capsule_clear,

};