From d645ec67005b7eac18ac5e330f090aa93944c64c Mon Sep 17 00:00:00 2001
From: Roberto Rossini <71787608+robomics@users.noreply.github.com>
Date: Sun, 20 Oct 2024 20:53:06 +0200
Subject: [PATCH] Validate pyarrow's version at runtime

---
 src/include/hictkpy/nanobind.hpp | 83 +++++++++++++++++++++++++++++++-
 src/pixel_selector.cpp           |  2 +-
 src/to_pyarrow.cpp               |  2 +-
 3 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/src/include/hictkpy/nanobind.hpp b/src/include/hictkpy/nanobind.hpp
index 3300554..d5676e6 100644
--- a/src/include/hictkpy/nanobind.hpp
+++ b/src/include/hictkpy/nanobind.hpp
@@ -28,15 +28,94 @@ HICTKPY_DISABLE_WARNING_USELESS_CAST
 HICTKPY_DISABLE_WARNING_POP
 // clang-format on
 
+#include <cassert>
+#include <mutex>
 #include <string>
 
 inline nanobind::module_ import_module_checked(const std::string& module_name) {
   try {
     return nanobind::module_::import_(module_name.c_str());
   } catch (nanobind::python_error& e) {
-    // NOLINTNEXTLINE(*-pro-type-vararg)
+    // NOLINTNEXTLINE(*-vararg)
     nanobind::raise_from(e, PyExc_ModuleNotFoundError,
-                         "To enable %s support, please install %s with: pip install 'hictkpy[%s]'",
+                         "To enable %s support, please install %s with: pip install 'hictkpy[%s]'\n"
+                         "Alternatively, you can install hictkpy with all its dependencies by "
+                         "running: pip install 'hictkpy[all]'",
                          module_name.c_str(), module_name.c_str(), module_name.c_str());
   }
 }
+
+// NOLINTNEXTLINE(*-avoid-magic-numbers)
+inline nanobind::module_ import_pyarrow_checked(int min_version_major = 16,
+                                                int min_version_minor = 0,
+                                                int min_version_patch = 0) {
+  assert(min_version_major >= 0);
+  assert(min_version_minor >= 0);
+  assert(min_version_patch >= 0);
+
+  static bool version_ok{false};
+  static std::mutex mtx{};
+
+  auto pa = import_module_checked("pyarrow");
+
+  [[maybe_unused]] const auto lck = std::scoped_lock(mtx);
+  if (version_ok) {
+    return pa;
+  }
+
+  static std::string error_msg{};
+  error_msg.clear();
+  try {
+    auto metadata = nanobind::module_::import_("importlib.metadata");
+
+    const auto version = nanobind::cast<std::vector<std::string>>(
+        metadata.attr("version")("pyarrow").attr("split")("."));
+    if (version.size() < 3) {
+      throw nanobind::import_error(
+          "unable to detect pyarrow version: assuming pyarrow's version is not compatible: please "
+          "install a compatible version of pyarrow with: pip install 'hictkpy[pyarrow]'");
+    }
+
+    const auto major_version_found = std::stoi(version[0]);
+    const auto minor_version_found = std::stoi(version[1]);
+    const auto patch_version_found = std::stoi(version[2]);
+
+    version_ok = major_version_found >= min_version_major;
+    version_ok |=
+        major_version_found == min_version_major && minor_version_found >= min_version_minor;
+    version_ok |= major_version_found == min_version_major &&
+                  minor_version_found == min_version_minor &&
+                  patch_version_found >= min_version_patch;
+
+    if (!version_ok) {
+      // Poor man's formatting
+      error_msg = "pyarrow ";
+      for (const auto& tok : version) {
+        error_msg += tok + ".";
+      }
+      if (version.size() > 1) {
+        error_msg.pop_back();
+      }
+      error_msg +=
+          " is too old to be used with hictkpy: please "
+          "install a compatible version with: pip install 'hictkpy[pyarrow]'";
+      throw nanobind::import_error(error_msg.c_str());
+    }
+  } catch (const nanobind::builtin_exception&) {
+    throw;
+  } catch (const std::exception& e) {
+    error_msg = "unable to parse pyarrow version: ";
+    error_msg += e.what();
+    error_msg +=
+        ". Assuming pyarrow's version is not compatible: please "
+        "install a compatible version of pyarrow with: pip install 'hictkpy[pyarrow]'";
+    throw nanobind::import_error(error_msg.c_str());
+  } catch (...) {
+    throw nanobind::import_error(
+        "unable to parse pyarrow version: Assuming pyarrow's version is not compatible: please "
+        "install a compatible version of pyarrow with: pip install 'hictkpy[pyarrow]'");
+  }
+
+  assert(version_ok);
+  return pa;
+}
diff --git a/src/pixel_selector.cpp b/src/pixel_selector.cpp
index 53a69ec..a3f71e3 100644
--- a/src/pixel_selector.cpp
+++ b/src/pixel_selector.cpp
@@ -185,7 +185,7 @@ template <typename N, typename PixelSelector>
 }
 
 nb::object PixelSelector::to_arrow(std::string_view span) const {
-  std::ignore = import_module_checked("pyarrow");
+  std::ignore = import_pyarrow_checked();
 
   const auto query_span = parse_span(span);
   auto table = std::visit(
diff --git a/src/to_pyarrow.cpp b/src/to_pyarrow.cpp
index fedc8e5..aa4ded1 100644
--- a/src/to_pyarrow.cpp
+++ b/src/to_pyarrow.cpp
@@ -114,7 +114,7 @@ static void release_arrow_array_streamPyCapsule(PyObject* capsule) {
 nb::object export_pyarrow_table(std::shared_ptr<arrow::Table> arrow_table) {
   assert(arrow_table);
 
-  const auto pa = import_module_checked("pyarrow");
+  const auto pa = import_pyarrow_checked();
 
   std::vector<nb::object> columns_py(static_cast<std::size_t>(arrow_table->num_columns()));
   std::vector<std::shared_ptr<arrow::ChunkedArray>> columns(columns_py.size());