From 1facfdb2ff93e3c8ac507a825927a64c6f9f8bd0 Mon Sep 17 00:00:00 2001
From: Dustin Spicuzza <dustin@virtualroadside.com>
Date: Fri, 30 Dec 2022 00:28:09 -0500
Subject: [PATCH] Fix inheriting from Sendable

- initSendable needs to be passed a reference
- The SendableRegistry needs to be cleared on shutdown to ensure that ntcore doesn't call into python
- Add tests to make sure it continues to work
---
 gen/Sendable.yml                         |   5 +
 tests/cpp/pyproject.toml                 |   1 +
 tests/cpp/wpiutil_test/module.cpp        |   4 +
 tests/cpp/wpiutil_test/sendable_test.cpp | 127 +++++++++++++++++++++++
 tests/test_sendable.py                   |  35 +++++++
 wpiutil/src/main.cpp                     |  13 ++-
 6 files changed, 184 insertions(+), 1 deletion(-)
 create mode 100644 tests/cpp/wpiutil_test/sendable_test.cpp
 create mode 100644 tests/test_sendable.py

diff --git a/gen/Sendable.yml b/gen/Sendable.yml
index d54e542..6b4c4f4 100644
--- a/gen/Sendable.yml
+++ b/gen/Sendable.yml
@@ -8,3 +8,8 @@ classes:
     shared_ptr: true
     methods:
       InitSendable:
+        virtual_xform: |
+          [&](py::function fn) {
+            auto builderHandle = py::cast(builder, py::return_value_policy::reference);
+            fn(builderHandle);
+          }
\ No newline at end of file
diff --git a/tests/cpp/pyproject.toml b/tests/cpp/pyproject.toml
index 4e13c51..a036400 100644
--- a/tests/cpp/pyproject.toml
+++ b/tests/cpp/pyproject.toml
@@ -10,6 +10,7 @@ extension = "module"
 depends = ["wpiutil"]
 sources = [
     "wpiutil_test/module.cpp",
+    "wpiutil_test/sendable_test.cpp",
 ]
 
 [tool.robotpy-build.metadata]
diff --git a/tests/cpp/wpiutil_test/module.cpp b/tests/cpp/wpiutil_test/module.cpp
index 596d387..1e3f9fa 100644
--- a/tests/cpp/wpiutil_test/module.cpp
+++ b/tests/cpp/wpiutil_test/module.cpp
@@ -144,9 +144,13 @@ wpi::json cast_json_val(std::function<wpi::json()> fn) {
     return fn();
 }
 
+void sendable_test(py::module &m);
+
 
 RPYBUILD_PYBIND11_MODULE(m) {
 
+    sendable_test(m);
+
     // array
     m.def("load_array_int", &load_array_int);
     m.def("load_array_int1", &load_array_int1);
diff --git a/tests/cpp/wpiutil_test/sendable_test.cpp b/tests/cpp/wpiutil_test/sendable_test.cpp
new file mode 100644
index 0000000..68c822e
--- /dev/null
+++ b/tests/cpp/wpiutil_test/sendable_test.cpp
@@ -0,0 +1,127 @@
+
+#include <robotpy_build.h>
+#include <pybind11/stl.h>
+#include <pybind11/functional.h>
+#include <wpi/sendable/SendableBuilder.h>
+#include <wpi/sendable/SendableRegistry.h>
+
+class MySendableBuilder : public wpi::SendableBuilder {
+public:
+  MySendableBuilder(py::dict keys) : keys(keys) {}
+
+  ~MySendableBuilder() {
+    // leak this so the python interpreter doesn't crash on shutdown
+    keys.release();
+  }
+
+  void SetSmartDashboardType(std::string_view type) override {}
+
+  void SetActuator(bool value) override {}
+
+  void SetSafeState(std::function<void()> func) override {}
+
+  void AddBooleanProperty(std::string_view key, std::function<bool()> getter,
+                          std::function<void(bool)> setter) override {}
+
+  void AddIntegerProperty(std::string_view key, std::function<int64_t()> getter,
+                          std::function<void(int64_t)> setter) override {}
+
+  void AddFloatProperty(std::string_view key, std::function<float()> getter,
+                        std::function<void(float)> setter) override {}
+
+  void AddDoubleProperty(std::string_view key, std::function<double()> getter,
+                         std::function<void(double)> setter) override {
+    py::gil_scoped_acquire gil;
+    py::object pykey = py::cast(key);
+    keys[pykey] = std::make_tuple(getter, setter);
+  }
+
+  void
+  AddStringProperty(std::string_view key, std::function<std::string()> getter,
+                    std::function<void(std::string_view)> setter) override {}
+
+  void AddBooleanArrayProperty(
+      std::string_view key, std::function<std::vector<int>()> getter,
+      std::function<void(std::span<const int>)> setter) override {}
+
+  void AddIntegerArrayProperty(
+      std::string_view key, std::function<std::vector<int64_t>()> getter,
+      std::function<void(std::span<const int64_t>)> setter) override {}
+
+  void AddFloatArrayProperty(
+      std::string_view key, std::function<std::vector<float>()> getter,
+      std::function<void(std::span<const float>)> setter) override {}
+
+  void AddDoubleArrayProperty(
+      std::string_view key, std::function<std::vector<double>()> getter,
+      std::function<void(std::span<const double>)> setter) override {}
+
+  void AddStringArrayProperty(
+      std::string_view key, std::function<std::vector<std::string>()> getter,
+      std::function<void(std::span<const std::string>)> setter) override {}
+
+  void AddRawProperty(
+      std::string_view key, std::string_view typeString,
+      std::function<std::vector<uint8_t>()> getter,
+      std::function<void(std::span<const uint8_t>)> setter) override {}
+
+  void AddSmallStringProperty(
+      std::string_view key,
+      std::function<std::string_view(wpi::SmallVectorImpl<char> &buf)> getter,
+      std::function<void(std::string_view)> setter) override {}
+
+  void AddSmallBooleanArrayProperty(
+      std::string_view key,
+      std::function<std::span<const int>(wpi::SmallVectorImpl<int> &buf)>
+          getter,
+      std::function<void(std::span<const int>)> setter) override {}
+
+  void AddSmallIntegerArrayProperty(
+      std::string_view key,
+      std::function<
+          std::span<const int64_t>(wpi::SmallVectorImpl<int64_t> &buf)>
+          getter,
+      std::function<void(std::span<const int64_t>)> setter) override {}
+
+  void AddSmallFloatArrayProperty(
+      std::string_view key,
+      std::function<std::span<const float>(wpi::SmallVectorImpl<float> &buf)>
+          getter,
+      std::function<void(std::span<const float>)> setter) override {}
+
+  void AddSmallDoubleArrayProperty(
+      std::string_view key,
+      std::function<std::span<const double>(wpi::SmallVectorImpl<double> &buf)>
+          getter,
+      std::function<void(std::span<const double>)> setter) override {}
+
+  void AddSmallStringArrayProperty(
+      std::string_view key,
+      std::function<
+          std::span<const std::string>(wpi::SmallVectorImpl<std::string> &buf)>
+          getter,
+      std::function<void(std::span<const std::string>)> setter) override {}
+
+  void AddSmallRawProperty(
+      std::string_view key, std::string_view typeString,
+      std::function<std::span<uint8_t>(wpi::SmallVectorImpl<uint8_t> &buf)>
+          getter,
+      std::function<void(std::span<const uint8_t>)> setter) override {}
+
+  wpi::SendableBuilder::BackendKind GetBackendKind() const override {
+    return wpi::SendableBuilder::BackendKind::kUnknown;
+  }
+
+  bool IsPublished() const override { return false; }
+  void Update() override {}
+  void ClearProperties() override {}
+
+  py::dict keys;
+};
+
+void Publish(wpi::SendableRegistry::UID sendableUid, py::dict keys) {
+  auto builder = std::make_unique<MySendableBuilder>(keys);
+  wpi::SendableRegistry::Publish(sendableUid, std::move(builder));
+}
+
+void sendable_test(py::module &m) { m.def("publish", Publish); }
\ No newline at end of file
diff --git a/tests/test_sendable.py b/tests/test_sendable.py
new file mode 100644
index 0000000..5f80e33
--- /dev/null
+++ b/tests/test_sendable.py
@@ -0,0 +1,35 @@
+import typing
+import wpiutil
+from wpiutil_test import module
+
+
+class MySendable(wpiutil.Sendable):
+    def __init__(self):
+        super().__init__()
+        wpiutil.SendableRegistry.add(self, "Test", 1)
+        self.value = 0
+
+    def initSendable(self, builder: wpiutil.SendableBuilder):
+        builder.addDoubleProperty("key", self._get, self._set)
+
+    def _set(self, value: float):
+        self.value = value
+
+    def _get(self) -> float:
+        return self.value
+
+
+def test_custom_sendable():
+    ms = MySendable()
+
+    uid = wpiutil.SendableRegistry.getUniqueId(ms)
+    keys = {}
+
+    module.publish(uid, keys)
+    assert ms.value == 0
+
+    getter, setter = keys["key"]
+    assert getter() == 0
+    setter(1)
+    assert getter() == 1
+    assert ms.value == 1
diff --git a/wpiutil/src/main.cpp b/wpiutil/src/main.cpp
index 943e5e4..76fb214 100644
--- a/wpiutil/src/main.cpp
+++ b/wpiutil/src/main.cpp
@@ -4,11 +4,22 @@
 void setup_stack_trace_hook(py::object fn);
 void cleanup_stack_trace_hook();
 
+namespace wpi::impl {
+void ResetSendableRegistry();
+} // namespace wpi::impl
+
 RPYBUILD_PYBIND11_MODULE(m) {
   initWrapper(m);
 
   static int unused;
-  py::capsule cleanup(&unused, [](void *) { cleanup_stack_trace_hook(); });
+  py::capsule cleanup(&unused, [](void *) {
+    {
+      py::gil_scoped_release unlock;
+      wpi::impl::ResetSendableRegistry();
+    }
+
+    cleanup_stack_trace_hook();
+  });
 
   m.def("_setup_stack_trace_hook", &setup_stack_trace_hook);
   m.add_object("_st_cleanup", cleanup);