From 19ab5ffdc7870543cfdb57bb3173c26898df0117 Mon Sep 17 00:00:00 2001
From: Matt Wozniski <mwozniski@bloomberg.net>
Date: Mon, 3 Jun 2024 22:28:54 -0400
Subject: [PATCH] Add Python 3.13 compatibility

- Use the new `PyEval_SetProfileAllThreads` public function, instead of
  our hacky backport of it.
- Handle `_Py_IsFinalizing` being promoted to a public API.
- Handle the prototype for `_PyMem_GetCurrentAllocatorName` being
  removed.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>
---
 setup.py                             |  1 +
 src/memray/_memray/compat.cpp        | 11 +++++++----
 src/memray/_memray/compat.h          | 10 ++++++++++
 src/memray/_memray/record_writer.cpp |  6 ++++++
 src/memray/_memray/tracking_api.cpp  |  2 +-
 5 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/setup.py b/setup.py
index f556897314..e1b6b1fc1a 100644
--- a/setup.py
+++ b/setup.py
@@ -318,6 +318,7 @@ def build_js_files(self):
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
         "Programming Language :: Python :: Implementation :: CPython",
         "Topic :: Software Development :: Debuggers",
     ],
diff --git a/src/memray/_memray/compat.cpp b/src/memray/_memray/compat.cpp
index e43ef6130c..4171cebfd6 100644
--- a/src/memray/_memray/compat.cpp
+++ b/src/memray/_memray/compat.cpp
@@ -6,17 +6,19 @@ void
 setprofileAllThreads(Py_tracefunc func, PyObject* arg)
 {
     assert(PyGILState_Check());
-
+#if PY_VERSION_HEX >= 0x030D0000
+    PyEval_SetProfileAllThreads(func, arg);
+#else
     PyThreadState* this_tstate = PyThreadState_Get();
     PyInterpreterState* interp = threadStateGetInterpreter(this_tstate);
     for (PyThreadState* tstate = PyInterpreterState_ThreadHead(interp); tstate != nullptr;
          tstate = PyThreadState_Next(tstate))
     {
-#if PY_VERSION_HEX >= 0x03090000
+#    if PY_VERSION_HEX >= 0x03090000
         if (_PyEval_SetProfile(tstate, func, arg) < 0) {
             _PyErr_WriteUnraisableMsg("in PyEval_SetProfileAllThreads", nullptr);
         }
-#else
+#    else
         // For 3.7 and 3.8, backport _PyEval_SetProfile from 3.9
         // https://github.com/python/cpython/blob/v3.9.13/Python/ceval.c#L4738-L4767
         PyObject* profileobj = tstate->c_profileobj;
@@ -33,8 +35,9 @@ setprofileAllThreads(Py_tracefunc func, PyObject* arg)
 
         /* Flag that tracing or profiling is turned on */
         tstate->use_tracing = (func != NULL) || (tstate->c_tracefunc != NULL);
-#endif
+#    endif
     }
+#endif
 }
 
 }  // namespace memray::compat
diff --git a/src/memray/_memray/compat.h b/src/memray/_memray/compat.h
index 1e8e4fb626..d5e30fdfa1 100644
--- a/src/memray/_memray/compat.h
+++ b/src/memray/_memray/compat.h
@@ -7,6 +7,16 @@
 
 namespace memray::compat {
 
+inline int
+isPythonFinalizing()
+{
+#if PY_VERSION_HEX >= 0x030D0000
+    return Py_IsFinalizing();
+#else
+    return _Py_IsFinalizing();
+#endif
+}
+
 inline bool
 isEntryFrame(PyFrameObject* frame)
 {
diff --git a/src/memray/_memray/record_writer.cpp b/src/memray/_memray/record_writer.cpp
index 8d979e0255..849f18afb7 100644
--- a/src/memray/_memray/record_writer.cpp
+++ b/src/memray/_memray/record_writer.cpp
@@ -11,6 +11,12 @@
 #include "records.h"
 #include "snapshot.h"
 
+#if PY_VERSION_HEX >= 0x030D0000
+// This function still exists in 3.13 but Python.h no longer has its prototype.
+extern "C" const char*
+_PyMem_GetCurrentAllocatorName();
+#endif
+
 namespace memray::tracking_api {
 
 using namespace std::chrono;
diff --git a/src/memray/_memray/tracking_api.cpp b/src/memray/_memray/tracking_api.cpp
index fcbfcd2e0a..a6fa414e3c 100644
--- a/src/memray/_memray/tracking_api.cpp
+++ b/src/memray/_memray/tracking_api.cpp
@@ -603,7 +603,7 @@ Tracker::~Tracker()
         d_patcher.restore_symbols();
     }
 
-    if (Py_IsInitialized() && !_Py_IsFinalizing()) {
+    if (Py_IsInitialized() && !compat::isPythonFinalizing()) {
         PyGILState_STATE gstate;
         gstate = PyGILState_Ensure();