Add printoptions (#3333)

python/src/print.cpp

@@ -0,0 +1,86 @@
+#include <cstdint>
+#include <cstring>
+#include <sstream>
+
+#include <nanobind/typing.h>
+
+#include "mlx/utils.h"
+#include "python/src/utils.h"
+
+#include "mlx/mlx.h"
+
+namespace mx = mlx::core;
+namespace nb = nanobind;
+using namespace nb::literals;
+
+struct PrintOptionsContext {
+  mx::PrintOptions old_options;
+  mx::PrintOptions new_options;
+  PrintOptionsContext(mx::PrintOptions p) : new_options(p) {}
+  PrintOptionsContext& enter() {
+    old_options = mx::get_global_formatter().format_options;
+    mx::set_printoptions(new_options);
+    return *this;
+  }
+  void exit(nb::args) {
+    mx::set_printoptions(old_options);
+  }
+};
+
+void init_print(nb::module_& m) {
+  // Set Python print formatting options
+  mx::get_global_formatter().capitalize_bool = true;
+  // Expose printing options to Python: allow setting global precision.
+  nb::class_<mx::PrintOptions>(m, "PrintOptions")
+      .def(nb::init<int>(), "precision"_a = -1)
+      .def_rw("precision", &mx::PrintOptions::precision);
+
+  m.def(
+      "set_printoptions",
+      [](int precision) { mx::set_printoptions({precision}); },
+      "precision"_a = mx::get_global_formatter().format_options.precision,
+      R"pbdoc(
+        Set global printing precision for array formatting.
+
+        Example:
+            >>> print(x)  # Uses default precision
+            >>> mx.set_printoptions(precision=3):
+            >>> print(x)  # Uses precision of 3
+            >>> print(x)  # Uses precision of 3 (again)
+
+        Args:
+            precision (int): Number of decimal places.
+        )pbdoc");
+  m.def(
+      "get_printoptions",
+      []() { return mx::get_global_formatter().format_options; },
+      R"pbdoc(
+        Get global printing precision for array formatting.
+
+        Returns:
+        PrintOptions: The format options used for printing arrays.
+        )pbdoc");
+
+  nb::class_<PrintOptionsContext>(m, "_PrintOptionsContext")
+      .def(nb::init<mx::PrintOptions>())
+      .def("__enter__", &PrintOptionsContext::enter)
+      .def("__exit__", &PrintOptionsContext::exit);
+
+  m.def(
+      "printoptions",
+      [](int precision) { return PrintOptionsContext({precision}); },
+      "precision"_a = mx::get_global_formatter().format_options.precision,
+      R"pbdoc(
+        Context manager for setting print options temporarily.
+
+        Example:
+            >>> print(x)  # Uses default precision
+            >>> with mx.printoptions(precision=3):
+            >>>     print(x)  # Uses precision of 3
+            >>> print(x)  # Back to default precision
+
+
+        Args:
+            precision (int): Number of decimal places. Use -1 for default
+        )pbdoc");
+}