fix: support CPython free-threaded builds

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

src/C/consts.jl

@@ -124,6 +124,20 @@ end
     type::Ptr{Cvoid} = C_NULL # really is Ptr{PyObject} or Ptr{PyTypeObject} but Julia 1.3 and below get the layout incorrect when circular types are involved
 end
 
+@kwdef struct PyMutex
+    bits::Cuchar = 0
+end
+
+@kwdef struct PyObjectFT
+    tid::Csize_t = 0
+    flags::Cushort = 0
+    mutex::PyMutex = PyMutex()
+    gc_bits::Cuchar = 0
+    ref_local::Cuint = 0
+    ref_shared::Py_ssize_t = 0
+    type::Ptr{Cvoid} = C_NULL # really is Ptr{PyObject} or Ptr{PyTypeObject} but Julia 1.3 and below get the layout incorrect when circular types are involved
+end
+
 const PyPtr = Ptr{PyObject}
 const PyNULL = PyPtr(0)
 
@@ -139,6 +153,11 @@ Base.unsafe_convert(::Type{PyPtr}, o::PyObjectRef) = o.ptr
     size::Py_ssize_t = 0
 end
 
+@kwdef struct PyVarObjectFT
+    ob_base::PyObjectFT = PyObjectFT()
+    size::Py_ssize_t = 0
+end
+
 @kwdef struct PyMethodDef
     name::Cstring = C_NULL
     meth::Ptr{Cvoid} = C_NULL
@@ -249,6 +268,16 @@ end
     weakreflist::PyPtr = PyNULL
 end
 
+@kwdef struct PyMemoryViewObjectFT
+    ob_base::PyVarObjectFT = PyVarObjectFT()
+    mbuf::PyPtr = PyNULL
+    hash::Py_hash_t = 0
+    flags::Cint = 0
+    exports::Py_ssize_t = 0
+    view::Py_buffer = Py_buffer()
+    weakreflist::PyPtr = PyNULL
+end
+
 @kwdef struct PyTypeObject
     ob_base::PyVarObject = PyVarObject()
     name::Cstring = C_NULL
@@ -327,6 +356,11 @@ end
     value::T
 end
 
+@kwdef struct PySimpleObjectFT{T}
+    ob_base::PyObjectFT = PyObjectFT()
+    value::T
+end
+
 @kwdef struct PyArrayInterface
     two::Cint = 0
     nd::Cint = 0

src/C/context.jl

@@ -17,6 +17,7 @@ A handle to a loaded instance of libpython, its interpreter, function pointers,
     pyhome_w::Any = missing
     which::Symbol = :unknown # :CondaPkg, :PyCall, :embedded or :unknown
     version::Union{VersionNumber,Missing} = missing
+    is_free_threaded::Bool = false
 end
 
 const CTX = Context()
@@ -312,10 +313,11 @@ function init_context()
     v"3.10" ≤ CTX.version < v"4" || error(
         "Only Python 3.10+ is supported, this is Python $(CTX.version) at $(CTX.exe_path===missing ? "unknown location" : CTX.exe_path).",
     )
+    CTX.is_free_threaded = occursin("free-threading build", verstr)
 
     launch_on_main_thread(Threads.threadid()) # makes on_main_thread usable
 
-    @debug "Initialized PythonCall.jl" CTX.is_embedded CTX.is_initialized CTX.exe_path CTX.lib_path CTX.lib_ptr CTX.pyprogname CTX.pyhome CTX.version
+    @debug "Initialized PythonCall.jl" CTX.is_embedded CTX.is_initialized CTX.exe_path CTX.lib_path CTX.lib_ptr CTX.pyprogname CTX.pyhome CTX.version CTX.is_free_threaded
 
     return
 end

src/C/extras.jl

@@ -1,44 +1,58 @@
 asptr(x) = Base.unsafe_convert(PyPtr, x)
 
-Py_Type(x) = Base.GC.@preserve x PyPtr(UnsafePtr(asptr(x)).type[!])
+Py_Type(x) = Base.GC.@preserve x begin
+    if CTX.is_free_threaded
+        PyPtr(UnsafePtr{PyObjectFT}(asptr(x)).type[!])
+    else
+        PyPtr(UnsafePtr{PyObject}(asptr(x)).type[!])
+    end
+end
 
 PyObject_Type(x) = Base.GC.@preserve x (t = Py_Type(asptr(x)); Py_IncRef(t); t)
 
 Py_TypeCheck(o, t) = Base.GC.@preserve o t PyType_IsSubtype(Py_Type(asptr(o)), asptr(t))
 Py_TypeCheckFast(o, f::Integer) = Base.GC.@preserve o PyType_IsSubtypeFast(Py_Type(asptr(o)), f)
 
 PyType_IsSubtypeFast(t, f::Integer) =
-    Base.GC.@preserve t Cint(!iszero(UnsafePtr{PyTypeObject}(asptr(t)).flags[] & f))
+    Base.GC.@preserve t Cint(!iszero(PyType_GetFlags(asptr(t)) & f))
 
-PyMemoryView_GET_BUFFER(m) = Base.GC.@preserve m Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObject}(asptr(m)).view)
+PyMemoryView_GET_BUFFER(m) = Base.GC.@preserve m begin
+    if CTX.is_free_threaded
+        Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObjectFT}(asptr(m)).view)
+    else
+        Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObject}(asptr(m)).view)
+    end
+end
 
 PyType_CheckBuffer(t) = Base.GC.@preserve t begin
-    p = UnsafePtr{PyTypeObject}(asptr(t)).as_buffer[]
-    return p != C_NULL && p.get[!] != C_NULL
+    getbuf = PyType_GetSlot(asptr(t), Py_bf_getbuffer)
+    return getbuf != C_NULL
 end
 
 PyObject_CheckBuffer(o) = Base.GC.@preserve o PyType_CheckBuffer(Py_Type(asptr(o)))
 
 PyObject_GetBuffer(_o, b, flags) = Base.GC.@preserve _o begin
     o = asptr(_o)
-    p = UnsafePtr{PyTypeObject}(Py_Type(o)).as_buffer[]
-    if p == C_NULL || p.get[!] == C_NULL
-        PyErr_SetString(
-            POINTERS.PyExc_TypeError,
-            "a bytes-like object is required, not '$(String(UnsafePtr{PyTypeObject}(Py_Type(o)).name[]))'",
-        )
+    getbuf = PyType_GetSlot(Py_Type(o), Py_bf_getbuffer)
+    if getbuf == C_NULL
+        msg = if CTX.is_free_threaded
+            "a bytes-like object is required"
+        else
+            "a bytes-like object is required, not '$(String(UnsafePtr{PyTypeObject}(Py_Type(o)).name[]))'"
+        end
+        PyErr_SetString(POINTERS.PyExc_TypeError, msg)
         return Cint(-1)
     end
-    return ccall(p.get[!], Cint, (PyPtr, Ptr{Py_buffer}, Cint), o, b, flags)
+    return ccall(getbuf, Cint, (PyPtr, Ptr{Py_buffer}, Cint), o, b, flags)
 end
 
 PyBuffer_Release(_b) = begin
     b = UnsafePtr(Base.unsafe_convert(Ptr{Py_buffer}, _b))
     o = b.obj[]
     o == C_NULL && return
-    p = UnsafePtr{PyTypeObject}(Py_Type(o)).as_buffer[]
-    if (p != C_NULL && p.release[!] != C_NULL)
-        ccall(p.release[!], Cvoid, (PyPtr, Ptr{Py_buffer}), o, b)
+    releasebuf = PyType_GetSlot(Py_Type(o), Py_bf_releasebuffer)
+    if releasebuf != C_NULL
+        ccall(releasebuf, Cvoid, (PyPtr, Ptr{Py_buffer}), o, b)
     end
     b.obj[] = C_NULL
     Py_DecRef(o)
@@ -65,7 +79,13 @@ function PyOS_RunInputHook()
 end
 
 function PySimpleObject_GetValue(::Type{T}, o) where {T}
-    Base.GC.@preserve o UnsafePtr{PySimpleObject{T}}(asptr(o)).value[!]
+    Base.GC.@preserve o begin
+        if CTX.is_free_threaded
+            UnsafePtr{PySimpleObjectFT{T}}(asptr(o)).value[!]
+        else
+            UnsafePtr{PySimpleObject{T}}(asptr(o)).value[!]
+        end
+    end
 end
 
 # FAST REFCOUNTING

src/C/pointers.jl

@@ -77,6 +77,8 @@ const CAPI_FUNC_SIGS = Dict{Symbol,Pair{Tuple,Type}}(
     :PyType_Ready => (PyPtr,) => Cint,
     :PyType_GenericNew => (PyPtr, PyPtr, PyPtr) => PyPtr,
     :PyType_FromSpec => (Ptr{Cvoid},) => PyPtr,
+    :PyType_GetFlags => (PyPtr,) => Culong,
+    :PyType_GetSlot => (PyPtr, Cint) => Ptr{Cvoid},
     # MAPPING
     :PyMapping_HasKeyString => (PyPtr, Ptr{Cchar}) => Cint,
     :PyMapping_SetItemString => (PyPtr, Ptr{Cchar}, PyPtr) => Cint,

src/JlWrap/C.jl

@@ -12,6 +12,12 @@ using Serialization: serialize, deserialize
     weaklist::C.PyPtr = C_NULL
 end
 
+@kwdef struct PyJuliaValueObjectFT
+    ob_base::C.PyObjectFT = C.PyObjectFT()
+    value::Int = 0
+    weaklist::C.PyPtr = C_NULL
+end
+
 const PyJuliaBase_Type = Ref(C.PyNULL)
 
 # we store the actual julia values here
@@ -20,22 +26,37 @@ const PYJLVALUES = []
 # unused indices in PYJLVALUES
 const PYJLFREEVALUES = Int[]
 
+# Choose the layout based on whether Python is free-threaded.
+pyjl_obj_type() = C.CTX.is_free_threaded ? PyJuliaValueObjectFT : PyJuliaValueObject
+
+function pyjl_obj_ptr(o::C.PyPtr)
+    if C.CTX.is_free_threaded
+        return UnsafePtr{PyJuliaValueObjectFT}(o)
+    else
+        return UnsafePtr{PyJuliaValueObject}(o)
+    end
+end
+
 function _pyjl_new(t::C.PyPtr, ::C.PyPtr, ::C.PyPtr)
-    o = ccall(UnsafePtr{C.PyTypeObject}(t).alloc[!], C.PyPtr, (C.PyPtr, C.Py_ssize_t), t, 0)
+    alloc = C.PyType_GetSlot(t, C.Py_tp_alloc)
+    alloc == C_NULL && return C.PyNULL
+    o = ccall(alloc, C.PyPtr, (C.PyPtr, C.Py_ssize_t), t, 0)
     o == C.PyNULL && return C.PyNULL
-    UnsafePtr{PyJuliaValueObject}(o).weaklist[] = C.PyNULL
-    UnsafePtr{PyJuliaValueObject}(o).value[] = 0
+    p = pyjl_obj_ptr(o)
+    p.weaklist[] = C.PyNULL
+    p.value[] = 0
     return o
 end
 
 function _pyjl_dealloc(o::C.PyPtr)
-    idx = UnsafePtr{PyJuliaValueObject}(o).value[]
+    idx = pyjl_obj_ptr(o).value[]
     if idx != 0
         PYJLVALUES[idx] = nothing
         push!(PYJLFREEVALUES, idx)
     end
-    UnsafePtr{PyJuliaValueObject}(o).weaklist[!] == C.PyNULL || C.PyObject_ClearWeakRefs(o)
-    ccall(UnsafePtr{C.PyTypeObject}(C.Py_Type(o)).free[!], Cvoid, (C.PyPtr,), o)
+    pyjl_obj_ptr(o).weaklist[!] == C.PyNULL || C.PyObject_ClearWeakRefs(o)
+    freeptr = C.PyType_GetSlot(C.Py_Type(o), C.Py_tp_free)
+    freeptr == C_NULL || ccall(freeptr, Cvoid, (C.PyPtr,), o)
     nothing
 end
 
@@ -314,12 +335,13 @@ function init_c()
 
     # Create members for weakref support
     empty!(_pyjlbase_members)
+    objT = pyjl_obj_type()
     push!(
         _pyjlbase_members,
         C.PyMemberDef(
             name = pointer(_pyjlbase_weaklistoffset_name),
             typ = C.Py_T_PYSSIZET,
-            offset = fieldoffset(PyJuliaValueObject, 3),
+            offset = fieldoffset(objT, 3),
             flags = C.Py_READONLY,
         ),
         C.PyMemberDef(), # NULL terminator
@@ -341,7 +363,7 @@ function init_c()
     # Create PyType_Spec
     _pyjlbase_spec[] = C.PyType_Spec(
         name = pointer(_pyjlbase_name),
-        basicsize = sizeof(PyJuliaValueObject),
+        basicsize = sizeof(objT),
         flags = C.Py_TPFLAGS_BASETYPE | C.Py_TPFLAGS_HAVE_VERSION_TAG,
         slots = pointer(_pyjlbase_slots),
     )
@@ -358,13 +380,14 @@ function __init__()
     init_c()
 end
 
-PyJuliaValue_IsNull(o) = Base.GC.@preserve o UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[] == 0
+PyJuliaValue_IsNull(o) = Base.GC.@preserve o pyjl_obj_ptr(C.asptr(o)).value[] == 0
 
-PyJuliaValue_GetValue(o) = Base.GC.@preserve o PYJLVALUES[UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[]]
+PyJuliaValue_GetValue(o) = Base.GC.@preserve o PYJLVALUES[pyjl_obj_ptr(C.asptr(o)).value[]]
 
 PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
     o = C.asptr(_o)
-    idx = UnsafePtr{PyJuliaValueObject}(o).value[]
+    p = pyjl_obj_ptr(o)
+    idx = p.value[]
     if idx == 0
         if isempty(PYJLFREEVALUES)
             push!(PYJLVALUES, v)
@@ -373,7 +396,7 @@ PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
             idx = pop!(PYJLFREEVALUES)
             PYJLVALUES[idx] = v
         end
-        UnsafePtr{PyJuliaValueObject}(o).value[] = idx
+        p.value[] = idx
     else
         PYJLVALUES[idx] = v
     end