post3: Deferred torch import (duck-typing via module name)

- Remove PYTORCH_TENSOR_TYPE static and look_up_pytorch_type()
- Replace type-based pytorch detection with duck-typing
- Check __module__ starts with 'torch' then verify numpy/cpu/detach methods
- Export PyObject_HasAttrString and PyObject_GetAttrString from ffi
- No longer imports torch at startup, reducing overhead for non-torch processes

Author: yilei

src/ffi/mod.rs

@@ -74,7 +74,8 @@ pub(crate) use pyo3_ffi::{
     PyMemoryView_Type, PyMethodDef, PyMethodDefPointer, PyModule_AddIntConstant,
     PyModule_AddObject, PyModuleDef, PyModuleDef_HEAD_INIT, PyModuleDef_Init, PyModuleDef_Slot,
     PyObject, PyObject_CallFunctionObjArgs, PyObject_CallMethodObjArgs, PyObject_GenericGetDict,
-    PyObject_GetAttr, PyObject_HasAttr, PyObject_Hash, PyObject_Vectorcall, PyTuple_New,
+    PyObject_GetAttr, PyObject_GetAttrString, PyObject_HasAttr, PyObject_HasAttrString,
+    PyObject_Hash, PyObject_Vectorcall, PyTuple_New,
     PyTuple_Type, PyTupleObject, PyType_Ready, PyType_Type, PyTypeObject, PyUnicode_AsUTF8AndSize,
     PyUnicode_FromStringAndSize, PyUnicode_InternFromString, PyUnicode_New, PyUnicode_Type,
     PyVarObject, PyVectorcall_NARGS,

src/serialize/obtype.rs

@@ -1,14 +1,14 @@
 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
 // Copyright ijl (2020-2025), Aviram Hassan (2020)
 
+use crate::ffi::{PyObject_HasAttrString, PyStrRef, PyTypeObject};
 use crate::opt::{
     Opt, PASSTHROUGH_DATACLASS, PASSTHROUGH_DATETIME, PASSTHROUGH_SUBCLASS, SERIALIZE_NUMPY,
 };
 use crate::serialize::per_type::{is_numpy_array, is_numpy_scalar};
 use crate::typeref::{
     BOOL_TYPE, DATACLASS_FIELDS_STR, DATE_TYPE, DATETIME_TYPE, DICT_TYPE, ENUM_TYPE, FLOAT_TYPE,
-    FRAGMENT_TYPE, INT_TYPE, LIST_TYPE, NONE_TYPE, PYTORCH_TENSOR_TYPE, STR_TYPE, TIME_TYPE,
-    TUPLE_TYPE, UUID_TYPE,
+    FRAGMENT_TYPE, INT_TYPE, LIST_TYPE, NONE_TYPE, STR_TYPE, TIME_TYPE, TUPLE_TYPE, UUID_TYPE,
 };
 
 #[repr(u32)]
@@ -110,10 +110,40 @@ pub(crate) fn pyobject_to_obtype_unlikely(
             return ObType::NumpyScalar;
         } else if is_numpy_array(ob_type) {
             return ObType::NumpyArray;
-        } else if is_class_by_type!(ob_type, PYTORCH_TENSOR_TYPE) {
+        } else if is_pytorch_tensor(ob_type) {
             return ObType::PyTorchTensor;
         }
     }
 
     ObType::Unknown
 }
+
+#[cold]
+fn is_pytorch_tensor(ob_type: *mut PyTypeObject) -> bool {
+    unsafe {
+        // Check if the type's __module__ starts with "torch" first,
+        // to avoid calling HasAttr on types like MagicMock
+        let ob_type_ptr = ob_type.cast::<crate::ffi::PyObject>();
+        let module = crate::ffi::PyObject_GetAttrString(ob_type_ptr, c"__module__".as_ptr());
+        if module.is_null() {
+            crate::ffi::PyErr_Clear();
+            return false;
+        }
+        let starts_with_torch = match PyStrRef::from_ptr(module) {
+            Ok(s) => match s.as_str() {
+                Some(s) => s.starts_with("torch"),
+                None => false,
+            },
+            Err(_) => false,
+        };
+        ffi!(Py_DECREF(module));
+        if !starts_with_torch {
+            return false;
+        }
+
+        // Verify it has the expected tensor methods
+        PyObject_HasAttrString(ob_type_ptr, c"numpy".as_ptr()) == 1
+            && PyObject_HasAttrString(ob_type_ptr, c"cpu".as_ptr()) == 1
+            && PyObject_HasAttrString(ob_type_ptr, c"detach".as_ptr()) == 1
+    }
+}

src/typeref.rs

@@ -40,7 +40,6 @@ pub(crate) static mut FIELD_TYPE: *mut PyTypeObject = null_mut();
 pub(crate) static mut FRAGMENT_TYPE: *mut PyTypeObject = null_mut();
 
 pub(crate) static mut ZONEINFO_TYPE: *mut PyTypeObject = null_mut();
-pub(crate) static mut PYTORCH_TENSOR_TYPE: *mut PyTypeObject = null_mut();
 
 pub(crate) static mut UTCOFFSET_METHOD_STR: *mut PyObject = null_mut();
 pub(crate) static mut NORMALIZE_METHOD_STR: *mut PyObject = null_mut();
@@ -175,26 +174,6 @@ fn _init_typerefs_impl() -> bool {
     true
 }
 
-#[cold]
-#[cfg_attr(feature = "optimize", optimize(size))]
-pub(crate) fn look_up_pytorch_type() {
-    unsafe {
-        let torch = PyImport_ImportModule(c"torch".as_ptr());
-        if torch.is_null() {
-            PyErr_Clear();
-            return;
-        }
-        let torch_module_dict = PyObject_GenericGetDict(torch, null_mut());
-        let tensor_type =
-            PyMapping_GetItemString(torch_module_dict, c"Tensor".as_ptr()).cast::<PyTypeObject>();
-        Py_XDECREF(torch_module_dict);
-        Py_XDECREF(torch);
-        if !tensor_type.is_null() {
-            PYTORCH_TENSOR_TYPE = tensor_type;
-        }
-    }
-}
-
 pub(crate) struct NumpyTypes {
     pub array: *mut PyTypeObject,
     pub float64: *mut PyTypeObject,
@@ -253,7 +232,6 @@ pub(crate) fn load_numpy_types() -> Box<Option<NonNull<NumpyTypes>>> {
         });
         Py_XDECREF(numpy_module_dict);
         Py_XDECREF(numpy);
-        look_up_pytorch_type();
         Box::new(Some(nonnull!(Box::<NumpyTypes>::into_raw(types))))
     }
 }