csonpath_python.c

#include <Python.h>
#include <stdio.h>

#define CSONPATH_JSON PyObject *

#define CSONPATH_NULL Py_None

#define CSONPATH_GET csonpath_python_get

#define CSONPATH_AT csonpath_python_at

#define CSONPATH_REMOVE(o) Py_XDECREF(o)

#define CSONPATH_NEW_OBJECT() PyDict_New()

#define CSONPATH_NEW_ARRAY() PyList_New(0)

#define CSONPATH_IS_OBJ(o) PyDict_Check(o)

#define CSONPATH_IS_ARRAY(o) PyList_Check(o)

#define CSONPATH_IS_STR(o) PyUnicode_Check(o)

#define CSONPATH_IS_NUM(o) PyLong_Check(o)

#define CSONPATH_IS_NULL(o) (o == Py_None)

#define CSONPATH_IS_BOOL(o) PyBool_Check(o)

/* that one assume CSONPATH_IS_BOOL have been called before calling CSONPATH_GET_BOOL */
#define CSONPATH_GET_BOOL(o) (o == Py_True)

#define CSONPATH_GET_STR(obj)			\
    PyUnicode_AsUTF8(obj)

#define CSONPATH_GET_NUM(obj)			\
    PyLong_AsLong(obj)

#define CSONPATH_EQUAL_NUM(obj, to_cmp)		\
    ({						\
	_Bool r = 0;				\
	if (PyLong_Check(obj)) {		\
	    r = PyLong_AsLong(obj) == to_cmp;	\
	}					\
	r;					\
    })

#define CSONPATH_EQUAL_STR(obj, to_cmp)	({			\
    _Bool r = 0;						\
      if (PyUnicode_Check(obj))	{				\
	const char *py_str = PyUnicode_AsUTF8(obj);		\
	r = !strcmp(py_str, to_cmp);				\
      }								\
      r;							\
    })

#define CSONPATH_FORMAT_EXCEPTION(ARGS...)	\
    PyErr_Format(PyExc_ValueError, ARGS)

#define CSONPATH_EXCEPTION(ARGS...)		\
    PyErr_Format(PyExc_ValueError, ARGS);	\
    return -1;

#define CSONPATH_CALLBACK PyObject *

#define CSONPATH_CALLBACK_DATA PyObject *

#define CSONPATH_CALL_CALLBACK(callback, ctx, child_info, tmp, udata) \
    ({								\
	int _cb_ok = 0;							\
	PyObject *arglist;						\
	PyObject *_cb_ret;						\
	if (child_info->type == CSONPATH_STR)				\
	    arglist = Py_BuildValue("(OsOO)", ctx, child_info->key, tmp, udata); \
	else								\
	    arglist = Py_BuildValue("(OiOO)", ctx, child_info->idx, tmp, udata); \
	_cb_ret = PyObject_CallObject(callback, arglist);		\
	Py_DECREF(arglist);						\
	if (!_cb_ret)							\
	    _cb_ok = -1;						\
	else								\
	    Py_DECREF(_cb_ret);						\
	_cb_ok;								\
    })

/* assuming each modification of the object need to go out of the loop */
#define CSONPATH_NEED_FOREACH_REDO(o)	1

#define CSONPATH_REMOVE_CHILD(obj, child_info)				\
  if (child_info.type == CSONPATH_INTEGER) {				\
    PyList_SetSlice(obj, child_info.idx, child_info.idx + 1, NULL);	\
  } else if (child_info.type == CSONPATH_STR) {				\
    PyDict_DelItemString(obj, child_info.key);				\
  }

#define CSONPATH_ARRAY_APPEND(list, el) PyList_Append(list, el)


static int pydict_try_setitemstring(PyObject *obj, const char * const at, PyObject *el, int do_incref)
{
    (void)do_incref;
    if (!PyDict_Check(obj)) {
	PyObject* repr = PyObject_Repr(obj);
	PyObject* str = PyUnicode_AsEncodedString(repr, "utf-8", "~E~");
	const char *bytes = PyBytes_AS_STRING(str);

	PyErr_Format(PyExc_ValueError, "Unable to follow path (%s) in %s: Dict expected", at, bytes);
	Py_XDECREF(repr);
	Py_XDECREF(str);
	return -1;
    }
    if (!at) {
        PyErr_SetString(PyExc_TypeError, "dict keys must be strings");
        return -1;
    }
    if (PyDict_SetItemString(obj, at, el) < 0)
	return -1;
    return 1;
}

static int python_set_or_insert_item(PyObject *array, Py_ssize_t at, PyObject *el, int do_incref)
{
    (void)do_incref;
    if (!PyList_Check(array)) {
	PyErr_Format(PyExc_ValueError, "Unable to follow path: List expected");
	return -1;
    }
    Py_ssize_t s = PyList_Size(array);
    if (at >= s) {
	for (;s < at; ++s)
	    PyList_Insert(array, s, Py_None);

	PyList_Insert(array, at, el);
    } else {
	Py_INCREF(el);
	if (PyList_SetItem(array, at, el) < 0) {
	    Py_DECREF(el);
	    return -1;
	}
    }
    return 1;
}

#define CSONPATH_APPEND_AT(array, at, el, do_incref)			\
    _Generic((at),							\
	     int: python_set_or_insert_item,				\
	     unsigned int: python_set_or_insert_item,			\
	     long: python_set_or_insert_item,				\
	     unsigned long: python_set_or_insert_item,			\
	     long long: python_set_or_insert_item,			\
	     unsigned long long: python_set_or_insert_item,		\
	     const char *: pydict_try_setitemstring,			\
	     char *: pydict_try_setitemstring				\
	) (array, at, el, do_incref)


#ifdef __TINYC__
#define CSONPATH_PRAGMA(...)
#else
#define CSONPATH_PRAGMA(...) _Pragma(__VA_ARGS__)
#endif


#if PY_MAJOR_VERSION > 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 13)
/* PyList_Clear was introduced in Python 3.13 */
#define CSONPATH_ARRAY_CLEAR(o) PyList_Clear(o)
#else
/* For older versions, manually clear the list */
#define CSONPATH_ARRAY_CLEAR(o)            \
    do {                                   \
        while (PyList_Size(o) > 0) {       \
            PyList_SetSlice(o, 0, PyList_Size(o), NULL); \
        }                                  \
    } while (0)
#endif

#define CSONPATH_OBJ_CLEAR(o) PyDict_Clear(o)

#define CSONPATH_FOREACH_ARRAY(obj, child, idx)				\
    CSONPATH_PRAGMA("GCC unroll 8")					\
    for (intptr_t array_len = PyList_Size(obj), idx = 0;		\
	 ({int r = idx < array_len; if (r) child = PyList_GetItem(obj, idx); r;}); ++idx)

#define CSONPATH_FOREACH_OBJ(obj, child, key)				\
    for (Py_ssize_t pos_ = 0; ({					\
		PyObject *key_;						\
		_Bool r = PyDict_Next(obj, &pos_, &key_, &child);	\
		if (r) key = PyUnicode_AsUTF8AndSize(key_, NULL);	\
		r;							\
	    });)

#define CSONPATH_FOREACH_EXT(obj, el, code, key_idx)			\
  if (PyDict_Check(obj)) {						\
    PyObject *key_;							\
    Py_ssize_t pos_ = 0;						\
    while (PyDict_Next(obj, &pos_, &key_, &el)) {			\
      const char *key_idx = PyUnicode_AsUTF8AndSize(key_, NULL);	\
      (void)key_idx;							\
      code								\
	}								\
  } else if (PyList_Check(obj)) {					\
    int array_len_ = PyList_Size(obj);					\
    CSONPATH_PRAGMA("GCC unroll 8")					\
    for (intptr_t key_idx = 0; key_idx < array_len_; ++key_idx) {	\
      el = PyList_GetItem(obj, key_idx);				\
      if (!el) {							\
        PyErr_SetString(PyExc_RuntimeError,				\
                        "list was modified during iteration");		\
        break;								\
      }									\
      code								\
	}								\
  }


static PyObject *csonpath_python_get(PyObject *obj, const char *key) {
    if (PyDict_Check(obj)) {
        PyObject *value = PyDict_GetItemString(obj, key);
        return value ? value : Py_None;
    }
    return Py_None;
}

static PyObject *csonpath_python_at(PyObject *obj, int at) {
    if (PyList_Check(obj)) {
        if (at >= 0 && at < PyList_Size(obj)) {
            PyObject *item = PyList_GetItem(obj, at);
            return item ? item : Py_None;
        }
    }
    return Py_None;
}

#include "csonpath.h"

typedef struct {
    PyObject_HEAD
    struct csonpath *cp;
} PyCsonPathObject;

#define BAD_ARG() ({fprintf(stderr, "bad argument\n"); PyErr_BadArgument(); return NULL;})


static PyObject *PyCsonPath_new(PyTypeObject *subtype, PyObject* args,
				PyObject* kwargs)
{
	static char *kwlist[] = {"path", "return_empty_array", "jq_like", NULL};
	PyCsonPathObject *self = (PyCsonPathObject *)subtype->tp_alloc(subtype, 0);
	const char *s;
	int return_empty_array = 0;
	int jq_like = 0;
	struct csonpath *ret = NULL;
	PyObject *py_ret = NULL;

	if (!self)
		BAD_ARG();

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|pp", kwlist,
					 &s, &return_empty_array, &jq_like)) {
		goto error;
					 }

	ret = csonpath_new_ex(s, CSONPATH_NO_DETROY);
	if (!ret) {
		PyErr_NoMemory();
		goto error;
	}
	if (ret->compile_error) {
		PyErr_Format(PyExc_ValueError, "compilation fail %s",
			     ret->compile_error);
		csonpath_destroy(ret);
		goto error;
	}

	ret->return_empty_array = return_empty_array;

	self->cp = ret;
	return (PyObject *)self;

	error:
		Py_DECREF(self);
	return py_ret;
}

static PyObject *find_all(PyCsonPathObject *self, PyObject* args)
{
    PyObject *json;

    if (!PyArg_ParseTuple(args, "O", &json))
        BAD_ARG();

    PyObject *ret = csonpath_find_all(self->cp, json);
    if (PyErr_Occurred())
        return NULL;

    return ret;
}

static PyObject *find_first(PyCsonPathObject *self, PyObject* args)
{
    PyObject *json;

    if (!PyArg_ParseTuple(args, "O", &json))
        BAD_ARG();

    PyObject *ret = csonpath_find_first(self->cp, json);
    if (PyErr_Occurred())
        return NULL;

    Py_INCREF(ret);
    return ret;
}

static PyObject *print_instructions(PyCsonPathObject *self, PyObject *args, PyObject *kwds)
{
    if (!self->cp) {
        PyErr_SetString(PyExc_RuntimeError, "compiled path is NULL");
        return NULL;
    }
    csonpath_print_instruction(self->cp);
    Py_RETURN_NONE;
}

static PyObject *callback(PyCsonPathObject *self, PyObject* args)
{
    PyObject *json, *callback, *udata = Py_None;

    if (!PyArg_ParseTuple(args, "OO|O", &json, &callback, &udata))
	BAD_ARG();
    int ret = csonpath_callback(self->cp, json, callback, udata);
    if (PyErr_Occurred() || ret < 0) {
        return NULL;
    }
    return PyLong_FromLong(ret);
}

static PyObject *do_remove(PyCsonPathObject *self, PyObject* args)
{
  PyObject *json;

  if (!PyArg_ParseTuple(args, "O", &json))
    BAD_ARG();
  int ret = csonpath_remove(self->cp, json);
  if (PyErr_Occurred() || ret < 0)
      return NULL;
  return PyLong_FromLong(ret);
}

static PyObject *update_or_create(PyCsonPathObject *self, PyObject* args)
{
  PyObject *json;
  PyObject *value;

  if (!PyArg_ParseTuple(args, "OO", &json, &value))
      BAD_ARG();
  int ret = csonpath_update_or_create(self->cp, json, value);
  if (PyErr_Occurred() || ret < 0) {
      return NULL;
  }
  return PyLong_FromLong(ret);
}

static PyObject *update_or_create_callback(PyCsonPathObject *self, PyObject* args)
{
    PyObject *json, *callback, *udata = Py_None;

    if (!PyArg_ParseTuple(args, "OO|O", &json, &callback, &udata))
	BAD_ARG();
    int ret = csonpath_update_or_create_callback(self->cp, json, callback, udata);
    if (PyErr_Occurred() || ret < 0) {
        return NULL;
    }
    return PyLong_FromLong(ret);
}

static void PyCsonPath_dealloc(PyCsonPathObject *self) {
  if (self->cp) {
    csonpath_destroy(self->cp);
    self->cp = NULL;
  }
  Py_TYPE(self)->tp_free((PyObject *)self);
}


static PyObject *PyCsonPath_set_path(PyCsonPathObject *self, PyObject* args) {
    const char *new_path;
    if (!PyArg_ParseTuple(args, "s", &new_path))
      return NULL;
    if (!new_path) return Py_False;

    struct csonpath *new_cjp = csonpath_new_ex(new_path, CSONPATH_NO_DETROY);
    if (!new_cjp) {
        PyErr_NoMemory();
        return NULL;
    }
    if (new_cjp->compile_error) {
        PyErr_Format(PyExc_ValueError, "compilation fail %s",
                     new_cjp->compile_error);
        csonpath_destroy(new_cjp);
        return NULL;
    }
    csonpath_destroy(self->cp);
    self->cp = new_cjp;
    return Py_True;
}

static PyMethodDef csonpath_py_method[] = {
    {"set_path", (PyCFunction)PyCsonPath_set_path, METH_VARARGS, "set_path"},
    {"callback", (PyCFunction)callback, METH_VARARGS, "callback"},
    {"print_instructions", (PyCFunction)print_instructions, METH_NOARGS, "print_instructions"},
    {"update_or_create_callback", (PyCFunction)update_or_create_callback, METH_VARARGS, "update_or_create_callback"},
    {"find_first", (PyCFunction)find_first, METH_VARARGS, "find first elems"},
    {"find_all", (PyCFunction)find_all, METH_VARARGS, "find all elems, if one found, pout it in an array"},
    {"remove", (PyCFunction)do_remove, METH_VARARGS, "remove all elems found"},
    {"update_or_create", (PyCFunction)update_or_create, METH_VARARGS, "update or create"},
    {NULL, NULL, 0, NULL}
};

static PyTypeObject PyCsonPathType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "csonpath.CsonPath",
    .tp_basicsize = sizeof(PyCsonPathObject),
    .tp_dealloc = (destructor)PyCsonPath_dealloc,
    .tp_flags = Py_TPFLAGS_DEFAULT,
    .tp_methods = csonpath_py_method,
    .tp_new = PyCsonPath_new
};

static struct PyModuleDef csonpath_py_mod = {
    PyModuleDef_HEAD_INIT,
    "csonpath",
    NULL,
    -1,
    NULL
};

PyMODINIT_FUNC PyInit_csonpath(void) {
  PyObject *m;
  if (PyType_Ready(&PyCsonPathType) < 0)
    return NULL;

  m = PyModule_Create(&csonpath_py_mod);
  if (!m) return NULL;

  PyModule_AddObject(m, "CsonPath", (PyObject *)&PyCsonPathType);

  return m;
}