<commit_before>/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <faiss/python/python_callbacks.h>
#include <faiss/impl/FaissAssert.h>

namespace {

struct PyThreadLock  {
    PyGILState_STATE gstate;

    PyThreadLock () {
        gstate = PyGILState_Ensure();
    }

    ~PyThreadLock () {
        PyGILState_Release(gstate);
    }
};


};


/***********************************************************
 * Callbacks for IO reader and writer
 ***********************************************************/

PyCallbackIOWriter::PyCallbackIOWriter(PyObject *callback, size_t bs):
    callback(callback), bs(bs) {
    PyThreadLock gil;
    Py_INCREF(callback);
    name = "PyCallbackIOWriter";
}



size_t PyCallbackIOWriter::operator()(const void *ptrv, size_t size, size_t nitems) {
    size_t ws = size * nitems;
    const char *ptr = (const char*)ptrv;
    PyThreadLock gil;
    while(ws > 0) {
        size_t wi = ws > bs ? bs : ws;
        PyObject * result = PyObject_CallFunction(
                callback, "(N)", PyBytes_FromStringAndSize(ptr, wi));
        if (result == NULL) {
            FAISS_THROW_MSG("py err");
        }
        // TODO check nb of bytes written
        ptr += wi;
        ws -= wi;
        Py_DECREF(result);
    }
    return nitems;
}

PyCallbackIOWriter::~PyCallbackIOWriter() {
    PyThreadLock gil;
    Py_DECREF(callback);
}


PyCallbackIOReader::PyCallbackIOReader(PyObject *callback, size_t bs):
    callback(callback), bs(bs) {
    PyThreadLock gil;
    Py_INCREF(callback);
    name = "PyCallbackIOReader";
}

size_t PyCallbackIOReader::operator()(void *ptrv, size_t size, size_t nitems)
{
    size_t rs = size * nitems;
    size_t nb = 0;
    char *ptr = (char*)ptrv;
    PyThreadLock gil;
    while(rs > 0) {
        size_t ri = rs > bs ? bs : rs;
        PyObject * result = PyObject_CallFunction(callback, "(n)", ri);
        if (result == NULL) {
            FAISS_THROW_MSG("propagate py error");
        }
        if(!PyBytes_Check(result)) {
            Py_DECREF(result);
            FAISS_THROW_MSG("read callback did not return a bytes object");
        }
        size_t sz = PyBytes_Size(result);
        if (sz == 0) {
            Py_DECREF(result);
            break;
        }
        nb += sz;
        if (sz > rs) {
            Py_DECREF(result);
            FAISS_THROW_FMT("read callback returned %ld bytes (asked %ld)",
                            sz, rs);
        }
        memcpy(ptr, PyBytes_AsString(result), sz);
        Py_DECREF(result);
        ptr += sz;
        rs -= sz;
    }
    return nb / size;
}

PyCallbackIOReader::~PyCallbackIOReader() {
    PyThreadLock gil;
    Py_DECREF(callback);
}
<commit_msg>Fix format specifier in python_callbacks. (#1361)<commit_after>/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <faiss/python/python_callbacks.h>
#include <faiss/impl/FaissAssert.h>

namespace {

struct PyThreadLock  {
    PyGILState_STATE gstate;

    PyThreadLock () {
        gstate = PyGILState_Ensure();
    }

    ~PyThreadLock () {
        PyGILState_Release(gstate);
    }
};


};


/***********************************************************
 * Callbacks for IO reader and writer
 ***********************************************************/

PyCallbackIOWriter::PyCallbackIOWriter(PyObject *callback, size_t bs):
    callback(callback), bs(bs) {
    PyThreadLock gil;
    Py_INCREF(callback);
    name = "PyCallbackIOWriter";
}



size_t PyCallbackIOWriter::operator()(const void *ptrv, size_t size, size_t nitems) {
    size_t ws = size * nitems;
    const char *ptr = (const char*)ptrv;
    PyThreadLock gil;
    while(ws > 0) {
        size_t wi = ws > bs ? bs : ws;
        PyObject * result = PyObject_CallFunction(
                callback, "(N)", PyBytes_FromStringAndSize(ptr, wi));
        if (result == NULL) {
            FAISS_THROW_MSG("py err");
        }
        // TODO check nb of bytes written
        ptr += wi;
        ws -= wi;
        Py_DECREF(result);
    }
    return nitems;
}

PyCallbackIOWriter::~PyCallbackIOWriter() {
    PyThreadLock gil;
    Py_DECREF(callback);
}


PyCallbackIOReader::PyCallbackIOReader(PyObject *callback, size_t bs):
    callback(callback), bs(bs) {
    PyThreadLock gil;
    Py_INCREF(callback);
    name = "PyCallbackIOReader";
}

size_t PyCallbackIOReader::operator()(void *ptrv, size_t size, size_t nitems)
{
    size_t rs = size * nitems;
    size_t nb = 0;
    char *ptr = (char*)ptrv;
    PyThreadLock gil;
    while(rs > 0) {
        size_t ri = rs > bs ? bs : rs;
        PyObject * result = PyObject_CallFunction(callback, "(n)", ri);
        if (result == NULL) {
            FAISS_THROW_MSG("propagate py error");
        }
        if(!PyBytes_Check(result)) {
            Py_DECREF(result);
            FAISS_THROW_MSG("read callback did not return a bytes object");
        }
        size_t sz = PyBytes_Size(result);
        if (sz == 0) {
            Py_DECREF(result);
            break;
        }
        nb += sz;
        if (sz > rs) {
            Py_DECREF(result);
            FAISS_THROW_FMT("read callback returned %zd bytes (asked %zd)",
                            sz, rs);
        }
        memcpy(ptr, PyBytes_AsString(result), sz);
        Py_DECREF(result);
        ptr += sz;
        rs -= sz;
    }
    return nb / size;
}

PyCallbackIOReader::~PyCallbackIOReader() {
    PyThreadLock gil;
    Py_DECREF(callback);
}

{'fraction_non_alphanumeric': 0.15254833040421792, 'fraction_numerical': 0.00210896309314587, 'mean_word_length': 2.5151327980234712, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3289994', 'n_tokens_mistral': 1887, 'n_tokens_neox': 1840, 'n_words': 433}