python_cluster_reader/src/RawFileReader.c

#include "RawFileReader.h"
#include "arr_desc.h"
#include "data_types.h"
#include "raw_reader.h"

#include <stdbool.h>

//clang-format off
typedef struct {
    PyObject_HEAD FILE *fp;
    // additional fields for size and decoder?
    int dtype;
    bool read_header;
    int detector;
} RawFileReader;
//clang-format on

// Constructor: sets the fp to NULL then tries to open the file
// raises python exception if something goes wrong
// returned object should mean file is open and ready to read
static int RawFileReader_init(RawFileReader *self, PyObject *args,
                              PyObject *kwds) {

    // Parse file name, accepts string or pathlike objects
    char *fname = NULL;
    PyObject *fname_obj = NULL;
    PyObject *fname_bytes = NULL;
    Py_ssize_t len;

    // Should we read the header
    self->read_header = false;
    self->detector = DT_GENERIC;

    static char *kwlist[] = {"fname", "detector_type", "header", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oi|p", kwlist, &fname_obj,
                                     &self->detector, &self->read_header)) {
        return -1;
    }

    if (fname_obj != Py_None)
        if (!PyUnicode_FSConverter(fname_obj, &fname_bytes))
            return -1;

    PyBytes_AsStringAndSize(fname_bytes, &fname, &len);

#ifdef CR_VERBOSE
    printf("fname: %s\n read_header: %d detector type: %d\n", fname,
           self->read_header, self->detector);
#endif
    self->fp = fopen((const char *)fname, "rb");

    // Keep the return code to not return before releasing buffer
    int rc = 0;

    // Raise python exception using information from errno
    if (self->fp == NULL) {
        PyErr_SetFromErrnoWithFilename(PyExc_OSError, fname);
        rc = -1;
    }
    // Release buffer
    Py_DECREF(fname_bytes);

    // Success or fail
    return rc;
}

// Custom destructor to make sure we close the file
static void RawFileReader_dealloc(RawFileReader *self) {
    if (self->fp) {
        fclose(self->fp);
        self->fp = NULL;
    }
    Py_TYPE(self)->tp_free((PyObject *)self);
}

// read method
static PyObject *RawFileReader_read(RawFileReader *self, PyObject *args) {

    self->dtype = NPY_UINT16;
    const int ndim = 3;
    Py_ssize_t n_frames = 1; // default number of frames to read
    if (!PyArg_ParseTuple(args, "|n", &n_frames))
        return NULL;

    npy_intp dims[3] = {n_frames, 400, 400};
    PyObject *frames = PyArray_SimpleNew(ndim, dims, self->dtype);
    char *out_buf = PyArray_DATA((PyArrayObject *)frames);
    PyObject *digital_frames = NULL;
    char *digital_out = NULL;

    PyArray_FILLWBYTE((PyArrayObject *)frames, 0);

    // Optional return the header
    PyObject *header = NULL;
    char *header_out = NULL;
    if (self->read_header) {
        header = PyArray_SimpleNewFromDescr(1, dims, frame_header_dt());
        header_out = PyArray_DATA((PyArrayObject *)header);
    }

    // Both analog and digital data
    if (self->detector == DT_MOENCH_04_AD) {
        digital_frames = PyArray_SimpleNew(ndim, dims, NPY_UINT8);
        PyArray_FILLWBYTE((PyArrayObject *)digital_frames, 0);
        digital_out = PyArray_DATA((PyArrayObject *)digital_frames);
    }

    int64_t n_read = 0;

    switch (self->detector) {
    case DT_MOENCH_03:
        n_read =
            read_raw_m03(self->fp, n_frames, out_buf, (Header *)header_out);
        break;
    case DT_MOENCH_04_A:
    case DT_MOENCH_04_AD:
        n_read = read_raw_m04(self->fp, n_frames, out_buf, digital_out,
                              (Header *)header_out);
        break;
    default:
        break;
    }

    if (n_read != n_frames) {
        // resize the array to match the number of read photons
        // this will reallocate memory

        // create a new_shape struct on the stack
        PyArray_Dims new_shape;

        // reuse dims for the shape
        dims[0] = n_read;
        new_shape.ptr = dims;
        new_shape.len = 3;

        // resize the array to match the number of clusters read
        PyArray_Resize((PyArrayObject *)frames, &new_shape, 1, NPY_ANYORDER);

        if (digital_frames) {
            PyArray_Resize((PyArrayObject *)digital_frames, &new_shape, 1,
                           NPY_ANYORDER);
        }

        // if we also read header we need to reshape the header
        if (self->read_header) {
            new_shape.len = 1;
            PyArray_Resize((PyArrayObject *)header, &new_shape, 1,
                           NPY_ANYORDER);
        }
    }

    // Build up a tuple with the return values
    PyObject *ret = PyTuple_Pack(1, frames);
    if (self->detector == DT_MOENCH_04_AD) {
        _PyTuple_Resize(&ret, 2);
        PyTuple_SET_ITEM(ret, 1, digital_frames);
    }

    if (self->read_header) {
        Py_ssize_t old_size = PyTuple_GET_SIZE(ret);
        _PyTuple_Resize(&ret, old_size + 1);
        PyTuple_SET_ITEM(ret, old_size, header);
    }

    // if we only have one item in the tuple lets return it instead of the tuple
    if (PyTuple_GET_SIZE(ret) == 1) {
        Py_DECREF(ret);
        return frames;
    } else {
        Py_DECREF(frames);
        return ret;
    }
}

// List all methods in our ClusterFileReader class
static PyMethodDef RawFileReader_methods[] = {
    {"read", (PyCFunction)RawFileReader_read, METH_VARARGS,
     "Read and decode frames"},
    {NULL, NULL, 0, NULL} /* Sentinel */
};

// Class defenition
static PyTypeObject RawFileReaderType = {
    PyVarObject_HEAD_INIT(NULL, 0).tp_name = "_creader.RawFileReader",
    .tp_doc = PyDoc_STR("RawFileReader implemented in C"),
    .tp_basicsize = sizeof(RawFileReader),
    .tp_itemsize = 0,
    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .tp_new = PyType_GenericNew,
    .tp_dealloc = (destructor)RawFileReader_dealloc,
    .tp_init = (initproc)RawFileReader_init,
    .tp_methods = RawFileReader_methods,
};

PyObject *init_RawFileReader(PyObject *m) {
    import_array();
    if (PyType_Ready(&RawFileReaderType) < 0)
        return NULL;

    Py_INCREF(&RawFileReaderType);
    if (PyModule_AddObject(m, "RawFileReader", (PyObject *)&RawFileReaderType) <
        0) {
        Py_DECREF(&RawFileReaderType);
        Py_DECREF(m);
        return NULL;
    }
    return m;
}