Merge #2904 #2921

2904: refactor docstring generation code r=mejrs a=davidhewitt

As a first step towards #2866 this is a tweak to the macro code which generates Python docstrings. This PR refactors the behaviour so that instead of always creating a `concat!` expression to generate a nul-terminated string, this will only happen if a Rust 1.54+ macro doc is present (e.g. `#[doc = include_str!(...)]`).

The default case now just collects all the `#[doc]` attributes into a single string.

This should make it easier to factor out the `text_signature` formatting, and avoids wasting compile time invoking the `concat!` macro when not necessary.

2921: Check to see if object is `None` before traversing r=davidhewitt a=neachdainn

Closes #2915

When using the C API directly, the intended way to call `visitproc` is via the `Py_VISIT` macro, which checks to see that the provided pointer is not null before passing it along to `visitproc`. Because PyO3 isn't using the macro, it needs to manually check that the pointer isn't null. Without this check, calling `visit.call(&obj)` where `let obj = None;` will segfault.


Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
Co-authored-by: Nate Kent <nate@nkent.net>

newsfragments/2921.fixed.md

@@ -0,0 +1 @@
+Traversal visit calls to `Option<T: AsPyPointer>` no longer segfaults when `None`.

pyo3-macros-backend/src/utils.rs

@@ -1,9 +1,9 @@
-use std::borrow::Cow;
+use std::{borrow::Cow, fmt::Write};
 
 // Copyright (c) 2017-present PyO3 Project and Contributors
 use proc_macro2::{Span, TokenStream};
 use quote::ToTokens;
-use syn::{spanned::Spanned, Ident};
+use syn::{punctuated::Punctuated, spanned::Spanned, Ident, Token};
 
 use crate::attributes::CrateAttribute;
 
@@ -57,70 +57,85 @@ pub fn option_type_argument(ty: &syn::Type) -> Option<&syn::Type> {
     None
 }
 
-/// A syntax tree which evaluates to a null-terminated docstring for Python.
+/// A syntax tree which evaluates to a nul-terminated docstring for Python.
 ///
-/// It's built as a `concat!` evaluation, so it's hard to do anything with this
+/// Typically the tokens will just be that string, but if the original docs included macro
+/// expressions then the tokens will be a concat!("...", "\n", "\0") expression of the strings and
+/// macro parts.
 /// contents such as parse the string contents.
 #[derive(Clone)]
 pub struct PythonDoc(TokenStream);
 
-/// Collects all #[doc = "..."] attributes into a TokenStream evaluating to a null-terminated string
-/// e.g. concat!("...", "\n", "\0")
+/// Collects all #[doc = "..."] attributes into a TokenStream evaluating to a null-terminated string.
 pub fn get_doc(
     attrs: &[syn::Attribute],
     text_signature: Option<(Cow<'_, Ident>, String)>,
 ) -> PythonDoc {
-    let mut tokens = TokenStream::new();
-    let comma = syn::token::Comma(Span::call_site());
-    let newline = syn::LitStr::new("\n", Span::call_site());
-
-    syn::Ident::new("concat", Span::call_site()).to_tokens(&mut tokens);
-    syn::token::Bang(Span::call_site()).to_tokens(&mut tokens);
-    syn::token::Bracket(Span::call_site()).surround(&mut tokens, |tokens| {
-        if let Some((python_name, text_signature)) = text_signature {
-            // create special doc string lines to set `__text_signature__`
-            let signature_lines = format!("{}{}\n--\n\n", python_name, text_signature);
-            signature_lines.to_tokens(tokens);
-            comma.to_tokens(tokens);
-        }
+    let mut parts = Punctuated::<TokenStream, Token![,]>::new();
+    let mut current_part = String::new();
+
+    if let Some((python_name, text_signature)) = text_signature {
+        // create special doc string lines to set `__text_signature__`
+        write!(
+            &mut current_part,
+            "{}{}\n--\n\n",
+            python_name, text_signature
+        )
+        .expect("error occurred while trying to format text_signature to string")
+    }
 
-        let mut first = true;
-
-        for attr in attrs.iter() {
-            if attr.path.is_ident("doc") {
-                if let Ok(DocArgs {
-                    _eq_token,
-                    token_stream,
-                }) = syn::parse2(attr.tokens.clone())
-                {
-                    if !first {
-                        newline.to_tokens(tokens);
-                        comma.to_tokens(tokens);
-                    } else {
-                        first = false;
-                    }
-                    if let Ok(syn::Lit::Str(lit_str)) = syn::parse2(token_stream.clone()) {
-                        // Strip single left space from literal strings, if needed.
-                        // e.g. `/// Hello world` expands to #[doc = " Hello world"]
-                        let doc_line = lit_str.value();
-                        doc_line
-                            .strip_prefix(' ')
-                            .map(|stripped| syn::LitStr::new(stripped, lit_str.span()))
-                            .unwrap_or(lit_str)
-                            .to_tokens(tokens);
-                    } else {
-                        // This is probably a macro doc from Rust 1.54, e.g. #[doc = include_str!(...)]
-                        token_stream.to_tokens(tokens)
-                    }
-                    comma.to_tokens(tokens);
+    let mut first = true;
+
+    for attr in attrs.iter() {
+        if attr.path.is_ident("doc") {
+            if let Ok(DocArgs {
+                _eq_token,
+                token_stream,
+            }) = syn::parse2(attr.tokens.clone())
+            {
+                if !first {
+                    current_part.push('\n');
+                } else {
+                    first = false;
+                }
+                if let Ok(syn::Lit::Str(lit_str)) = syn::parse2(token_stream.clone()) {
+                    // Strip single left space from literal strings, if needed.
+                    // e.g. `/// Hello world` expands to #[doc = " Hello world"]
+                    let doc_line = lit_str.value();
+                    current_part.push_str(doc_line.strip_prefix(' ').unwrap_or(&doc_line));
+                } else {
+                    // This is probably a macro doc from Rust 1.54, e.g. #[doc = include_str!(...)]
+                    // Reset the string buffer, write that part, and then push this macro part too.
+                    parts.push(current_part.to_token_stream());
+                    current_part.clear();
+                    parts.push(token_stream);
                 }
             }
         }
+    }
 
-        syn::LitStr::new("\0", Span::call_site()).to_tokens(tokens);
-    });
+    if !parts.is_empty() {
+        // Doc contained macro pieces - return as `concat!` expression
+        if !current_part.is_empty() {
+            parts.push(current_part.to_token_stream());
+        }
 
-    PythonDoc(tokens)
+        let mut tokens = TokenStream::new();
+
+        syn::Ident::new("concat", Span::call_site()).to_tokens(&mut tokens);
+        syn::token::Bang(Span::call_site()).to_tokens(&mut tokens);
+        syn::token::Bracket(Span::call_site()).surround(&mut tokens, |tokens| {
+            parts.to_tokens(tokens);
+            syn::token::Comma(Span::call_site()).to_tokens(tokens);
+            syn::LitStr::new("\0", Span::call_site()).to_tokens(tokens);
+        });
+
+        PythonDoc(tokens)
+    } else {
+        // Just a string doc - return directly with nul terminator
+        current_part.push('\0');
+        PythonDoc(current_part.to_token_stream())
+    }
 }
 
 impl quote::ToTokens for PythonDoc {

src/pyclass/gc.rs

@@ -23,11 +23,16 @@ impl<'p> PyVisit<'p> {
     where
         T: AsPyPointer,
     {
-        let r = unsafe { (self.visit)(obj.as_ptr(), self.arg) };
-        if r == 0 {
-            Ok(())
+        let ptr = obj.as_ptr();
+        if !ptr.is_null() {
+            let r = unsafe { (self.visit)(ptr, self.arg) };
+            if r == 0 {
+                Ok(())
+            } else {
+                Err(PyTraverseError(r))
+            }
         } else {
-            Err(PyTraverseError(r))
+            Ok(())
         }
     }
 

tests/test_gc.rs

@@ -120,6 +120,44 @@ fn gc_integration() {
     });
 }
 
+#[pyclass]
+struct GcNullTraversal {
+    cycle: Option<Py<Self>>,
+    null: Option<Py<Self>>,
+}
+
+#[pymethods]
+impl GcNullTraversal {
+    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.cycle)?;
+        visit.call(&self.null)?; // Should not segfault
+        Ok(())
+    }
+
+    fn __clear__(&mut self) {
+        self.cycle = None;
+        self.null = None;
+    }
+}
+
+#[test]
+fn gc_null_traversal() {
+    Python::with_gil(|py| {
+        let obj = Py::new(
+            py,
+            GcNullTraversal {
+                cycle: None,
+                null: None,
+            },
+        )
+        .unwrap();
+        obj.borrow_mut(py).cycle = Some(obj.clone_ref(py));
+
+        // the object doesn't have to be cleaned up, it just needs to be traversed.
+        py.run("import gc; gc.collect()", None, None).unwrap();
+    });
+}
+
 #[pyclass(subclass)]
 struct BaseClassWithDrop {
     data: Option<Arc<AtomicBool>>,