From a9e5d84d27141dcbbbe8261948c15c6632cbba04 Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Fri, 9 Dec 2022 17:15:45 +0100
Subject: [PATCH] notify: Fix crash when using set_ostream_ptr from Python

Also allow passing in None

Fixes #1371 (and maybe also #319)
---
 dtool/src/prc/notify_ext.cxx          | 48 +++++++++++++++++++++++++++
 dtool/src/prc/notify_ext.h            | 37 +++++++++++++++++++++
 dtool/src/prc/p3prc_ext_composite.cxx |  1 +
 dtool/src/prc/pnotify.h               |  4 +++
 4 files changed, 90 insertions(+)
 create mode 100644 dtool/src/prc/notify_ext.cxx
 create mode 100644 dtool/src/prc/notify_ext.h

diff --git a/dtool/src/prc/notify_ext.cxx b/dtool/src/prc/notify_ext.cxx
new file mode 100644
index 0000000000..b3dd5a57cf
--- /dev/null
+++ b/dtool/src/prc/notify_ext.cxx
@@ -0,0 +1,48 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file notify_ext.cxx
+ * @author rdb
+ * @date 2022-12-09
+ */
+
+#include "notify_ext.h"
+
+#ifdef HAVE_PYTHON
+
+/**
+ * Changes the ostream that all subsequent Notify messages will be written to.
+ * If the previous ostream was set with delete_later = true, this will delete
+ * the previous ostream.  If ostream_ptr is None, this resets the default to
+ * cerr.
+ */
+void Extension<Notify>::
+set_ostream_ptr(PyObject *ostream_ptr, bool delete_later) {
+  extern struct Dtool_PyTypedObject Dtool_std_ostream;
+
+  if (ostream_ptr == Py_None) {
+    _this->set_ostream_ptr(nullptr, false);
+    return;
+  }
+
+  std::ostream *ptr = (std::ostream *)DTOOL_Call_GetPointerThisClass(ostream_ptr, &Dtool_std_ostream, 1, "Notify.set_ostream_ptr", false, true);
+  if (ptr == nullptr) {
+    return;
+  }
+
+  // Since we now have a reference to this class on the C++ end, make sure
+  // that the ostream isn't being destructed when its Python wrapper expires.
+  // Note that this may cause a memory leak if delete_later is not set, but
+  // since these pointers are usually set once for the rest of time, this is
+  // considered less of a problem than having the Python object destroy the
+  // object while C++ is still using it.  See GitHub #1371.
+  ((Dtool_PyInstDef *)ostream_ptr)->_memory_rules = false;
+  _this->set_ostream_ptr(ptr, delete_later);
+}
+
+#endif  // HAVE_PYTHON
diff --git a/dtool/src/prc/notify_ext.h b/dtool/src/prc/notify_ext.h
new file mode 100644
index 0000000000..ced2ac193f
--- /dev/null
+++ b/dtool/src/prc/notify_ext.h
@@ -0,0 +1,37 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file notify_ext.h
+ * @author rdb
+ * @date 2022-12-09
+ */
+
+#ifndef NOTIFY_EXT_H
+#define NOTIFY_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "pnotify.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for Notify, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<>
+class Extension<Notify> : public ExtensionBase<Notify> {
+public:
+  void set_ostream_ptr(PyObject *ostream_ptr, bool delete_later);
+};
+
+#endif  // HAVE_PYTHON
+
+#endif  // NOTIFY_EXT_H
diff --git a/dtool/src/prc/p3prc_ext_composite.cxx b/dtool/src/prc/p3prc_ext_composite.cxx
index 223ef504cd..9cec7a1e2d 100644
--- a/dtool/src/prc/p3prc_ext_composite.cxx
+++ b/dtool/src/prc/p3prc_ext_composite.cxx
@@ -1,2 +1,3 @@
+#include "notify_ext.cxx"
 #include "streamReader_ext.cxx"
 #include "streamWriter_ext.cxx"
diff --git a/dtool/src/prc/pnotify.h b/dtool/src/prc/pnotify.h
index 6fe488862e..8ae8a48a08 100644
--- a/dtool/src/prc/pnotify.h
+++ b/dtool/src/prc/pnotify.h
@@ -35,7 +35,11 @@ PUBLISHED:
   Notify();
   ~Notify();
 
+#if defined(CPPPARSER) && defined(HAVE_PYTHON)
+  EXTEND void set_ostream_ptr(PyObject *ostream_ptr, bool delete_later);
+#else
   void set_ostream_ptr(std::ostream *ostream_ptr, bool delete_later);
+#endif
   std::ostream *get_ostream_ptr() const;
 
   typedef bool AssertHandler(const char *expression, int line,