1
/*
2
 * This file implements an object encapsulating a business day
3
 * calendar object for accelerating NumPy datetime business day functions.
4
 *
5
 * Written by Mark Wiebe (mwwiebe@gmail.com)
6
 * Copyright (c) 2011 by Enthought, Inc.
7
 *
8
 * See LICENSE.txt for the license.
9
 */
10

11
#define PY_SSIZE_T_CLEAN
12
#include <Python.h>
13

14
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
15
#define _MULTIARRAYMODULE
16
#include <numpy/arrayobject.h>
17

18
#include "npy_config.h"
19
#include "npy_pycompat.h"
20

21
#include "common.h"
22
#include "numpy/arrayscalars.h"
23
#include "lowlevel_strided_loops.h"
24
#include "_datetime.h"
25
#include "datetime_busday.h"
26
#include "datetime_busdaycal.h"
27

28
NPY_NO_EXPORT int
29 1
PyArray_WeekMaskConverter(PyObject *weekmask_in, npy_bool *weekmask)
30
{
31 1
    PyObject *obj = weekmask_in;
32

33
    /* Make obj into an ASCII string if it is UNICODE */
34 1
    Py_INCREF(obj);
35 1
    if (PyUnicode_Check(obj)) {
36
        /* accept unicode input */
37
        PyObject *obj_str;
38 1
        obj_str = PyUnicode_AsASCIIString(obj);
39 1
        if (obj_str == NULL) {
40 0
            Py_DECREF(obj);
41
            return 0;
42
        }
43 1
        Py_DECREF(obj);
44
        obj = obj_str;
45
    }
46

47 1
    if (PyBytes_Check(obj)) {
48
        char *str;
49
        Py_ssize_t len;
50
        int i;
51

52 1
        if (PyBytes_AsStringAndSize(obj, &str, &len) < 0) {
53 0
            Py_DECREF(obj);
54 1
            return 0;
55
        }
56

57
        /* Length 7 is a string like "1111100" */
58 1
        if (len == 7) {
59 1
            for (i = 0; i < 7; ++i) {
60 1
                switch(str[i]) {
61 1
                    case '0':
62 1
                        weekmask[i] = 0;
63 1
                        break;
64 1
                    case '1':
65 1
                        weekmask[i] = 1;
66 1
                        break;
67
                    default:
68
                        goto general_weekmask_string;
69
                }
70
            }
71

72 1
            goto finish;
73
        }
74

75 1
general_weekmask_string:
76
        /* a string like "SatSun" or "Mon Tue Wed" */
77 1
        memset(weekmask, 0, 7);
78 1
        for (i = 0; i < len; i += 3) {
79 1
            while (isspace(str[i]))
80 1
                ++i;
81

82 1
            if (i == len) {
83
                goto finish;
84
            }
85 1
            else if (i + 2 >= len) {
86
                goto invalid_weekmask_string;
87
            }
88

89 1
            switch (str[i]) {
90 1
                case 'M':
91 1
                    if (str[i+1] == 'o' && str[i+2] == 'n') {
92 1
                        weekmask[0] = 1;
93
                    }
94
                    else {
95
                        goto invalid_weekmask_string;
96
                    }
97 1
                    break;
98 1
                case 'T':
99 1
                    if (str[i+1] == 'u' && str[i+2] == 'e') {
100 1
                        weekmask[1] = 1;
101
                    }
102 1
                    else if (str[i+1] == 'h' && str[i+2] == 'u') {
103 1
                        weekmask[3] = 1;
104
                    }
105
                    else {
106
                        goto invalid_weekmask_string;
107
                    }
108
                    break;
109 1
                case 'W':
110 1
                    if (str[i+1] == 'e' && str[i+2] == 'd') {
111 1
                        weekmask[2] = 1;
112
                    }
113
                    else {
114
                        goto invalid_weekmask_string;
115
                    }
116 1
                    break;
117 1
                case 'F':
118 1
                    if (str[i+1] == 'r' && str[i+2] == 'i') {
119 1
                        weekmask[4] = 1;
120
                    }
121
                    else {
122
                        goto invalid_weekmask_string;
123
                    }
124 1
                    break;
125 1
                case 'S':
126 1
                    if (str[i+1] == 'a' && str[i+2] == 't') {
127 1
                        weekmask[5] = 1;
128
                    }
129 1
                    else if (str[i+1] == 'u' && str[i+2] == 'n') {
130 1
                        weekmask[6] = 1;
131
                    }
132
                    else {
133
                        goto invalid_weekmask_string;
134
                    }
135
                    break;
136
                default:
137
                    goto invalid_weekmask_string;
138
            }
139
        }
140

141
        goto finish;
142

143 1
invalid_weekmask_string:
144 1
        PyErr_Format(PyExc_ValueError,
145
                "Invalid business day weekmask string \"%s\"",
146
                str);
147 1
        Py_DECREF(obj);
148
        return 0;
149
    }
150
    /* Something like [1,1,1,1,1,0,0] */
151 1
    else if (PySequence_Check(obj)) {
152 1
        if (PySequence_Size(obj) != 7 ||
153 1
                        (PyArray_Check(obj) &&
154 0
                         PyArray_NDIM((PyArrayObject *)obj) != 1)) {
155 0
            PyErr_SetString(PyExc_ValueError,
156
                "A business day weekmask array must have length 7");
157 0
            Py_DECREF(obj);
158
            return 0;
159
        }
160
        else {
161
            int i;
162

163 1
            for (i = 0; i < 7; ++i) {
164
                long val;
165 1
                PyObject *f = PySequence_GetItem(obj, i);
166 1
                if (f == NULL) {
167 0
                    Py_DECREF(obj);
168
                    return 0;
169
                }
170

171 1
                val = PyLong_AsLong(f);
172 1
                if (error_converting(val)) {
173 0
                    Py_DECREF(f);
174 0
                    Py_DECREF(obj);
175
                    return 0;
176
                }
177 1
                if (val == 0) {
178 1
                    weekmask[i] = 0;
179
                }
180 0
                else if (val == 1) {
181 0
                    weekmask[i] = 1;
182
                }
183
                else {
184 0
                    PyErr_SetString(PyExc_ValueError,
185
                        "A business day weekmask array must have all "
186
                        "1's and 0's");
187 0
                    Py_DECREF(f);
188 0
                    Py_DECREF(obj);
189
                    return 0;
190
                }
191 1
                Py_DECREF(f);
192
            }
193

194
            goto finish;
195
        }
196
    }
197

198 0
    PyErr_SetString(PyExc_ValueError,
199
            "Couldn't convert object into a business day weekmask");
200 0
    Py_DECREF(obj);
201
    return 0;
202

203 1
finish:
204 1
    Py_DECREF(obj);
205
    return 1;
206
}
207

208
static int
209 1
qsort_datetime_compare(const void *elem1, const void *elem2)
210
{
211 1
    npy_datetime e1 = *(const npy_datetime *)elem1;
212 1
    npy_datetime e2 = *(const npy_datetime *)elem2;
213

214 1
    return (e1 < e2) ? -1 : (e1 == e2) ? 0 : 1;
215
}
216

217
/*
218
 * Sorts the array of dates provided in place and removes
219
 * NaT, duplicates and any date which is already excluded on account
220
 * of the weekmask.
221
 *
222
 * Returns the number of dates left after removing weekmask-excluded
223
 * dates.
224
 */
225
NPY_NO_EXPORT void
226 1
normalize_holidays_list(npy_holidayslist *holidays, npy_bool *weekmask)
227
{
228 1
    npy_datetime *dates = holidays->begin;
229 1
    npy_intp count = holidays->end - dates;
230

231 1
    npy_datetime lastdate = NPY_DATETIME_NAT;
232
    npy_intp trimcount, i;
233
    int day_of_week;
234

235
    /* Sort the dates */
236 1
    qsort(dates, count, sizeof(npy_datetime), &qsort_datetime_compare);
237

238
    /* Sweep through the array, eliminating unnecessary values */
239 1
    trimcount = 0;
240 1
    for (i = 0; i < count; ++i) {
241 1
        npy_datetime date = dates[i];
242

243
        /* Skip any NaT or duplicate */
244 1
        if (date != NPY_DATETIME_NAT && date != lastdate) {
245
            /* Get the day of the week (1970-01-05 is Monday) */
246 1
            day_of_week = (int)((date - 4) % 7);
247 1
            if (day_of_week < 0) {
248 0
                day_of_week += 7;
249
            }
250

251
            /*
252
             * If the holiday falls on a possible business day,
253
             * then keep it.
254
             */
255 1
            if (weekmask[day_of_week] == 1) {
256 1
                dates[trimcount++] = date;
257 1
                lastdate = date;
258
            }
259
        }
260
    }
261

262
    /* Adjust the end of the holidays array */
263 1
    holidays->end = dates + trimcount;
264
}
265

266
/*
267
 * Converts a Python input into a non-normalized list of holidays.
268
 *
269
 * IMPORTANT: This function can't do the normalization, because it doesn't
270
 *            know the weekmask. You must call 'normalize_holiday_list'
271
 *            on the result before using it.
272
 */
273
NPY_NO_EXPORT int
274 1
PyArray_HolidaysConverter(PyObject *dates_in, npy_holidayslist *holidays)
275
{
276 1
    PyArrayObject *dates = NULL;
277 1
    PyArray_Descr *date_dtype = NULL;
278
    npy_intp count;
279

280
    /* Make 'dates' into an array */
281 1
    if (PyArray_Check(dates_in)) {
282 0
        dates = (PyArrayObject *)dates_in;
283 0
        Py_INCREF(dates);
284
    }
285
    else {
286
        PyArray_Descr *datetime_dtype;
287

288
        /* Use the datetime dtype with generic units so it fills it in */
289 1
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
290 1
        if (datetime_dtype == NULL) {
291
            goto fail;
292
        }
293

294
        /* This steals the datetime_dtype reference */
295 1
        dates = (PyArrayObject *)PyArray_FromAny(dates_in, datetime_dtype,
296
                                                0, 0, 0, NULL);
297 1
        if (dates == NULL) {
298
            goto fail;
299
        }
300
    }
301

302 1
    date_dtype = create_datetime_dtype_with_unit(NPY_DATETIME, NPY_FR_D);
303 1
    if (date_dtype == NULL) {
304
        goto fail;
305
    }
306

307 1
    if (!PyArray_CanCastTypeTo(PyArray_DESCR(dates),
308
                                    date_dtype, NPY_SAFE_CASTING)) {
309 0
        PyErr_SetString(PyExc_ValueError, "Cannot safely convert "
310
                        "provided holidays input into an array of dates");
311 0
        goto fail;
312
    }
313 1
    if (PyArray_NDIM(dates) != 1) {
314 0
        PyErr_SetString(PyExc_ValueError, "holidays must be a provided "
315
                        "as a one-dimensional array");
316 0
        goto fail;
317
    }
318

319
    /* Allocate the memory for the dates */
320 1
    count = PyArray_DIM(dates, 0);
321 1
    holidays->begin = PyArray_malloc(sizeof(npy_datetime) * count);
322 1
    if (holidays->begin == NULL) {
323 0
        PyErr_NoMemory();
324 0
        goto fail;
325
    }
326 1
    holidays->end = holidays->begin + count;
327

328
    /* Cast the data into a raw date array */
329 1
    if (PyArray_CastRawArrays(count,
330 1
                            PyArray_BYTES(dates), (char *)holidays->begin,
331
                            PyArray_STRIDE(dates, 0), sizeof(npy_datetime),
332
                            PyArray_DESCR(dates), date_dtype,
333
                            0) != NPY_SUCCEED) {
334
        goto fail;
335
    }
336

337 1
    Py_DECREF(dates);
338 1
    Py_DECREF(date_dtype);
339

340
    return 1;
341

342 0
fail:
343 0
    Py_XDECREF(dates);
344 0
    Py_XDECREF(date_dtype);
345
    return 0;
346
}
347

348
static PyObject *
349 0
busdaycalendar_new(PyTypeObject *subtype,
350
                    PyObject *NPY_UNUSED(args), PyObject *NPY_UNUSED(kwds))
351
{
352
    NpyBusDayCalendar *self;
353

354 0
    self = (NpyBusDayCalendar *)subtype->tp_alloc(subtype, 0);
355 0
    if (self != NULL) {
356
        /* Start with an empty holidays list */
357 0
        self->holidays.begin = NULL;
358 0
        self->holidays.end = NULL;
359

360
        /* Set the weekmask to the default */
361 0
        self->busdays_in_weekmask = 5;
362 0
        self->weekmask[0] = 1;
363 0
        self->weekmask[1] = 1;
364 0
        self->weekmask[2] = 1;
365 0
        self->weekmask[3] = 1;
366 0
        self->weekmask[4] = 1;
367 0
        self->weekmask[5] = 0;
368 0
        self->weekmask[6] = 0;
369
    }
370

371 0
    return (PyObject *)self;
372
}
373

374
static int
375 1
busdaycalendar_init(NpyBusDayCalendar *self, PyObject *args, PyObject *kwds)
376
{
377
    static char *kwlist[] = {"weekmask", "holidays", NULL};
378
    int i, busdays_in_weekmask;
379

380
    /* Clear the holidays if necessary */
381 1
    if (self->holidays.begin != NULL) {
382 0
        PyArray_free(self->holidays.begin);
383 0
        self->holidays.begin = NULL;
384 0
        self->holidays.end = NULL;
385
    }
386

387
    /* Reset the weekmask to the default */
388 1
    self->busdays_in_weekmask = 5;
389 1
    self->weekmask[0] = 1;
390 1
    self->weekmask[1] = 1;
391 1
    self->weekmask[2] = 1;
392 1
    self->weekmask[3] = 1;
393 1
    self->weekmask[4] = 1;
394 1
    self->weekmask[5] = 0;
395 1
    self->weekmask[6] = 0;
396

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

405
    /* Count the number of business days in a week */
406
    busdays_in_weekmask = 0;
407 1
    for (i = 0; i < 7; ++i) {
408 1
        busdays_in_weekmask += self->weekmask[i];
409
    }
410 1
    self->busdays_in_weekmask = busdays_in_weekmask;
411

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

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

422
    return 0;
423
}
424

425
static void
426 1
busdaycalendar_dealloc(NpyBusDayCalendar *self)
427
{
428
    /* Clear the holidays */
429 1
    if (self->holidays.begin != NULL) {
430 1
        PyArray_free(self->holidays.begin);
431 1
        self->holidays.begin = NULL;
432 1
        self->holidays.end = NULL;
433
    }
434

435 1
    Py_TYPE(self)->tp_free((PyObject*)self);
436
}
437

438
static PyObject *
439 1
busdaycalendar_weekmask_get(NpyBusDayCalendar *self)
440
{
441
    PyArrayObject *ret;
442 1
    npy_intp size = 7;
443

444
    /* Allocate a 7-element boolean array */
445 1
    ret = (PyArrayObject *)PyArray_SimpleNew(1, &size, NPY_BOOL);
446 1
    if (ret == NULL) {
447
        return NULL;
448
    }
449

450
    /* Copy the weekmask data */
451 1
    memcpy(PyArray_DATA(ret), self->weekmask, 7);
452

453 1
    return (PyObject *)ret;
454
}
455

456
static PyObject *
457 1
busdaycalendar_holidays_get(NpyBusDayCalendar *self)
458
{
459
    PyArrayObject *ret;
460
    PyArray_Descr *date_dtype;
461 1
    npy_intp size = self->holidays.end - self->holidays.begin;
462

463
    /* Create a date dtype */
464 1
    date_dtype = create_datetime_dtype_with_unit(NPY_DATETIME, NPY_FR_D);
465 1
    if (date_dtype == NULL) {
466
        return NULL;
467
    }
468

469
    /* Allocate a date array (this steals the date_dtype reference) */
470 1
    ret = (PyArrayObject *)PyArray_SimpleNewFromDescr(1, &size, date_dtype);
471 1
    if (ret == NULL) {
472
        return NULL;
473
    }
474

475
    /* Copy the holidays */
476 1
    if (size > 0) {
477 1
        memcpy(PyArray_DATA(ret), self->holidays.begin,
478
                    size * sizeof(npy_datetime));
479
    }
480

481
    return (PyObject *)ret;
482
}
483

484
static PyGetSetDef busdaycalendar_getsets[] = {
485
    {"weekmask",
486
        (getter)busdaycalendar_weekmask_get,
487
        NULL, NULL, NULL},
488
    {"holidays",
489
        (getter)busdaycalendar_holidays_get,
490
        NULL, NULL, NULL},
491

492
    {NULL, NULL, NULL, NULL, NULL}
493
};
494

495
NPY_NO_EXPORT PyTypeObject NpyBusDayCalendar_Type = {
496
    PyVarObject_HEAD_INIT(NULL, 0)
497
    .tp_name = "numpy.busdaycalendar",
498
    .tp_basicsize = sizeof(NpyBusDayCalendar),
499
    .tp_dealloc = (destructor)busdaycalendar_dealloc,
500
    .tp_flags = Py_TPFLAGS_DEFAULT,
501
    .tp_getset = busdaycalendar_getsets,
502
    .tp_init = (initproc)busdaycalendar_init,
503
    .tp_new = busdaycalendar_new,
504
};

Read our documentation on viewing source code .

Loading