KIKIMR-19287: add task_stats_drawing script

1 files changed, 203 insertions, 0 deletions

diff --git a/contrib/python/matplotlib/py2/src/_contour_wrapper.cpp b/contrib/python/matplotlib/py2/src/_contour_wrapper.cpp
new file mode 100644
index 0000000000..eedc8a1aec
--- /dev/null
+++ b/contrib/python/matplotlib/py2/src/_contour_wrapper.cpp
@@ -0,0 +1,203 @@
+#include "src/_contour.h"
+#include "src/mplutils.h"
+#include "src/py_exceptions.h"
+
+/* QuadContourGenerator */
+
+typedef struct
+{
+    PyObject_HEAD
+    QuadContourGenerator* ptr;
+} PyQuadContourGenerator;
+
+static PyTypeObject PyQuadContourGeneratorType;
+
+static PyObject* PyQuadContourGenerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds)
+{
+    PyQuadContourGenerator* self;
+    self = (PyQuadContourGenerator*)type->tp_alloc(type, 0);
+    self->ptr = NULL;
+    return (PyObject*)self;
+}
+
+const char* PyQuadContourGenerator_init__doc__ =
+    "QuadContourGenerator(x, y, z, mask, corner_mask, chunk_size)\n"
+    "\n"
+    "Create a new C++ QuadContourGenerator object\n";
+
+static int PyQuadContourGenerator_init(PyQuadContourGenerator* self, PyObject* args, PyObject* kwds)
+{
+    QuadContourGenerator::CoordinateArray x, y, z;
+    QuadContourGenerator::MaskArray mask;
+    int corner_mask;
+    long chunk_size;
+
+    if (!PyArg_ParseTuple(args, "O&O&O&O&il",
+                          &x.converter_contiguous, &x,
+                          &y.converter_contiguous, &y,
+                          &z.converter_contiguous, &z,
+                          &mask.converter_contiguous, &mask,
+                          &corner_mask,
+                          &chunk_size)) {
+        return -1;
+    }
+
+    if (x.empty() || y.empty() || z.empty() ||
+        y.dim(0) != x.dim(0) || z.dim(0) != x.dim(0) ||
+        y.dim(1) != x.dim(1) || z.dim(1) != x.dim(1)) {
+        PyErr_SetString(PyExc_ValueError,
+            "x, y and z must all be 2D arrays with the same dimensions");
+        return -1;
+    }
+
+    if (z.dim(0) < 2 || z.dim(1) < 2) {
+        PyErr_SetString(PyExc_ValueError,
+            "x, y and z must all be at least 2x2 arrays");
+        return -1;
+    }
+
+    // Mask array is optional, if set must be same size as other arrays.
+    if (!mask.empty() && (mask.dim(0) != x.dim(0) || mask.dim(1) != x.dim(1))) {
+        PyErr_SetString(PyExc_ValueError,
+            "If mask is set it must be a 2D array with the same dimensions as x.");
+        return -1;
+    }
+
+    CALL_CPP_INIT("QuadContourGenerator",
+                  (self->ptr = new QuadContourGenerator(
+                       x, y, z, mask, corner_mask, chunk_size)));
+    return 0;
+}
+
+static void PyQuadContourGenerator_dealloc(PyQuadContourGenerator* self)
+{
+    delete self->ptr;
+    Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+const char* PyQuadContourGenerator_create_contour__doc__ =
+    "create_contour(level)\n"
+    "\n"
+    "Create and return a non-filled contour.";
+
+static PyObject* PyQuadContourGenerator_create_contour(PyQuadContourGenerator* self, PyObject* args, PyObject* kwds)
+{
+    double level;
+    if (!PyArg_ParseTuple(args, "d:create_contour", &level)) {
+        return NULL;
+    }
+
+    PyObject* result;
+    CALL_CPP("create_contour", (result = self->ptr->create_contour(level)));
+    return result;
+}
+
+const char* PyQuadContourGenerator_create_filled_contour__doc__ =
+    "create_filled_contour(lower_level, upper_level)\n"
+    "\n"
+    "Create and return a filled contour";
+
+static PyObject* PyQuadContourGenerator_create_filled_contour(PyQuadContourGenerator* self, PyObject* args, PyObject* kwds)
+{
+    double lower_level, upper_level;
+    if (!PyArg_ParseTuple(args, "dd:create_filled_contour",
+                          &lower_level, &upper_level)) {
+        return NULL;
+    }
+
+    if (lower_level >= upper_level)
+    {
+        PyErr_SetString(PyExc_ValueError,
+            "filled contour levels must be increasing");
+        return NULL;
+    }
+
+    PyObject* result;
+    CALL_CPP("create_filled_contour",
+             (result = self->ptr->create_filled_contour(lower_level,
+                                                        upper_level)));
+    return result;
+}
+
+static PyTypeObject* PyQuadContourGenerator_init_type(PyObject* m, PyTypeObject* type)
+{
+    static PyMethodDef methods[] = {
+        {"create_contour", (PyCFunction)PyQuadContourGenerator_create_contour, METH_VARARGS, PyQuadContourGenerator_create_contour__doc__},
+        {"create_filled_contour", (PyCFunction)PyQuadContourGenerator_create_filled_contour, METH_VARARGS, PyQuadContourGenerator_create_filled_contour__doc__},
+        {NULL}
+    };
+
+    memset(type, 0, sizeof(PyTypeObject));
+    type->tp_name = "matplotlib.QuadContourGenerator";
+    type->tp_doc = PyQuadContourGenerator_init__doc__;
+    type->tp_basicsize = sizeof(PyQuadContourGenerator);
+    type->tp_dealloc = (destructor)PyQuadContourGenerator_dealloc;
+    type->tp_flags = Py_TPFLAGS_DEFAULT;
+    type->tp_methods = methods;
+    type->tp_new = PyQuadContourGenerator_new;
+    type->tp_init = (initproc)PyQuadContourGenerator_init;
+
+    if (PyType_Ready(type) < 0) {
+        return NULL;
+    }
+
+    if (PyModule_AddObject(m, "QuadContourGenerator", (PyObject*)type)) {
+        return NULL;
+    }
+
+    return type;
+}
+
+
+/* Module */
+
+extern "C" {
+
+#if PY3K
+static struct PyModuleDef moduledef = {
+    PyModuleDef_HEAD_INIT,
+    "_contour",
+    NULL,
+    0,
+    NULL,
+    NULL,
+    NULL,
+    NULL,
+    NULL
+};
+
+#define INITERROR return NULL
+
+PyMODINIT_FUNC PyInit__contour(void)
+
+#else
+#define INITERROR return
+
+PyMODINIT_FUNC init_contour(void)
+#endif
+
+{
+    PyObject *m;
+
+#if PY3K
+    m = PyModule_Create(&moduledef);
+#else
+    m = Py_InitModule3("_contour", NULL, NULL);
+#endif
+
+    if (m == NULL) {
+        INITERROR;
+    }
+
+    if (!PyQuadContourGenerator_init_type(m, &PyQuadContourGeneratorType)) {
+        INITERROR;
+    }
+
+    import_array();
+
+#if PY3K
+    return m;
+#endif
+}
+
+} // extern "C"