From f9b87c976e6185a1c17d3c1ef0d18fe1fb7c8fc4 Mon Sep 17 00:00:00 2001
From: Franco Cipollone <53065142+francocipollone@users.noreply.github.com>
Date: Fri, 10 Dec 2021 18:26:33 -0300
Subject: [PATCH] Adds set_zlim() method

Adds set_zlim() which allows setting the minimum and maximum limit in the z-axis.
---
 matplotlibcpp.h | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/matplotlibcpp.h b/matplotlibcpp.h
index d95d46a..2843458 100644
--- a/matplotlibcpp.h
+++ b/matplotlibcpp.h
@@ -2522,6 +2522,55 @@ inline void set_zlabel(const std::string &str, const std::map<std::string, std::
     if (res) Py_DECREF(res);
 }
 
+inline void set_zlim(double left, double right)
+{
+    detail::_interpreter::get();
+
+    // Same as with plot_surface: We lazily load the modules here the first time
+    // this function is called because I'm not sure that we can assume "matplotlib
+    // installed" implies "mpl_toolkits installed" on all platforms, and we don't
+    // want to require it for people who don't need 3d plots.
+    static PyObject *mpl_toolkitsmod = nullptr, *axis3dmod = nullptr;
+    if (!mpl_toolkitsmod) {
+        PyObject* mpl_toolkits = PyString_FromString("mpl_toolkits");
+        PyObject* axis3d = PyString_FromString("mpl_toolkits.mplot3d");
+        if (!mpl_toolkits || !axis3d) { throw std::runtime_error("couldnt create string"); }
+
+        mpl_toolkitsmod = PyImport_Import(mpl_toolkits);
+        Py_DECREF(mpl_toolkits);
+        if (!mpl_toolkitsmod) { throw std::runtime_error("Error loading module mpl_toolkits!"); }
+
+        axis3dmod = PyImport_Import(axis3d);
+        Py_DECREF(axis3d);
+        if (!axis3dmod) { throw std::runtime_error("Error loading module mpl_toolkits.mplot3d!"); }
+    }
+
+    PyObject* list = PyList_New(2);
+    PyList_SetItem(list, 0, PyFloat_FromDouble(left));
+    PyList_SetItem(list, 1, PyFloat_FromDouble(right));
+
+    PyObject* args = PyTuple_New(1);
+    PyTuple_SetItem(args, 0, list);
+
+    PyObject *ax =
+    PyObject_CallObject(detail::_interpreter::get().s_python_function_gca,
+      detail::_interpreter::get().s_python_empty_tuple);
+    if (!ax) throw std::runtime_error("Call to gca() failed.");
+    Py_INCREF(ax);
+
+    PyObject *zlim = PyObject_GetAttrString(ax, "set_zlim");
+    if (!zlim) throw std::runtime_error("Attribute set_zlim not found.");
+    Py_INCREF(zlim);
+
+    PyObject *res = PyObject_Call(zlim, args, {});
+    if (!res) throw std::runtime_error("Call to set_zlim() failed.");
+    Py_DECREF(zlim);
+
+    Py_DECREF(ax);
+    Py_DECREF(args);
+    if (res) Py_DECREF(res);
+}
+
 inline void grid(bool flag)
 {
     detail::_interpreter::get();