Exemple #1
0
static int
busdaycalendar_init(NpyBusDayCalendar *self, PyObject *args, PyObject *kwds)
{
    static char *kwlist[] = {"weekmask", "holidays", NULL};
    int i, busdays_in_weekmask;

    /* Clear the holidays if necessary */
    if (self->holidays.begin != NULL) {
        PyArray_free(self->holidays.begin);
        self->holidays.begin = NULL;
        self->holidays.end = NULL;
    }

    /* Reset the weekmask to the default */
    self->busdays_in_weekmask = 5;
    self->weekmask[0] = 1;
    self->weekmask[1] = 1;
    self->weekmask[2] = 1;
    self->weekmask[3] = 1;
    self->weekmask[4] = 1;
    self->weekmask[5] = 0;
    self->weekmask[6] = 0;

    /* Parse the parameters */
    if (!PyArg_ParseTupleAndKeywords(args, kwds,
                        "|O&O&:busdaycal", kwlist,
                        &PyArray_WeekMaskConverter, &self->weekmask[0],
                        &PyArray_HolidaysConverter, &self->holidays)) {
        return -1;
    }

    /* Count the number of business days in a week */
    busdays_in_weekmask = 0;
    for (i = 0; i < 7; ++i) {
        busdays_in_weekmask += self->weekmask[i];
    }
    self->busdays_in_weekmask = busdays_in_weekmask;

    /* Normalize the holidays list */
    normalize_holidays_list(&self->holidays, self->weekmask);

    if (self->busdays_in_weekmask == 0) {
        PyErr_SetString(PyExc_ValueError,
                "Cannot construct a numpy.busdaycal with a weekmask of "
                "all zeros");
        return -1;
    }

    return 0;
}
Exemple #2
0
/*
 * This is the 'busday_offset' function exposed for calling
 * from Python.
 */
NPY_NO_EXPORT PyObject *
array_busday_offset(PyObject *NPY_UNUSED(self),
                      PyObject *args, PyObject *kwds)
{
    char *kwlist[] = {"dates", "offsets", "roll",
                      "weekmask", "holidays", "busdaycal", "out", NULL};

    PyObject *dates_in = NULL, *offsets_in = NULL, *out_in = NULL;

    PyArrayObject *dates = NULL, *offsets = NULL, *out = NULL, *ret;
    NPY_BUSDAY_ROLL roll = NPY_BUSDAY_RAISE;
    npy_bool weekmask[7] = {2, 1, 1, 1, 1, 0, 0};
    NpyBusDayCalendar *busdaycal = NULL;
    int i, busdays_in_weekmask;
    npy_holidayslist holidays = {NULL, NULL};
    int allocated_holidays = 1;

    if (!PyArg_ParseTupleAndKeywords(args, kwds,
                                    "OO|O&O&O&O!O:busday_offset", kwlist,
                                    &dates_in,
                                    &offsets_in,
                                    &PyArray_BusDayRollConverter, &roll,
                                    &PyArray_WeekMaskConverter, &weekmask[0],
                                    &PyArray_HolidaysConverter, &holidays,
                                    &NpyBusDayCalendar_Type, &busdaycal,
                                    &out_in)) {
        goto fail;
    }

    /* Make sure only one of the weekmask/holidays and busdaycal is supplied */
    if (busdaycal != NULL) {
        if (weekmask[0] != 2 || holidays.begin != NULL) {
            PyErr_SetString(PyExc_ValueError,
                    "Cannot supply both the weekmask/holidays and the "
                    "busdaycal parameters to busday_offset()");
            goto fail;
        }

        /* Indicate that the holidays weren't allocated by us */
        allocated_holidays = 0;

        /* Copy the private normalized weekmask/holidays data */
        holidays = busdaycal->holidays;
        busdays_in_weekmask = busdaycal->busdays_in_weekmask;
        memcpy(weekmask, busdaycal->weekmask, 7);
    }
    else {
        /*
         * Fix up the weekmask from the uninitialized
         * signal value to a proper default.
         */
        if (weekmask[0] == 2) {
            weekmask[0] = 1;
        }

        /* Count the number of business days in a week */
        busdays_in_weekmask = 0;
        for (i = 0; i < 7; ++i) {
            busdays_in_weekmask += weekmask[i];
        }

        /* The holidays list must be normalized before using it */
        normalize_holidays_list(&holidays, weekmask);
    }

    /* Make 'dates' into an array */
    if (PyArray_Check(dates_in)) {
        dates = (PyArrayObject *)dates_in;
        Py_INCREF(dates);
    }
    else {
        PyArray_Descr *datetime_dtype;

        /* Use the datetime dtype with generic units so it fills it in */
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
        if (datetime_dtype == NULL) {
            goto fail;
        }

        /* This steals the datetime_dtype reference */
        dates = (PyArrayObject *)PyArray_FromAny(dates_in, datetime_dtype,
                                                0, 0, 0, dates_in);
        if (dates == NULL) {
            goto fail;
        }
    }

    /* Make 'offsets' into an array */
    offsets = (PyArrayObject *)PyArray_FromAny(offsets_in,
                            PyArray_DescrFromType(NPY_INT64),
                            0, 0, 0, offsets_in);
    if (offsets == NULL) {
        goto fail;
    }

    /* Make sure 'out' is an array if it's provided */
    if (out_in != NULL) {
        if (!PyArray_Check(out_in)) {
            PyErr_SetString(PyExc_ValueError,
                    "busday_offset: must provide a NumPy array for 'out'");
            goto fail;
        }
        out = (PyArrayObject *)out_in;
    }

    ret = business_day_offset(dates, offsets, out, roll,
                    weekmask, busdays_in_weekmask,
                    holidays.begin, holidays.end);

    Py_DECREF(dates);
    Py_DECREF(offsets);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return out == NULL ? PyArray_Return(ret) : (PyObject *)ret;

fail:
    Py_XDECREF(dates);
    Py_XDECREF(offsets);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return NULL;
}