diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b32d8f23..77c86448 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -38,7 +38,7 @@ repos: hooks: - id: cpplint - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.0 + rev: v0.16.1 hooks: - id: ruff-check args: [--fix, --exit-non-zero-on-fix] diff --git a/CHANGELOG.md b/CHANGELOG.md index 76b6bb6b..45553613 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Define `Py_GIL_DISABLED` for free-threaded debug builds on Windows when building the C extension to work around an upstream CMake `FindPython` bug by [@XuehaiPan](https://github.com/XuehaiPan) in [#285](https://github.com/metaopt/optree/pull/285). +- Fix a deadlock when registering or unregistering a pytree node concurrently with flattening, caused by releasing the GIL while holding the registry lock by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a deadlock when a duplicate registration or an unregistration of an unregistered type formats its error message, which runs the type's `__repr__` while the registry lock is held and wedges any thread that is flattening by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix registering an already-registered `collections.namedtuple` subclass or `PyStructSequence` emitting the override warning even though the registration is rejected, which raised `UserWarning` instead of `ValueError` under warnings-as-errors by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a deadlock when unregistering a pytree node drops the last reference to its flatten or unflatten function, so a `__del__` or weakref callback re-entered the registry while its lock was held by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a deadlock while flattening on free-threading builds, caused by acquiring the non-recursive dictionary insertion order lock twice by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a spurious `SystemError` from `optree._C.get_registry_size()` when a concurrent registration slipped between its two reads by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `optree.tree_iter()` hanging uninterruptibly when the `is_leaf` predicate or a custom flatten function advances the same iterator, now reported as a `RuntimeError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the type caches handing an interpreter a value owned by another one, by keying them on the interpreter in addition to the type address by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the type caches retaining entries for a finalized interpreter, which leaked immortal keys and could be inherited by a fresh interpreter reusing the same ID by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a deadlock in the `PyStructSequence` field cache, caused by re-acquiring the GIL while still holding the cache lock by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the type caches publishing an entry that a failed cleanup registration would orphan, leaving a value owned by an interpreter with no callback able to evict it, whether the publish preceded that registration or raced another thread still performing it by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `optree.structseq_fields()` and the struct sequence accessors reporting the trailing hidden field names for unnamed sequence slots, by mapping each field to the slot its member offset names; the pure Python fallback, which cannot read those offsets, recovers the positions from a probe instance instead of assuming the unnamed slots trail the named ones by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the treespec repr printing the raw unnamed-field marker as if it were a keyword argument, now rendered as `` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `StructSequenceEntry` reporting an unnamed sequence slot as `field='unnamed field'`, which reads as a real field name, now reported by its index by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeSpec.__setstate__()` accepting a malformed state, which could read out of bounds or abort the interpreter when the treespec was later used by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeSpec.__getstate__()` returning the treespec's internal mutable containers, so mutating the pickled state corrupted an otherwise immutable treespec by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeSpec.__setstate__()` borrowing the key lists from the state it was given rather than copying them, so mutating that state afterwards silently reordered the restored treespec by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix pickling a `PyTreeSpec` with protocol 0 or 1 aborting the interpreter, by reducing through `copyreg.__newobj__` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeSpec.compose()`, `PyTreeSpec.broadcast_to_common_suffix()`, `treespec_transform()`, and `treespec_from_collection()` silently rebinding a custom node to a different registration when an empty namespace adopted a non-empty one by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeSpec.broadcast_to_common_suffix()` sorting the argument treespec's dictionary keys in place while building its key-mismatch error message, corrupting a treespec the caller still holds by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the dictionary key order being lost when a treespec built under the global namespace is promoted to an insertion-ordered namespace by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `treespec_from_collection()` on a leaf reporting an escalated `UserWarning` as a confusing `SystemError`, by checking the return value of `PyErr_WarnEx()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `treespec_from_collection()` keeping the caller's namespace on a leaf or `None` root, which made otherwise identical treespecs compare unequal by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `treespec_is_prefix()` and `treespec_is_suffix()` comparing against a stale subtree when a dictionary node's keys had been reordered by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the tree iterator never reporting or clearing its `is_leaf` predicate to the garbage collector, leaking any reference cycle that passes through it by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a list shrunk by an `is_leaf` predicate part way through flattening being read out of bounds on Python versions before 3.13.0a4, now raising `IndexError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a reference cycle passing through a registered custom type not being collectable once the registry no longer holds the registration, by reporting the registration's members to the garbage collector when a treespec's own nodes hold every reference to it; a cycle spanning several treespecs remains uncollectable by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a deeply nested treespec overflowing the native stack and crashing in `treespec_paths()`, `treespec_accessors()`, and `PyTreeSpec.broadcast_to_common_suffix()` instead of raising `RecursionError` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `DataclassEntry` and `AttrsEntry` resolving an integer entry against every declared field rather than the children the registration emitted, which returned the wrong attribute for a class holding a metadata or non-`init` field, now read from the record `register_node()` leaves on the class itself by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `optree.dataclasses.register_node()` and `optree.integrations.attrs.register_node()` leaving a class in a half-registered state that could never be registered again: the registry entry and the field marker are now committed together, and a failure of either rolls the other back by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `optree.dataclasses.register_node()` silently dropping `InitVar` pseudo-fields, which are neither children nor metadata and cannot round-trip, now rejected with a pointer to the generic API by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `tree_broadcast_common()` and `broadcast_common()` applying the caller's `is_leaf` predicate to an internal sentinel tree, which could raise from the predicate or collapse a filled subtree and under-replicate by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeEntry` equality and hashing comparing the bytecode of `__call__()` and `codify()` rather than the methods themselves, so two entry classes that happened to share an implementation compared equal by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `PyTreeAccessor` overriding `__eq__()` while inheriting `tuple.__ne__()`, so comparing one with a plain tuple of equal entries reported `False` for both `==` and `!=` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `GetAttrEntry.codify()` emitting invalid attribute access for a field name that is not an identifier, now rendered as a `getattr()` call by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `prefix_errors()` raising `AssertionError` instead of reporting an error for a custom node whose per-instance entries differ while its metadata does not, a pair `broadcast_prefix()` accepts by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `prefix_errors()` raising `TypeError` while formatting a dictionary key mismatch when the keys have different types, now ordered with `total_order_sorted()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `tree_transpose_map()`, `tree_transpose_map_with_path()`, `tree_transpose_map_with_accessor()`, and `tree_partition()` rejecting a leafless outer structure when an explicit `inner_treespec` leaves nothing to infer by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix the `tree_broadcast_map()` family flattening a custom node twice when called with a single input tree, which broke a one-shot flatten function and contradicted the documented equivalence with `tree_map()` by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `treespec_entry()` and `treespec_child()` annotating their index as `int` while the runtime accepts any `SupportsInt` or `SupportsIndex` object by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `optree.utils.total_order_sorted()` mistaking a `TypeError` raised by the caller's `key` callback for a comparison failure and silently returning the sequence unsorted, and calling the callback twice per element by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix a dictionary whose keys cannot be compared being flattened in a partially sorted order instead of insertion order, when a comparison raised part way through the sort by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). +- Fix `tree_broadcast_common()` documenting that its results share one structure and `broadcast_common()` documenting that it returns two pytrees rather than two lists of leaves by [@XuehaiPan](https://github.com/XuehaiPan) in [#290](https://github.com/metaopt/optree/pull/290). ### Removed diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index 3687fbbd..c7f534ce 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -2,6 +2,7 @@ abc ABCMeta accessor accessors +antisymmetric arg args arity @@ -70,7 +71,9 @@ optree OrderedDict param params +picklable pragma +preorder py pypy pytree diff --git a/include/optree/hashing.h b/include/optree/hashing.h index 862612ee..405b4164 100644 --- a/include/optree/hashing.h +++ b/include/optree/hashing.h @@ -22,9 +22,11 @@ limitations under the License. #include // std::string #include // std::pair +#include + #include -#include "optree/pymacros.h" // Py_ALWAYS_INLINE +#include "optree/pymacros.h" // Py_ALWAYS_INLINE, interpid_t namespace py = pybind11; @@ -89,6 +91,26 @@ struct std::not_equal_to> { std::not_equal_to{}(lhs.second, rhs.second); } }; +template <> +struct std::equal_to> { + using is_transparent = void; + inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair &lhs, + const std::pair &rhs) + const noexcept(noexcept(std::equal_to{}(lhs.first, rhs.first))) { + return std::equal_to{}(lhs.first, rhs.first) && + std::equal_to{}(lhs.second, rhs.second); + } +}; +template <> +struct std::not_equal_to> { + using is_transparent = void; + inline constexpr Py_ALWAYS_INLINE bool operator()(const std::pair &lhs, + const std::pair &rhs) + const noexcept(noexcept(std::not_equal_to{}(lhs.first, rhs.first))) { + return std::not_equal_to{}(lhs.first, rhs.first) || + std::not_equal_to{}(lhs.second, rhs.second); + } +}; template struct std::hash> { using is_transparent = void; diff --git a/include/optree/pymacros.h b/include/optree/pymacros.h index e0469047..2162338d 100644 --- a/include/optree/pymacros.h +++ b/include/optree/pymacros.h @@ -71,6 +71,19 @@ inline constexpr Py_ALWAYS_INLINE bool Py_IsConstant(PyObject *x) noexcept { } #define Py_IsConstant(x) Py_IsConstant(x) +// `PyStructSequence_UnnamedField` is declared `extern` with hidden visibility (so it is not an +// exported dynamic symbol for extension modules) before Python 3.11.0a2, where it became +// `PyAPI_DATA`. Referencing it directly leaves an undefined symbol that makes the module fail to +// import on those versions. Its value is the stable marker "unnamed field", and callers only ever +// use it by value (never by pointer identity), so fall back to that literal there. +inline const char *PyStructSequenceUnnamedField() noexcept { +#if PY_VERSION_HEX >= 0x030B00A2 // Python 3.11.0a2 + return PyStructSequence_UnnamedField; +#else + return "unnamed field"; +#endif +} + using interpid_t = decltype(PyInterpreterState_GetID(nullptr)); #if defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT) && \ @@ -81,15 +94,17 @@ using interpid_t = decltype(PyInterpreterState_GetID(nullptr)); } [[nodiscard]] inline interpid_t GetCurrentPyInterpreterID() { + // This sits on the flatten path (twice per leaf, via the `Is*Instance` caches), so it tests the + // documented failure returns rather than probing the global error state. PyInterpreterState *interp = PyInterpreterState_Get(); - if (PyErr_Occurred() != nullptr) [[unlikely]] { - throw py::error_already_set(); - } if (interp == nullptr) [[unlikely]] { + if (PyErr_Occurred() != nullptr) [[unlikely]] { + throw py::error_already_set(); + } throw std::runtime_error("Failed to get the current Python interpreter state."); } const interpid_t interpid = PyInterpreterState_GetID(interp); - if (PyErr_Occurred() != nullptr) [[unlikely]] { + if (interpid < 0) [[unlikely]] { throw py::error_already_set(); } return interpid; diff --git a/include/optree/pytypes.h b/include/optree/pytypes.h index a7cac3ea..3ef2483b 100644 --- a/include/optree/pytypes.h +++ b/include/optree/pytypes.h @@ -17,11 +17,14 @@ limitations under the License. #pragma once +#include // std::size_t, offsetof #include // std::rethrow_exception, std::current_exception +#include // std::optional #include // std::string -#include // std::enable_if_t, std::is_base_of_v +#include // std::enable_if_t, std::is_same_v, std::is_base_of_v, std::conditional_t #include // std::unordered_map -#include // std::move, std::pair, std::make_pair +#include // std::forward, std::pair, std::make_pair, std::move +#include // std::vector #include @@ -38,20 +41,19 @@ limitations under the License. namespace py = pybind11; -inline Py_ALWAYS_INLINE std::string PyStr(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE std::string PyStr(const py::handle &object) { return EVALUATE_WITH_LOCK_HELD(static_cast(py::str(object)), object); } -inline Py_ALWAYS_INLINE std::string PyStr(const std::string &string) { return string; } -inline Py_ALWAYS_INLINE std::string PyRepr(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE std::string PyStr(const std::string &string) { + return string; +} +[[nodiscard]] inline Py_ALWAYS_INLINE std::string PyRepr(const py::handle &object) { return EVALUATE_WITH_LOCK_HELD(static_cast(py::repr(object)), object); } -inline Py_ALWAYS_INLINE std::string PyRepr(const std::string &string) { +[[nodiscard]] inline Py_ALWAYS_INLINE std::string PyRepr(const std::string &string) { return static_cast(py::repr(py::str(string))); } -// The maximum size of the type cache. -constexpr py::ssize_t MAX_TYPE_CACHE_SIZE = 4096; - #define PyNoneTypeObject \ (py::reinterpret_borrow(reinterpret_cast(Py_TYPE(Py_None)))) #define PyTupleTypeObject \ @@ -67,7 +69,7 @@ constexpr py::ssize_t MAX_TYPE_CACHE_SIZE = 4096; #define PyDefaultDict_Type (reinterpret_cast(PyDefaultDictTypeObject.ptr())) #define PyDeque_Type (reinterpret_cast(PyDequeTypeObject.ptr())) -inline const py::object &ImportOrderedDict() { +[[nodiscard]] inline const py::object &ImportOrderedDict() { PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; return storage .call_once_and_store_result([]() -> py::object { @@ -75,7 +77,7 @@ inline const py::object &ImportOrderedDict() { }) .get_stored(); } -inline const py::object &ImportDefaultDict() { +[[nodiscard]] inline const py::object &ImportDefaultDict() { PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; return storage .call_once_and_store_result([]() -> py::object { @@ -83,7 +85,7 @@ inline const py::object &ImportDefaultDict() { }) .get_stored(); } -inline const py::object &ImportDeque() { +[[nodiscard]] inline const py::object &ImportDeque() { PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; return storage .call_once_and_store_result( @@ -91,13 +93,13 @@ inline const py::object &ImportDeque() { .get_stored(); } -inline Py_ALWAYS_INLINE py::ssize_t TupleGetSize(const py::handle &tuple) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::ssize_t TupleGetSize(const py::handle &tuple) { return PyTuple_GET_SIZE(tuple.ptr()); } -inline Py_ALWAYS_INLINE py::ssize_t ListGetSize(const py::handle &list) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::ssize_t ListGetSize(const py::handle &list) { return PyList_GET_SIZE(list.ptr()); } -inline Py_ALWAYS_INLINE py::ssize_t DictGetSize(const py::handle &dict) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::ssize_t DictGetSize(const py::handle &dict) { #if defined(PyDict_GET_SIZE) return PyDict_GET_SIZE(dict.ptr()); #else @@ -106,14 +108,17 @@ inline Py_ALWAYS_INLINE py::ssize_t DictGetSize(const py::handle &dict) { } template >> -inline Py_ALWAYS_INLINE T TupleGetItemAs(const py::handle &tuple, const py::ssize_t &index) { +[[nodiscard]] inline Py_ALWAYS_INLINE T TupleGetItemAs(const py::handle &tuple, + const py::ssize_t &index) { return py::reinterpret_borrow(PyTuple_GET_ITEM(tuple.ptr(), index)); } -inline Py_ALWAYS_INLINE py::object TupleGetItem(const py::handle &tuple, const py::ssize_t &index) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::object TupleGetItem(const py::handle &tuple, + const py::ssize_t &index) { return TupleGetItemAs(tuple, index); } template >> -inline Py_ALWAYS_INLINE T ListGetItemAs(const py::handle &list, const py::ssize_t &index) { +[[nodiscard]] inline Py_ALWAYS_INLINE T ListGetItemAs(const py::handle &list, + const py::ssize_t &index) { #if PY_VERSION_HEX >= 0x030D00A4 // Python 3.13.0a4 PyObject * const item = PyList_GetItemRef(list.ptr(), index); if (item == nullptr) [[unlikely]] { @@ -121,14 +126,22 @@ inline Py_ALWAYS_INLINE T ListGetItemAs(const py::handle &list, const py::ssize_ } return py::reinterpret_steal(item); #else + // Bounds-check like `PyList_GetItemRef` does: callers read the length once and then run user + // code, so the list can shrink mid-loop and the unchecked macro would read out of bounds. + if (index < 0 || index >= PyList_GET_SIZE(list.ptr())) [[unlikely]] { + py::set_error(PyExc_IndexError, "list index out of range"); + throw py::error_already_set(); + } return py::reinterpret_borrow(PyList_GET_ITEM(list.ptr(), index)); #endif } -inline Py_ALWAYS_INLINE py::object ListGetItem(const py::handle &list, const py::ssize_t &index) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::object ListGetItem(const py::handle &list, + const py::ssize_t &index) { return ListGetItemAs(list, index); } template >> -inline Py_ALWAYS_INLINE T DictGetItemAs(const py::handle &dict, const py::handle &key) { +[[nodiscard]] inline Py_ALWAYS_INLINE T DictGetItemAs(const py::handle &dict, + const py::handle &key) { #if PY_VERSION_HEX >= 0x030D00A1 // Python 3.13.0a1 PyObject *value = nullptr; if (PyDict_GetItemRef(dict.ptr(), key.ptr(), &value) < 0) [[unlikely]] { @@ -143,7 +156,8 @@ inline Py_ALWAYS_INLINE T DictGetItemAs(const py::handle &dict, const py::handle return py::reinterpret_borrow(PyDict_GetItem(dict.ptr(), key.ptr())); #endif } -inline Py_ALWAYS_INLINE py::object DictGetItem(const py::handle &dict, const py::handle &key) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::object DictGetItem(const py::handle &dict, + const py::handle &key) { return DictGetItemAs(dict, key); } @@ -165,6 +179,25 @@ inline Py_ALWAYS_INLINE void DictSetItem(const py::handle &dict, } } +// Shallow copies through the C API. A Python-level `.copy()` would cost an attribute lookup, a +// bound-method allocation and a vectorcall per call. +[[nodiscard]] inline py::list ListCopy(const py::handle &list) { + const scoped_critical_section cs{list}; + auto copy = py::reinterpret_steal(PyList_GetSlice(list.ptr(), 0, ListGetSize(list))); + if (!copy) [[unlikely]] { + throw py::error_already_set(); + } + return copy; +} +[[nodiscard]] inline py::dict DictCopy(const py::handle &dict) { + const scoped_critical_section cs{dict}; + auto copy = py::reinterpret_steal(PyDict_Copy(dict.ptr())); + if (!copy) [[unlikely]] { + throw py::error_already_set(); + } + return copy; +} + inline Py_ALWAYS_INLINE void AssertExactList(const py::handle &object) { if (!PyList_CheckExact(object.ptr())) [[unlikely]] { throw py::value_error("Expected an instance of list, got " + PyRepr(object) + "."); @@ -213,8 +246,194 @@ inline Py_ALWAYS_INLINE void AssertExactDeque(const py::handle &object) { } } +// A process-global cache mapping a Python object (in practice a type) to a value computed from it, +// e.g. whether a type is a namedtuple. It is a function-local static shared by every interpreter +// and outlives the Python runtime. +// +// Entries are keyed by `(interpreter_id, object address)` rather than the address alone, because a +// shared key can map to a per-interpreter-owned result: `int` is immortal and lives at the same +// address in every interpreter, while a computed `py::tuple` belongs to the interpreter that made +// it. An address-only key would hand that value to another interpreter to use after the owner frees +// it on finalization. +// +// A per-entry weakref evicts an entry when its key is collected, so a later key reusing that +// address cannot read a stale value. A per-interpreter `atexit` callback clears that interpreter's +// entries, covering what the weakref cannot: an immortal key is never collected, and interpreter +// ids restart from 0 after a `Py_Finalize`/`Py_Initialize` cycle, so a fresh interpreter must not +// inherit a finalized one's entries. +// +// `ValueType` may be a value such as `bool`, or a pybind11 reference whose entry owns one reference +// that is dropped on eviction. +template +class WeakKeyCache { +public: + explicit WeakKeyCache(const std::size_t &max_size) : m_max_size{max_size} {} + ~WeakKeyCache() = default; + + WeakKeyCache(const WeakKeyCache &) = delete; + WeakKeyCache(WeakKeyCache &&) = delete; + WeakKeyCache &operator=(const WeakKeyCache &) = delete; + WeakKeyCache &operator=(WeakKeyCache &&) = delete; + + // Return the value cached for `key`, computing and inserting it via `compute` on a miss. + // `compute` is a nullary callable returning `ValueType`, invoked with the GIL held and the + // cache lock NOT held. + template + // NOLINTNEXTLINE[readability-function-cognitive-complexity] + [[nodiscard]] ValueType LookupOrInsert(const py::handle &key, Compute &&compute) { + // Read the interpreter id (part of the cache key) while the GIL is still held, before the + // read lock below releases it: `GetCurrentPyInterpreterID()` needs a valid thread state. + const interpid_t interpreter_id = GetCurrentPyInterpreterID(); + const CacheKey cache_key{interpreter_id, key}; + std::optional cached_value{}; + { +#if !defined(Py_GIL_DISABLED) + const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_read_lock lock{m_mutex}; + const auto it = m_cache.find(cache_key); + if (it != m_cache.end()) [[likely]] { + cached_value = it->second; + } + } + // The read lock is released and the GIL re-acquired (in that destruction order) BEFORE the + // borrowed object is touched, so the GIL is never (re-)acquired while the lock is held. + // Doing so would invert the lock order against the weakref eviction callback (which holds + // the GIL, then takes the write lock) and could deadlock. `key` stays alive for the whole + // call, so its entry cannot be evicted and `cached_value` stays valid. + if (cached_value.has_value()) [[likely]] { + if constexpr (std::is_same_v) { + // A value or a `py::handle`: the stored type is the value type. + return *cached_value; + } else { + // An owning object stored as a borrowed `py::handle`: return a fresh owning borrow. + return py::reinterpret_borrow(*cached_value); + } + } + + ValueType value = std::forward(compute)(); + + // Register the per-interpreter cleanup BEFORE publishing an entry. It runs Python and can + // raise, so the interpreter is claimed first and only marked done once it succeeds: marking + // done first would leave the interpreter believing a callback exists and never retry it. + bool claimed = false; + bool registered = false; + { +#if !defined(Py_GIL_DISABLED) + const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_write_lock lock{m_mutex}; + const auto [it, inserted] = m_cleanup_registered.try_emplace(interpreter_id, false); + claimed = inserted; + registered = it->second; + } + if (claimed) [[unlikely]] { + try { + RegisterInterpreterCleanup(interpreter_id); + } catch (...) { + // Take the lock with the GIL held, the same order the eviction callback uses. + const scoped_write_lock lock{m_mutex}; + m_cleanup_registered.erase(interpreter_id); + throw; + } + + { + const scoped_write_lock lock{m_mutex}; + m_cleanup_registered[interpreter_id] = true; + registered = true; + } + } + + // Publish only against a completed registration: publishing while another thread is still + // registering would outlive the owner if that registration raises. Skipping is safe. + bool inserted = false; + if (registered) [[likely]] { +#if !defined(Py_GIL_DISABLED) + const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_write_lock lock{m_mutex}; + if (m_cache.size() < m_max_size) [[likely]] { + // The GIL is released here, so store the value without touching any refcount (a + // reference value is stored as a borrowed `py::handle` and owned by the `inc_ref()` + // below). + inserted = m_cache.emplace(cache_key, StoredType{value}).second; + } + } + if (inserted) [[likely]] { + // The GIL is held here, so we can safely increment the reference count and create the + // weakref. If the weakref cannot be created, drop the entry again: a published entry + // whose value is unowned and whose key has no eviction callback would be read back + // after the value is freed. + if constexpr (kValueIsPyReference) { + value.inc_ref(); + } + try { + (void)py::weakref(key, + py::cpp_function([this, cache_key](py::handle weakref) -> void { + const scoped_write_lock lock{m_mutex}; + const auto it = m_cache.find(cache_key); + if (it != m_cache.end()) [[likely]] { + if constexpr (kValueIsPyReference) { + it->second.dec_ref(); + } + m_cache.erase(it); + } + weakref.dec_ref(); + })) + .release(); + } catch (...) { + { + const scoped_write_lock lock{m_mutex}; + m_cache.erase(cache_key); + } + if constexpr (kValueIsPyReference) { + value.dec_ref(); + } + throw; + } + } + return value; + } + +private: + static constexpr bool kValueIsPyReference = std::is_base_of_v; + using StoredType = std::conditional_t; + using CacheKey = std::pair; + + // Register (once per interpreter, with the GIL held and WITHOUT the cache lock, mirroring + // `PyTreeTypeRegistry::Init`) an `atexit` callback that evicts this interpreter's entries on + // shutdown. + void RegisterInterpreterCleanup(const interpid_t &interpreter_id) { + auto atexit_register = py::getattr(py::module_::import("atexit"), "register"); + atexit_register(py::cpp_function([this, interpreter_id]() -> void { + const scoped_write_lock lock{m_mutex}; + for (auto it = m_cache.begin(); it != m_cache.end();) { + if (it->first.first == interpreter_id) [[likely]] { + if constexpr (kValueIsPyReference) { + it->second.dec_ref(); + } + it = m_cache.erase(it); + } else [[unlikely]] { + ++it; + } + } + m_cleanup_registered.erase(interpreter_id); + })); + } + + std::unordered_map m_cache{}; + // Interpreter id -> whether its `atexit` callback is registered. An id maps to `false` while + // the claiming thread is still registering it. + std::unordered_map m_cleanup_registered{}; + const std::size_t m_max_size{}; + mutable read_write_mutex m_mutex{}; +}; + +// The maximum size of a type cache. +constexpr std::size_t MAX_TYPE_CACHE_SIZE = 4096; + // NOLINTNEXTLINE[readability-function-cognitive-complexity] -inline bool IsNamedTupleClassImpl(const py::handle &type) { +[[nodiscard]] inline bool IsNamedTupleClassImpl(const py::handle &type) { // We can only identify namedtuples heuristically, here by the presence of a _fields attribute. if (PyType_FastSubclass(reinterpret_cast(type.ptr()), Py_TPFLAGS_TUPLE_SUBCLASS)) [[unlikely]] { @@ -252,50 +471,20 @@ inline bool IsNamedTupleClassImpl(const py::handle &type) { } return false; } -inline bool IsNamedTupleClass(const py::handle &type) { +[[nodiscard]] inline bool IsNamedTupleClass(const py::handle &type) { if (!PyType_Check(type.ptr())) [[unlikely]] { return false; } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - return it->second; - } - } - - const bool result = EVALUATE_WITH_LOCK_HELD(IsNamedTupleClassImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, result).second; - } - } - if (cache_inserted) [[likely]] { - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - cache.erase(type); - weakref.dec_ref(); - })) - .release(); - } - return result; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> bool { + return EVALUATE_WITH_LOCK_HELD(IsNamedTupleClassImpl(type), type); + }); } -inline Py_ALWAYS_INLINE bool IsNamedTupleInstance(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE bool IsNamedTupleInstance(const py::handle &object) { return IsNamedTupleClass(py::type::handle_of(object)); } -inline Py_ALWAYS_INLINE bool IsNamedTuple(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE bool IsNamedTuple(const py::handle &object) { const py::handle type = (PyType_Check(object.ptr()) ? object : py::type::handle_of(object)); return IsNamedTupleClass(type); } @@ -305,7 +494,7 @@ inline Py_ALWAYS_INLINE void AssertExactNamedTuple(const py::handle &object) { PyRepr(object) + "."); } } -inline py::tuple NamedTupleGetFields(const py::handle &object) { +[[nodiscard]] inline py::tuple NamedTupleGetFields(const py::handle &object) { py::handle type; if (PyType_Check(object.ptr())) [[unlikely]] { type = object; @@ -323,7 +512,7 @@ inline py::tuple NamedTupleGetFields(const py::handle &object) { return EVALUATE_WITH_LOCK_HELD(py::getattr(type, "_fields"), type); } -inline bool IsStructSequenceClassImpl(const py::handle &type) { +[[nodiscard]] inline bool IsStructSequenceClassImpl(const py::handle &type) { // We can only identify PyStructSequences heuristically, here by the presence of // n_fields, n_sequence_fields, n_unnamed_fields attributes. auto * const type_object = reinterpret_cast(type.ptr()); @@ -363,50 +552,20 @@ inline bool IsStructSequenceClassImpl(const py::handle &type) { } return false; } -inline bool IsStructSequenceClass(const py::handle &type) { +[[nodiscard]] inline bool IsStructSequenceClass(const py::handle &type) { if (!PyType_Check(type.ptr())) [[unlikely]] { return false; } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - return it->second; - } - } - - const bool result = EVALUATE_WITH_LOCK_HELD(IsStructSequenceClassImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, result).second; - } - } - if (cache_inserted) [[likely]] { - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - cache.erase(type); - weakref.dec_ref(); - })) - .release(); - } - return result; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> bool { + return EVALUATE_WITH_LOCK_HELD(IsStructSequenceClassImpl(type), type); + }); } -inline Py_ALWAYS_INLINE bool IsStructSequenceInstance(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE bool IsStructSequenceInstance(const py::handle &object) { return IsStructSequenceClass(py::type::handle_of(object)); } -inline Py_ALWAYS_INLINE bool IsStructSequence(const py::handle &object) { +[[nodiscard]] inline Py_ALWAYS_INLINE bool IsStructSequence(const py::handle &object) { const py::handle type = (PyType_Check(object.ptr()) ? object : py::type::handle_of(object)); return IsStructSequenceClass(type); } @@ -416,7 +575,7 @@ inline Py_ALWAYS_INLINE void AssertExactStructSequence(const py::handle &object) PyRepr(object) + "."); } } -inline py::tuple StructSequenceGetFieldsImpl(const py::handle &type) { +[[nodiscard]] inline py::tuple StructSequenceGetFieldsImpl(const py::handle &type) { #if defined(PYPY_VERSION) py::list fields{}; py::exec( @@ -424,28 +583,55 @@ inline py::tuple StructSequenceGetFieldsImpl(const py::handle &type) { import sys StructSequenceFieldType = type(type(sys.version_info).major) - indices_by_name = { - name: member.index + # PyPy has no unnamed fields; a descriptor's `.index` is its sequence position. + # Map index -> name and defensively fill any missing (unnamed) slot with the marker. + names_by_index = { + member.index: name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) } - fields.extend(sorted(indices_by_name, key=indices_by_name.get)[:cls.n_sequence_fields]) + fields.extend( + names_by_index.get(index, unnamed_field) for index in range(cls.n_sequence_fields) + ) )py", - py::dict(py::arg("cls") = type, py::arg("fields") = fields)); + py::dict(py::arg("cls") = type, + py::arg("fields") = fields, + py::arg("unnamed_field") = py::str(PyStructSequenceUnnamedField()))); return py::tuple{fields}; #else const auto n_sequence_fields = thread_safe_cast( EVALUATE_WITH_LOCK_HELD(py::getattr(type, "n_sequence_fields"), type)); const auto * const members = reinterpret_cast(type.ptr())->tp_members; + // `tp_members` lists only the NAMED fields, but each carries a byte `offset` encoding its + // sequence index, relative to `offsetof(PyTupleObject, ob_item)`. Map each member back to its + // slot by offset: indexing `members[i]` by position mislabels every slot after the first + // unnamed one (e.g. `os.stat_result` slots 7/8/9 reported as `st_atime`/`st_mtime`/`st_ctime`). py::tuple fields{n_sequence_fields}; + // Fill the named slots first, then default the remaining (unnamed) slots to the marker. + // Pre-filling every slot with the marker and overwriting the named ones would leak each + // overwritten marker: `TupleSetItem` uses `PyTuple_SET_ITEM`, which does not decref what it + // replaces. + std::vector named(n_sequence_fields, false); + for (const PyMemberDef *member = members; member != nullptr && member->name != nullptr; + // NOLINTNEXTLINE[cppcoreguidelines-pro-bounds-pointer-arithmetic] + ++member) { + const py::ssize_t index = + (member->offset - py::ssize_t_cast(offsetof(PyTupleObject, ob_item))) / + py::ssize_t_cast(sizeof(PyObject *)); + if (index >= 0 && index < n_sequence_fields) [[likely]] { + TupleSetItem(fields, index, py::str(member->name)); + named[index] = true; + } + } for (py::ssize_t i = 0; i < n_sequence_fields; ++i) { - // NOLINTNEXTLINE[cppcoreguidelines-pro-bounds-pointer-arithmetic] - TupleSetItem(fields, i, py::str(members[i].name)); + if (!named[i]) [[unlikely]] { + TupleSetItem(fields, i, py::str(PyStructSequenceUnnamedField())); + } } return fields; #endif } -inline py::tuple StructSequenceGetFields(const py::handle &object) { +[[nodiscard]] inline py::tuple StructSequenceGetFields(const py::handle &object) { py::handle type; if (PyType_Check(object.ptr())) [[unlikely]] { type = object; @@ -460,90 +646,59 @@ inline py::tuple StructSequenceGetFields(const py::handle &object) { } } - static auto cache = std::unordered_map{}; - static read_write_mutex mutex{}; - bool cache_inserted = false; - - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_read_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_acquire_simple gil_acquire{}; -#endif - return py::reinterpret_borrow(it->second); - } - } - - const py::tuple fields = EVALUATE_WITH_LOCK_HELD(StructSequenceGetFieldsImpl(type), type); - { -#if !defined(Py_GIL_DISABLED) - const py::gil_scoped_release_simple gil_release{}; -#endif - const scoped_write_lock lock{mutex}; - if (cache.size() < MAX_TYPE_CACHE_SIZE) [[likely]] { - cache_inserted = cache.emplace(type, fields).second; - } - } - if (cache_inserted) [[likely]] { - fields.inc_ref(); - (void)py::weakref(type, py::cpp_function([type](py::handle weakref) -> void { - const scoped_write_lock lock{mutex}; - const auto it = cache.find(type); - if (it != cache.end()) [[likely]] { - it->second.dec_ref(); - cache.erase(it); - } - weakref.dec_ref(); - })) - .release(); - } - return fields; + static WeakKeyCache cache{MAX_TYPE_CACHE_SIZE}; + return cache.LookupOrInsert(type, [&type]() -> py::tuple { + return EVALUATE_WITH_LOCK_HELD(StructSequenceGetFieldsImpl(type), type); + }); } +// `list.sort()` leaves the list partially reordered when a comparison raises, so each attempt sorts +// a copy and only a fully sorted one is committed (mirrors `optree.utils.total_order_sorted`). inline void TotalOrderSort(py::list &list) { // NOLINT[runtime/references] + py::list sorted = ListCopy(list); try { // Sort directly if possible. - if (static_cast(EVALUATE_WITH_LOCK_HELD(PyList_Sort(list.ptr()), list))) + if (static_cast(EVALUATE_WITH_LOCK_HELD(PyList_Sort(sorted.ptr()), sorted))) [[unlikely]] { throw py::error_already_set(); } + list = std::move(sorted); + return; } catch (py::error_already_set &ex1) { - if (ex1.matches(PyExc_TypeError)) [[likely]] { - // Found incomparable keys (e.g. `int` vs. `str`, or user-defined types). - try { - // Sort with `(f'{obj.__class__.__module__}.{obj.__class__.__qualname__}', obj)` - const auto sort_key_fn = py::cpp_function([](const py::object &obj) -> py::tuple { - const py::handle cls = py::type::handle_of(obj); - const py::str qualname{ - EVALUATE_WITH_LOCK_HELD(PyStr(py::getattr(cls, "__module__")) + "." + - PyStr(py::getattr(cls, "__qualname__")), - cls)}; - return py::make_tuple(qualname, obj); - }); - { - const scoped_critical_section cs{list}; - py::getattr(list, "sort")(py::arg("key") = sort_key_fn); - } - } catch (py::error_already_set &ex2) { - if (ex2.matches(PyExc_TypeError)) [[likely]] { - // Found incomparable user-defined key types. - // The keys remain in the insertion order. - PyErr_Clear(); - } else [[unlikely]] { - std::rethrow_exception(std::current_exception()); - } - } + if (!ex1.matches(PyExc_TypeError)) [[unlikely]] { + std::rethrow_exception(std::current_exception()); + } + // Found incomparable keys (e.g. `int` vs. `str`, or user-defined types). + } + + sorted = ListCopy(list); + try { + // Sort with `(f'{obj.__class__.__module__}.{obj.__class__.__qualname__}', obj)` + const auto sort_key_fn = py::cpp_function([](const py::object &obj) -> py::tuple { + const py::handle cls = py::type::handle_of(obj); + const py::str qualname{ + EVALUATE_WITH_LOCK_HELD(PyStr(py::getattr(cls, "__module__")) + "." + + PyStr(py::getattr(cls, "__qualname__")), + cls)}; + return py::make_tuple(qualname, obj); + }); + { + const scoped_critical_section cs{sorted}; + py::getattr(sorted, "sort")(py::arg("key") = sort_key_fn); + } + list = std::move(sorted); + } catch (py::error_already_set &ex2) { + if (ex2.matches(PyExc_TypeError)) [[likely]] { + // Found incomparable user-defined key types. + // The keys remain in the insertion order. + PyErr_Clear(); } else [[unlikely]] { std::rethrow_exception(std::current_exception()); } } } -inline Py_ALWAYS_INLINE py::list DictKeys(const py::dict &dict) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::list DictKeys(const py::dict &dict) { const scoped_critical_section cs{dict}; return py::reinterpret_steal(PyDict_Keys(dict.ptr())); } @@ -551,7 +706,7 @@ inline Py_ALWAYS_INLINE py::list DictKeys(const py::dict &dict) { // Equivalent to Python's `dict.fromkeys(iterable)`: returns a new `dict[Key, None]` with the // keys taken from `iterable` in order. When `iterable` is itself a dict, this hits CPython's // `dict_dict_fromkeys` fast path (contiguous bucket copy, no per-key rehash). -inline Py_ALWAYS_INLINE py::dict DictFromKeys(const py::handle &iterable) { +[[nodiscard]] inline Py_ALWAYS_INLINE py::dict DictFromKeys(const py::handle &iterable) { const scoped_critical_section cs{iterable}; // NOLINTNEXTLINE[cppcoreguidelines-pro-type-vararg] PyObject *result = PyObject_CallMethod(reinterpret_cast(&PyDict_Type), @@ -564,13 +719,15 @@ inline Py_ALWAYS_INLINE py::dict DictFromKeys(const py::handle &iterable) { return py::reinterpret_steal(result); } -inline py::list SortedDictKeys(const py::dict &dict) { +[[nodiscard]] inline py::list SortedDictKeys(const py::dict &dict) { py::list keys = DictKeys(dict); TotalOrderSort(keys); return keys; } -inline bool DictKeysEqual(const py::list & /*unique*/ keys, const py::dict &dict) { +// Test whether `keys` and the keys of `dict` are the same set. +// Precondition: `keys` holds no duplicates; the length shortcut below relies on that. +[[nodiscard]] inline bool DictKeysEqual(const py::list &keys, const py::dict &dict) { const scoped_critical_section2 cs{keys, dict}; const py::ssize_t list_len = ListGetSize(keys); const py::ssize_t dict_len = DictGetSize(dict); @@ -590,8 +747,10 @@ inline bool DictKeysEqual(const py::list & /*unique*/ keys, const py::dict &dict return true; } -inline std::pair DictKeysDifference(const py::list & /*unique*/ keys, - const py::dict &dict) { +// Return the keys of `keys` missing from `dict` and the keys of `dict` not in `keys`. +// Precondition: `keys` holds no duplicates. +[[nodiscard]] inline std::pair DictKeysDifference(const py::list &keys, + const py::dict &dict) { const py::set expected_keys = EVALUATE_WITH_LOCK_HELD(py::set{keys}, keys); const py::set got_keys = EVALUATE_WITH_LOCK_HELD(py::set{dict}, dict); py::list missing_keys{expected_keys - got_keys}; @@ -600,3 +759,12 @@ inline std::pair DictKeysDifference(const py::list & /*uniqu TotalOrderSort(extra_keys); return std::make_pair(std::move(missing_keys), std::move(extra_keys)); } + +[[nodiscard]] inline py::ssize_t DistinctCount(const py::handle &iterable) { + const scoped_critical_section cs{iterable}; + const auto set = py::reinterpret_steal(PySet_New(iterable.ptr())); + if (!set) [[unlikely]] { // a non-iterable or an unhashable element raised here + throw py::error_already_set(); + } + return PySet_GET_SIZE(set.ptr()); +} diff --git a/include/optree/registry.h b/include/optree/registry.h index 5d2be589..8dcfbcb9 100644 --- a/include/optree/registry.h +++ b/include/optree/registry.h @@ -17,7 +17,7 @@ limitations under the License. #pragma once -#include // std::uint8_t +#include // std::uint8_t, UINT8_MAX #include // std::shared_ptr #include // std::optional, std::nullopt #include // std::string @@ -56,6 +56,12 @@ enum class PyTreeKind : std::uint8_t { NumKinds, // Number of kinds (placed at the end) }; +// A new pytree kind must keep `NumKinds` within the `std::uint8_t` underlying type; otherwise the +// enum cannot represent every value and code that narrows a kind to `std::uint8_t` (e.g. pickle +// deserialization in `PyTreeSpec::FromPicklable`) would silently wrap. +static_assert(static_cast(PyTreeKind::NumKinds) <= UINT8_MAX, + "PyTreeKind::NumKinds overflows its std::uint8_t underlying type."); + constexpr PyTreeKind kCustom = PyTreeKind::Custom; constexpr PyTreeKind kLeaf = PyTreeKind::Leaf; constexpr PyTreeKind kNone = PyTreeKind::None; @@ -133,8 +139,12 @@ class PyTreeTypeRegistry { auto ®istry1 = GetSingleton(); auto ®istry2 = GetSingleton(); - const ssize_t count1 = registry1.Size(registry_namespace); - const ssize_t count2 = registry2.Size(registry_namespace); + // Read both registries under a single lock so the two counts form a consistent snapshot. + // Two separate `Size()` calls each drop the lock, letting a concurrent (un)registration + // slip between them and spuriously trip the invariant check below. + const scoped_read_lock lock{sm_mutex}; + const ssize_t count1 = registry1.SizeImpl(registry_namespace); + const ssize_t count2 = registry2.SizeImpl(registry_namespace); EXPECT_EQ(count1, count2 + 1, "The number of registered types in the two registries should match " @@ -165,12 +175,36 @@ class PyTreeTypeRegistry { [[nodiscard]] static inline Py_ALWAYS_INLINE bool IsDictInsertionOrdered( const std::string ®istry_namespace, const bool &inherit_global_namespace = true) { + const auto flags = GetDictInsertionOrderedFlags(registry_namespace); + return inherit_global_namespace ? flags.with_inherited_global_namespace + : flags.in_current_namespace; + } + + // Whether dictionary key insertion order is preserved during flattening. Both flags are + // computed together under a single lock so callers see a consistent snapshot; this also avoids + // recursively read-locking `sm_dict_order_mutex` (which is not a recursive mutex) that would + // happen if a caller held the lock while calling `IsDictInsertionOrdered`. + struct DictInsertionOrderedFlags { + // Whether the given namespace itself preserves insertion order. + bool in_current_namespace; + // Whether the given namespace, or the inherited global namespace, preserves insertion + // order. + bool with_inherited_global_namespace; + }; + + [[nodiscard]] static inline Py_ALWAYS_INLINE DictInsertionOrderedFlags + GetDictInsertionOrderedFlags(const std::string ®istry_namespace) { const scoped_read_lock lock{sm_dict_order_mutex}; const auto interpid = GetCurrentPyInterpreterID(); const auto &namespaces = sm_dict_insertion_ordered_namespaces; - return (namespaces.find({interpid, registry_namespace}) != namespaces.end()) || - (inherit_global_namespace && namespaces.find({interpid, ""}) != namespaces.end()); + const bool in_current_namespace = + namespaces.find({interpid, registry_namespace}) != namespaces.end(); + return { + .in_current_namespace = in_current_namespace, + .with_inherited_global_namespace = + in_current_namespace || namespaces.find({interpid, ""}) != namespaces.end(), + }; } // Set the namespace to preserve the insertion order of the dictionary keys during flattening. @@ -194,16 +228,29 @@ class PyTreeTypeRegistry { template [[nodiscard]] static PyTreeTypeRegistry &GetSingleton(); - template - static void RegisterImpl(const py::object &cls, - const py::function &flatten_func, - const py::function &unflatten_func, - const py::object &path_entry_type, - const std::string ®istry_namespace); + // Why an (un)registration failed. Formatting the message needs `PyRepr`, which runs user Python + // (e.g. a metaclass `__repr__`); doing that under `sm_mutex` inverts the GIL <-> `sm_mutex` + // lock order against a concurrent flatten and deadlocks, so the caller raises after unlocking. + enum class RegistryStatus : std::uint8_t { + Ok = 0, + BuiltinType, + AlreadyRegistered, + NotRegistered, + }; - template - [[nodiscard]] static RegistrationPtr UnregisterImpl(const py::object &cls, - const std::string ®istry_namespace); + [[nodiscard]] RegistryStatus RegisterImpl(const py::object &cls, + const py::function &flatten_func, + const py::function &unflatten_func, + const py::object &path_entry_type, + const std::string ®istry_namespace); + + [[nodiscard]] RegistryStatus UnregisterImpl( + const py::object &cls, + const std::string ®istry_namespace, + RegistrationPtr ®istration); // NOLINT[runtime/references] + + // Get the number of registered types without locking. The caller must hold `sm_mutex`. + [[nodiscard]] ssize_t SizeImpl(const std::optional ®istry_namespace) const; // Initialize the registry for the current interpreter. static void Init(); diff --git a/include/optree/treespec.h b/include/optree/treespec.h index 1e132a51..05197908 100644 --- a/include/optree/treespec.h +++ b/include/optree/treespec.h @@ -17,6 +17,7 @@ limitations under the License. #pragma once +#include // std::atomic #include // std::unique_ptr #include // std::optional, std::nullopt #include // std::string @@ -25,10 +26,13 @@ limitations under the License. #include // std::pair #include // std::vector +#include + #include #include "optree/exceptions.h" #include "optree/hashing.h" +#include "optree/pymacros.h" #include "optree/registry.h" #include "optree/synchronization.h" @@ -39,7 +43,15 @@ using size_t = py::size_t; using ssize_t = py::ssize_t; // The maximum depth of a pytree. -#if defined(Py_DEBUG) || defined(PYPY_VERSION) || defined(MS_WINDOWS) || \ +#if defined(MS_WINDOWS) && (defined(Py_DEBUG) || defined(Py_GIL_DISABLED)) +// A debug or free-threading build on Windows combines large frames with a 1MB default stack, and +// the walkers still need room to unwind the `RecursionError` thrown at the guard. +# if PY_VERSION_HEX < 0x030A0000 // Python 3.10 +constexpr ssize_t MAX_RECURSION_DEPTH = 100; +# else +constexpr ssize_t MAX_RECURSION_DEPTH = 250; +# endif +#elif defined(Py_DEBUG) || defined(PYPY_VERSION) || defined(MS_WINDOWS) || \ (defined(__wasm__) || defined(__wasm32__) || defined(__wasm64__) || defined(__wasi__) || \ defined(__EMSCRIPTEN__)) constexpr ssize_t MAX_RECURSION_DEPTH = 500; @@ -261,6 +273,7 @@ class PyTreeSpec { friend void BuildModule(py::module_ &mod); // NOLINT[runtime/references] private: + using Registration = PyTreeTypeRegistry::Registration; using RegistrationPtr = PyTreeTypeRegistry::RegistrationPtr; using ThreadedIdentity = std::pair; @@ -314,7 +327,7 @@ class PyTreeSpec { // Manufacture an instance of a node given its children. [[nodiscard]] static py::object MakeNode( const Node &node, - const py::object children[], // NOLINT[hicpp-avoid-c-arrays] + const py::object children[], // NOLINT[cppcoreguidelines-avoid-c-arrays] const size_t &num_children); // Identify the path entry class for a node. @@ -364,7 +377,15 @@ class PyTreeSpec { const std::vector &traversal, const ssize_t &pos, const std::vector &other_traversal, - const ssize_t &other_pos); + const ssize_t &other_pos, + const ssize_t &depth); + + // Return the type of the first custom node whose registration under `target_namespace` is no + // longer the one it holds, either re-registered or unregistered altogether, or a null object + // if all custom nodes are consistent. Used by namespace merges (compose / broadcast) to detect + // an unsafe re-tagging; the caller raises the error. + [[nodiscard]] std::optional FindStaleCustomType( + const std::string &target_namespace) const; template [[nodiscard]] py::object WalkImpl( @@ -431,11 +452,14 @@ class PyTreeIter { private: py::object m_root; std::vector> m_agenda; - const std::optional m_leaf_predicate; + std::optional m_leaf_predicate; const bool m_none_is_leaf; const std::string m_namespace; const bool m_is_dict_insertion_ordered; mutable mutex m_mutex{}; + // The thread currently running `Next()`, or a default-constructed id when idle. Used to reject + // re-entrant iteration instead of deadlocking on the non-recursive `m_mutex`. + std::atomic m_running_thread_id{}; template [[nodiscard]] py::object NextImpl(); diff --git a/optree/_C.pyi b/optree/_C.pyi index 2643289c..94d0f059 100644 --- a/optree/_C.pyi +++ b/optree/_C.pyi @@ -20,7 +20,7 @@ import enum import sys from collections.abc import Callable, Collection, Iterable, Iterator from types import MappingProxyType -from typing import Any, ClassVar, Final, final +from typing import Any, ClassVar, Final, SupportsIndex, SupportsInt, final, overload from typing_extensions import Self # Python 3.11+ from optree.typing import ( @@ -37,6 +37,9 @@ from optree.typing import ( # Set if the type allows subclassing (see CPython's Include/object.h) Py_TPFLAGS_BASETYPE: Final[int] # (1UL << 10) +# The name reported for an unnamed PyStructSequence slot (see CPython's Include/structseq.h) +PyStructSequence_UnnamedField: Final[str] # 'unnamed field' + # Meta-information during build-time BUILDTIME_METADATA: Final[MappingProxyType[str, Any]] PY_VERSION: Final[str] @@ -99,26 +102,44 @@ class PyTreeSpec: f_leaf: Callable[[Self], Self] | None = None, ) -> Self: ... def compose(self, inner: Self, /) -> Self: ... + @overload + def traverse( + self, + leaves: Iterable[T], + /, + f_node: None = None, + f_leaf: Callable[[T], U] | None = None, + ) -> PyTree[U]: ... + @overload def traverse( self, leaves: Iterable[T], /, - f_node: Callable[[Collection[U]], U] | None = None, + f_node: Callable[[Collection[U]], U] = ..., f_leaf: Callable[[T], U] | None = None, ) -> U: ... + @overload + def walk( + self, + leaves: Iterable[T], + /, + f_node: None = None, + f_leaf: Callable[[T], U] | None = None, + ) -> PyTree[U]: ... + @overload def walk( self, leaves: Iterable[T], /, - f_node: Callable[[builtins.type, MetaData, tuple[U, ...]], U] | None = None, + f_node: Callable[[builtins.type, MetaData, tuple[U, ...]], U] = ..., f_leaf: Callable[[T], U] | None = None, ) -> U: ... def paths(self, /) -> list[tuple[Any, ...]]: ... def accessors(self, /) -> list[PyTreeAccessor]: ... def entries(self, /) -> list[Any]: ... - def entry(self, index: int, /) -> Any: ... + def entry(self, index: SupportsInt | SupportsIndex, /) -> Any: ... def children(self, /) -> list[Self]: ... - def child(self, index: int, /) -> Self: ... + def child(self, index: SupportsInt | SupportsIndex, /) -> Self: ... def one_level(self, /) -> Self | None: ... def is_leaf(self, /, *, strict: bool = True) -> bool: ... def is_one_level(self, /) -> bool: ... @@ -188,11 +209,11 @@ def is_leaf( ) -> bool: ... def is_namedtuple(obj: object | type, /) -> bool: ... def is_namedtuple_instance(obj: object, /) -> bool: ... -def is_namedtuple_class(cls: type, /) -> bool: ... +def is_namedtuple_class(cls: object, /) -> bool: ... def namedtuple_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ... def is_structseq(obj: object | type, /) -> bool: ... def is_structseq_instance(obj: object, /) -> bool: ... -def is_structseq_class(cls: type, /) -> bool: ... +def is_structseq_class(cls: object, /) -> bool: ... def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ... # Registration functions diff --git a/optree/accessors.py b/optree/accessors.py index 181b3df8..6fcd2399 100644 --- a/optree/accessors.py +++ b/optree/accessors.py @@ -89,15 +89,15 @@ def __eq__(self, other: object, /) -> bool: self.entry, self.type, self.kind, - self.__class__.__call__.__code__.co_code, - self.__class__.codify.__code__.co_code, + self.__class__.__call__, + self.__class__.codify, ) == ( other.entry, other.type, other.kind, - other.__class__.__call__.__code__.co_code, - other.__class__.codify.__code__.co_code, + other.__class__.__call__, + other.__class__.codify, ) ) @@ -108,8 +108,8 @@ def __hash__(self, /) -> int: self.entry, self.type, self.kind, - self.__class__.__call__.__code__.co_code, - self.__class__.codify.__code__.co_code, + self.__class__.__call__, + self.__class__.codify, ), ) @@ -215,7 +215,9 @@ def __call__(self, obj: Any, /) -> Any: def codify(self, /, node: str = '') -> str: """Generate code for accessing the path entry.""" - return f'{node}.{self.name}' + for name in self.name.split('.'): + node = f'{node}.{name}' if name.isidentifier() else f'getattr({node}, {name!r})' + return node class FlattenedEntry(PyTreeEntry): # pylint: disable=too-few-public-methods @@ -319,12 +321,26 @@ def field(self, /) -> str: """Get the field name.""" return self.fields[self.entry] + @property + def is_unnamed(self, /) -> bool: + """Whether the entry addresses an unnamed :class:`PyStructSequence` slot.""" + # pylint: disable-next=import-outside-toplevel + from optree.typing import PyStructSequence_UnnamedField + + return self.field == PyStructSequence_UnnamedField + def __repr__(self, /) -> str: """Get the representation of the path entry.""" + if self.is_unnamed: # pragma: pypy no cover + # The marker is not a field name, so show the index it stands for. + return f'{self.__class__.__name__}(entry={self.entry!r}, type={self.type!r})' return f'{self.__class__.__name__}(field={self.field!r}, type={self.type!r})' def codify(self, /, node: str = '') -> str: """Generate code for accessing the path entry.""" + if self.is_unnamed: # pragma: pypy no cover + # An unnamed slot is only reachable by index. + return super().codify(node) return f'{node}.{self.field}' @@ -345,11 +361,31 @@ def init_fields(self, /) -> tuple[str, ...]: """Get the init field names.""" return tuple(f.name for f in dataclasses.fields(self.type) if f.init) + @property + def children_fields(self, /) -> tuple[str, ...]: + """Get the field names that an integer entry indexes. + + An integer entry addresses a child by position among those the registered flatten function + emitted. :func:`optree.dataclasses.register_node` records exactly which fields those are on + the class, so they are used verbatim. A class registered through the generic + :func:`optree.register_pytree_node` keeps no such record and its flatten function may emit + anything, so the declared field order is used instead. + """ + # pylint: disable-next=import-outside-toplevel,cyclic-import + from optree.dataclasses import _FIELDS + + # Read the class's own `__dict__`: a subclass of a registered dataclass inherits the + # attribute but not the registration. + registration = self.type.__dict__.get(_FIELDS) + if registration is None: + return self.fields + return tuple(registration[0]) + @property def field(self, /) -> str: """Get the field name.""" if isinstance(self.entry, int): - return self.init_fields[self.entry] + return self.children_fields[self.entry] return self.entry @property @@ -418,6 +454,10 @@ def __eq__(self, other: object, /) -> bool: """Check if the accessors are equal.""" return isinstance(other, PyTreeAccessor) and super().__eq__(other) + def __ne__(self, other: object, /) -> bool: + """Check if the accessors are not equal.""" + return not self == other + def __hash__(self, /) -> int: """Get the hash of the accessor.""" return super().__hash__() diff --git a/optree/dataclasses.py b/optree/dataclasses.py index 298397e9..e41b653f 100644 --- a/optree/dataclasses.py +++ b/optree/dataclasses.py @@ -268,6 +268,9 @@ def dataclass( # noqa: C901,D417 # pylint: disable=function-redefined ) -> _TypeT | Callable[[_TypeT], _TypeT]: """Dataclass decorator with PyTree integration. + See also :func:`register_node` for how instances are reconstructed and when a class needs an + explicit registration with custom flatten/unflatten functions instead. + Args: cls (type or None, optional): The class to decorate. If :data:`None`, return a decorator. namespace (str): The registry namespace used for the PyTree registration. @@ -502,6 +505,14 @@ def register_node( # noqa: C901 # pylint: disable=function-redefined,too-many-b :data:`True`) are treated as children, while init fields with ``metadata['pytree_node']`` set to :data:`False` are treated as metadata. + .. note:: + These generated functions cover the straightforward case where a class merely stores its + fields. Instances are rebuilt with ``cls(**fields)``, which re-runs ``__init__`` and any + ``__post_init__``. The :func:`tree_unflatten` round-trip is exact only when they leave each + field value unchanged from what it is given. Register a class that needs any other + reconstruction behavior explicitly with :func:`optree.register_pytree_node` or + :func:`optree.register_pytree_node_class` using custom flatten/unflatten functions. + Usage:: # Direct function call @@ -549,7 +560,10 @@ def decorator(cls: _TypeT, /) -> _TypeT: raise TypeError(f'{cls!r} is not a dataclass.') if _FIELDS in cls.__dict__: raise TypeError( - f'Cannot register {cls.__name__} as a pytree node more than once.', + f'Cannot register {cls.__name__} as a pytree node more than once with ' + f'`{__name__}.register_node()`. ' + 'Use `optree.register_pytree_node()` or `optree.register_pytree_node_class()` ' + 'with explicit flatten/unflatten functions to register it in a different namespace.', ) if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str): raise TypeError(f'The namespace must be a string, got {namespace!r}.') @@ -564,6 +578,25 @@ def decorator(cls: _TypeT, /) -> _TypeT: stacklevel=2, ) + # `InitVar` pseudo-fields are excluded from `dataclasses.fields()` (so they are neither children + # nor metadata) but remain required `__init__` parameters that `cls(**kwargs)` cannot restore. + # `__dataclass_fields__` includes them, tagged with the private `_FIELD_INITVAR` field type. + init_var_names = [ + name + # pylint: disable-next=protected-access + for name, dc_field in getattr(cls, dataclasses._FIELDS).items() # type: ignore[attr-defined] + # pylint: disable-next=protected-access + if dc_field._field_type is dataclasses._FIELD_INITVAR # type: ignore[attr-defined] + ] + if init_var_names: + raise TypeError( + f'Dataclass {cls.__name__!r} has `InitVar` field(s) {init_var_names!r}, which the ' + 'auto-generated flatten/unflatten functions cannot round-trip ' + f'(`{__name__}.register_node()` reconstructs instances with `cls(**kwargs)` and cannot ' + 'restore `InitVar` values). Use `optree.register_pytree_node()` or ' + '`optree.register_pytree_node_class()` with explicit flatten/unflatten functions instead.', + ) + children_fields = {} metadata_fields = {} for f in dataclasses.fields(cls): @@ -578,9 +611,6 @@ def decorator(cls: _TypeT, /) -> _TypeT: metadata_fields[f.name] = f children_field_names = tuple(children_fields) - children_fields_proxy = MappingProxyType(children_fields) - metadata_fields_proxy = MappingProxyType(metadata_fields) - setattr(cls, _FIELDS, (children_fields_proxy, metadata_fields_proxy)) def flatten_func( obj: _T, @@ -600,7 +630,8 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. kwargs.update(metadata) return cls(**kwargs) # type: ignore[return-value] - from optree.registry import register_pytree_node # pylint: disable=import-outside-toplevel + # pylint: disable-next=import-outside-toplevel + from optree.registry import register_pytree_node, unregister_pytree_node register_pytree_node( cls, # type: ignore[arg-type] @@ -609,4 +640,18 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. path_entry_type=DataclassEntry, namespace=namespace, ) + # Mark the class as registered only AFTER `register_pytree_node()` succeeds: `_FIELDS` is the + # "already registered" guard, so setting it before a failed registration would leave the class + # impossible to register ever again. The registration and the marker must land together: a + # registered but unmarked class cannot be registered again, so roll the registration back if + # setting the marker fails. + try: + setattr( + cls, + _FIELDS, + (MappingProxyType(children_fields), MappingProxyType(metadata_fields)), + ) + except BaseException: + unregister_pytree_node(cls, namespace=namespace) + raise return cls diff --git a/optree/integrations/attrs.py b/optree/integrations/attrs.py index 8a43d582..98ad6b02 100644 --- a/optree/integrations/attrs.py +++ b/optree/integrations/attrs.py @@ -143,11 +143,28 @@ def init_fields(self, /) -> tuple[str, ...]: """Get the init field names.""" return tuple(a.name for a in attrs.fields(self.type) if a.init) + @property + def children_fields(self, /) -> tuple[str, ...]: + """Get the field names that an integer entry indexes. + + An integer entry addresses a child by position among those the registered flatten function + emitted. :func:`optree.integrations.attrs.register_node` records exactly which fields those + are on the class, so they are used verbatim. A class registered through the generic + :func:`optree.register_pytree_node` keeps no such record and its flatten function may emit + anything, so the declared field order is used instead. + """ + # Read the class's own `__dict__`: a subclass of a registered attrs class inherits the + # attribute but not the registration. + registration = self.type.__dict__.get(_FIELDS) + if registration is None: + return self.fields + return tuple(registration[0]) + @property def field(self, /) -> str: """Get the field name.""" if isinstance(self.entry, int): - return self.init_fields[self.entry] + return self.children_fields[self.entry] return self.entry @property @@ -232,6 +249,9 @@ def define( # pylint: disable=function-redefined This is a wrapper around :func:`attrs.define` that also registers the class as a pytree node. + See also :func:`register_node` for how instances are reconstructed and when a class needs an + explicit registration with custom flatten/unflatten functions instead. + Args: cls (type or None, optional): The class to decorate. If :data:`None`, return a decorator. namespace (str): The registry namespace used for the PyTree registration. @@ -365,6 +385,15 @@ def register_node( # noqa: C901 # pylint: disable=function-redefined,too-many-b :data:`True`) are treated as children, while init fields with ``metadata['pytree_node']`` set to :data:`False` are treated as metadata. + .. note:: + These generated functions cover the straightforward case where a class merely stores its + fields. Instances are rebuilt with ``cls(**fields)``, which re-runs the attrs-generated + ``__init__`` and hence its field converters, validation, and ``__attrs_post_init__``. The + :func:`tree_unflatten` round-trip is exact only when ``__init__`` returns each field value + unchanged from what it is given. Register a class that needs any other reconstruction + behavior explicitly with :func:`optree.register_pytree_node` or + :func:`optree.register_pytree_node_class` using custom flatten/unflatten functions. + Usage:: # Direct function call @@ -413,7 +442,10 @@ def decorator(cls: _TypeT, /) -> _TypeT: raise TypeError(f'{cls!r} is not an attrs-decorated class.') if _FIELDS in cls.__dict__: raise TypeError( - f'Cannot register {cls.__name__} as a pytree node more than once.', + f'Cannot register {cls.__name__} as a pytree node more than once with ' + f'`{__name__}.register_node()`. ' + 'Use `optree.register_pytree_node()` or `optree.register_pytree_node_class()` ' + 'with explicit flatten/unflatten functions to register it in a different namespace.', ) if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str): raise TypeError(f'The namespace must be a string, got {namespace!r}.') @@ -444,9 +476,6 @@ def decorator(cls: _TypeT, /) -> _TypeT: children_field_names = tuple(children_fields) children_aliases = tuple(a.alias for a in children_fields.values()) - children_fields_proxy = MappingProxyType(children_fields) - metadata_fields_proxy = MappingProxyType(metadata_fields) - setattr(cls, _FIELDS, (children_fields_proxy, metadata_fields_proxy)) def flatten_func( obj: _T, @@ -466,7 +495,8 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. kwargs.update(metadata) return cls(**kwargs) # type: ignore[return-value] - from optree.registry import register_pytree_node # pylint: disable=import-outside-toplevel + # pylint: disable-next=import-outside-toplevel + from optree.registry import register_pytree_node, unregister_pytree_node register_pytree_node( cls, # type: ignore[arg-type] @@ -475,4 +505,18 @@ def unflatten_func(metadata: tuple[tuple[str, Any], ...], children: tuple[_U, .. path_entry_type=AttrsEntry, namespace=namespace, ) + # Mark the class as registered only AFTER `register_pytree_node()` succeeds: `_FIELDS` is the + # "already registered" guard, so setting it before a failed registration would leave the class + # impossible to register ever again. The registration and the marker must land together: a + # registered but unmarked class cannot be registered again, so roll the registration back if + # setting the marker fails. + try: + setattr( + cls, + _FIELDS, + (MappingProxyType(children_fields), MappingProxyType(metadata_fields)), + ) + except BaseException: + unregister_pytree_node(cls, namespace=namespace) + raise return cls diff --git a/optree/ops.py b/optree/ops.py index 8db8183c..3e08fd0d 100644 --- a/optree/ops.py +++ b/optree/ops.py @@ -23,11 +23,21 @@ import itertools import textwrap from collections import OrderedDict, defaultdict, deque -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, overload +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Generic, + SupportsIndex, + SupportsInt, + overload, +) import optree._C as _C from optree.accessors import PyTreeAccessor from optree.typing import NamedTuple, T, is_namedtuple_instance, is_structseq_instance +from optree.utils import total_order_sorted if TYPE_CHECKING: @@ -1320,7 +1330,9 @@ def tree_transpose_map( and ``xs`` is the tuple of values at corresponding nodes in ``rests``. """ leaves, outer_treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace) - if outer_treespec.num_leaves == 0: + # A leaf is only needed to infer the inner structure from the first output. With an explicit + # `inner_treespec`, a leafless outer structure transposes fine and `func` is never called. + if inner_treespec is None and outer_treespec.num_leaves == 0: raise ValueError(f'The outer structure must have at least one leaf. Got: {outer_treespec}.') flat_args = [leaves] + [outer_treespec.flatten_up_to(r) for r in rests] outputs = list(map(func, *flat_args)) @@ -1336,7 +1348,7 @@ def tree_transpose_map( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) + transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1407,7 +1419,7 @@ def tree_transpose_map_with_path( leaf in ``tree`` and ``xs`` is the tuple of values at corresponding nodes in ``rests``. """ # pylint: disable=line-too-long paths, leaves, outer_treespec = _C.flatten_with_path(tree, is_leaf, none_is_leaf, namespace) - if outer_treespec.num_leaves == 0: + if inner_treespec is None and outer_treespec.num_leaves == 0: raise ValueError(f'The outer structure must have at least one leaf. Got: {outer_treespec}.') flat_args = [leaves] + [outer_treespec.flatten_up_to(r) for r in rests] outputs = list(map(func, paths, *flat_args)) @@ -1423,7 +1435,7 @@ def tree_transpose_map_with_path( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) + transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1521,7 +1533,7 @@ def tree_transpose_map_with_accessor( leaf in ``tree`` and ``xs`` is the tuple of values at corresponding nodes in ``rests``. """ # pylint: disable=line-too-long leaves, outer_treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace) - if outer_treespec.num_leaves == 0: + if inner_treespec is None and outer_treespec.num_leaves == 0: raise ValueError(f'The outer structure must have at least one leaf. Got: {outer_treespec}.') flat_args = [leaves] + [outer_treespec.flatten_up_to(r) for r in rests] outputs = list(map(func, outer_treespec.accessors(), *flat_args)) @@ -1537,7 +1549,7 @@ def tree_transpose_map_with_accessor( raise ValueError(f'The inner structure must have at least one leaf. Got: {inner_treespec}.') grouped = [inner_treespec.flatten_up_to(o) for o in outputs] - transposed = zip(*grouped) + transposed = zip(*grouped) if grouped else [()] * inner_treespec.num_leaves subtrees = map(outer_treespec.unflatten, transposed) return inner_treespec.unflatten(subtrees) # type: ignore[arg-type] @@ -1706,6 +1718,64 @@ def add_leaves(x: T, subtree: PyTree[S]) -> None: return result +def _tree_broadcast_common_with_treespec( + tree: PyTree[T], + /, + *rests: PyTree[T], + is_leaf: Callable[[T], bool] | None = None, + none_is_leaf: bool = False, + namespace: str = '', +) -> tuple[tuple[PyTree[T], ...], PyTreeSpec]: + """Broadcast to a common suffix and return the trees together with that common structure. + + The structure is derived from the input trees, where ``is_leaf`` describes the caller's data. + Re-deriving it from the broadcast result instead would let a predicate that classifies by leaf + value treat a subtree broadcasting just created as a leaf, silently mapping over fewer leaves + than the common suffix has. Each input is flattened once and broadcast straight to the common + suffix of all of them. + """ + if not rests: # pragma: no cover + return (tree,), tree_structure( + tree, + is_leaf=is_leaf, + none_is_leaf=none_is_leaf, + namespace=namespace, + ) + + flattened = [_C.flatten(t, is_leaf, none_is_leaf, namespace) for t in (tree, *rests)] + common_suffix_treespec: PyTreeSpec = flattened[0][1] + for _, treespec in flattened[1:]: + common_suffix_treespec = common_suffix_treespec.broadcast_to_common_suffix(treespec) + + sentinel: T = object() # type: ignore[assignment] + common_suffix_tree: PyTree[T] = common_suffix_treespec.unflatten( + itertools.repeat(sentinel, common_suffix_treespec.num_leaves), + ) + + def broadcast_leaves(x: T, subtree: PyTree[T]) -> PyTree[T]: + # `subtree`'s leaves are the private `sentinel`, not real values, so the user's `is_leaf` + # must not see them: its structure is already fixed by `common_suffix_treespec`. + subtreespec = tree_structure( + subtree, + is_leaf=None, + none_is_leaf=none_is_leaf, + namespace=namespace, + ) + return subtreespec.unflatten(itertools.repeat(x, subtreespec.num_leaves)) + + broadcasted: tuple[PyTree[T], ...] = tuple( + treespec.unflatten( + map( + broadcast_leaves, # type: ignore[arg-type] + leaves, + treespec.flatten_up_to(common_suffix_tree), + ), + ) + for leaves, treespec in flattened + ) + return broadcasted, common_suffix_treespec + + def tree_broadcast_common( tree: PyTree[T], other_tree: PyTree[T], @@ -1722,10 +1792,18 @@ def tree_broadcast_common( If a ``suffix_tree`` is a suffix of a ``tree``, this means the ``suffix_tree`` can be constructed by replacing the leaves of ``tree`` with appropriate **subtrees**. - This function returns two pytrees with the same structure. The tree structure is the common - suffix structure of ``tree`` and ``other_tree``. The leaves are replicated from ``tree`` and - ``other_tree``. The number of replicas is determined by the corresponding subtree in the suffix - structure. + This function returns two pytrees broadcasted to the common suffix structure of ``tree`` and + ``other_tree``. Each is valid input to that common treespec, but the two need not share a + treespec with each other: a dictionary node keeps the key order and node type of its own input. + The leaves are replicated from ``tree`` and ``other_tree``. The number of replicas is determined + by the corresponding subtree in the suffix structure. + + .. note:: + If ``is_leaf`` classifies nodes by their leaf **value** rather than by type or structure + (e.g., treating an integer tuple as a leaf), broadcasting may replicate a value into a slot + whose filled form the predicate then re-classifies, yielding two trees that re-flatten to + different structures under the same ``is_leaf``. Prefer type- or structure-based predicates + when broadcasting. >>> tree_broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) @@ -1765,37 +1843,12 @@ def tree_broadcast_common( Returns: Two pytrees of common suffix structure of ``tree`` and ``other_tree`` with broadcasted subtrees. """ - leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace) - other_leaves, other_treespec = _C.flatten(other_tree, is_leaf, none_is_leaf, namespace) - common_suffix_treespec = treespec.broadcast_to_common_suffix(other_treespec) - - sentinel: T = object() # type: ignore[assignment] - common_suffix_tree: PyTree[T] = common_suffix_treespec.unflatten( - itertools.repeat(sentinel, common_suffix_treespec.num_leaves), - ) - - def broadcast_leaves(x: T, subtree: PyTree[T]) -> PyTree[T]: - subtreespec = tree_structure( - subtree, - is_leaf=is_leaf, - none_is_leaf=none_is_leaf, - namespace=namespace, - ) - return subtreespec.unflatten(itertools.repeat(x, subtreespec.num_leaves)) - - broadcasted_tree: PyTree[T] = treespec.unflatten( - map( - broadcast_leaves, # type: ignore[arg-type] - leaves, - treespec.flatten_up_to(common_suffix_tree), - ), - ) - other_broadcasted_tree: PyTree[T] = other_treespec.unflatten( - map( - broadcast_leaves, # type: ignore[arg-type] - other_leaves, - other_treespec.flatten_up_to(common_suffix_tree), - ), + (broadcasted_tree, other_broadcasted_tree), _ = _tree_broadcast_common_with_treespec( + tree, + other_tree, + is_leaf=is_leaf, + none_is_leaf=none_is_leaf, + namespace=namespace, ) return broadcasted_tree, other_broadcasted_tree @@ -1816,10 +1869,10 @@ def broadcast_common( If a ``suffix_tree`` is a suffix of a ``tree``, this means the ``suffix_tree`` can be constructed by replacing the leaves of ``tree`` with appropriate **subtrees**. - This function returns two pytrees with the same structure. The tree structure is the common - suffix structure of ``tree`` and ``other_tree``. The leaves are replicated from ``tree`` and - ``other_tree``. The number of replicas is determined by the corresponding subtree in the suffix - structure. + This function returns two lists of leaves of the same length, both in the leaf order of the + common suffix structure of ``tree`` and ``other_tree``. The leaves are replicated from ``tree`` + and ``other_tree``. The number of replicas is determined by the corresponding subtree in the + suffix structure. >>> broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) @@ -1858,67 +1911,20 @@ def broadcast_common( Two lists of leaves in ``tree`` and ``other_tree`` broadcasted to match the number of leaves in the common suffix structure. """ # pylint: disable=line-too-long - broadcasted_tree, other_broadcasted_tree = tree_broadcast_common( + (broadcasted_tree, other_broadcasted_tree), treespec = _tree_broadcast_common_with_treespec( tree, other_tree, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, ) - - broadcasted_leaves: list[T] = [] - other_broadcasted_leaves: list[T] = [] - - def add_leaves(x: T, y: T) -> None: - broadcasted_leaves.append(x) - other_broadcasted_leaves.append(y) - - tree_map_( - add_leaves, - broadcasted_tree, - other_broadcasted_tree, - is_leaf=is_leaf, - none_is_leaf=none_is_leaf, - namespace=namespace, + # The common suffix treespec bottoms out at the caller's leaves, so these are `T`, not subtrees. + return ( # type: ignore[return-value] + treespec.flatten_up_to(broadcasted_tree), + treespec.flatten_up_to(other_broadcasted_tree), ) - return broadcasted_leaves, other_broadcasted_leaves - -def _tree_broadcast_common( - tree: PyTree[T], - /, - *rests: PyTree[T], - is_leaf: Callable[[T], bool] | None = None, - none_is_leaf: bool = False, - namespace: str = '', -) -> tuple[PyTree[T], ...]: - if not rests: - return (tree,) - if len(rests) == 1: - return tree_broadcast_common( - tree, - rests[0], - is_leaf=is_leaf, - none_is_leaf=none_is_leaf, - namespace=namespace, - ) - - broadcasted_tree = tree - broadcasted_rests = list(rests) - for _ in range(2): - for i, rest in enumerate(rests): - broadcasted_tree, broadcasted_rests[i] = tree_broadcast_common( - broadcasted_tree, - rest, - is_leaf=is_leaf, - none_is_leaf=none_is_leaf, - namespace=namespace, - ) - return (broadcasted_tree, *broadcasted_rests) - - -# pylint: disable-next=too-many-locals def tree_broadcast_map( func: Callable[..., U], tree: PyTree[T], @@ -1974,22 +1980,26 @@ def tree_broadcast_map( corresponding leaf (may be broadcasted) in ``tree`` and ``xs`` is the tuple of values at corresponding leaves (may be broadcasted) in ``rests``. """ - return tree_map( - func, - *_tree_broadcast_common( + if not rests: + return tree_map( + func, tree, - *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, - ), + ) + + broadcasted, treespec = _tree_broadcast_common_with_treespec( + tree, + *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, ) + flat_args = [treespec.flatten_up_to(broadcasted_tree) for broadcasted_tree in broadcasted] + return treespec.unflatten(map(func, *flat_args)) -# pylint: disable-next=too-many-locals def tree_broadcast_map_with_path( func: Callable[..., U], tree: PyTree[T], @@ -2053,19 +2063,24 @@ def tree_broadcast_map_with_path( value at the corresponding leaf (may be broadcasted) in ``tree`` and ``xs`` is the tuple of values at corresponding leaves (may be broadcasted) in ``rests``. """ - return tree_map_with_path( - func, - *_tree_broadcast_common( + if not rests: + return tree_map_with_path( + func, tree, - *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, - ), + ) + + broadcasted, treespec = _tree_broadcast_common_with_treespec( + tree, + *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, ) + flat_args = [treespec.flatten_up_to(broadcasted_tree) for broadcasted_tree in broadcasted] + return treespec.unflatten(map(func, treespec.paths(), *flat_args)) def tree_broadcast_map_with_accessor( @@ -2146,19 +2161,24 @@ def tree_broadcast_map_with_accessor( and value at the corresponding leaf (may be broadcasted) in ``tree`` and ``xs`` is the tuple of values at corresponding leaves (may be broadcasted) in ``rests``. """ - return tree_map_with_accessor( - func, - *_tree_broadcast_common( + if not rests: + return tree_map_with_accessor( + func, tree, - *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, - ), + ) + + broadcasted, treespec = _tree_broadcast_common_with_treespec( + tree, + *rests, is_leaf=is_leaf, none_is_leaf=none_is_leaf, namespace=namespace, ) + flat_args = [treespec.flatten_up_to(broadcasted_tree) for broadcasted_tree in broadcasted] + return treespec.unflatten(map(func, treespec.accessors(), *flat_args)) # pylint: disable-next=missing-class-docstring,too-few-public-methods @@ -2761,7 +2781,7 @@ def treespec_entries(treespec: PyTreeSpec, /) -> list[Any]: return treespec.entries() -def treespec_entry(treespec: PyTreeSpec, index: int, /) -> Any: +def treespec_entry(treespec: PyTreeSpec, index: SupportsInt | SupportsIndex, /) -> Any: """Return the entry of a treespec at the given index. See also :func:`treespec_entries`, :func:`treespec_children`, and :meth:`PyTreeSpec.entry`. @@ -2784,7 +2804,7 @@ def treespec_children(treespec: PyTreeSpec, /) -> list[PyTreeSpec]: return treespec.children() -def treespec_child(treespec: PyTreeSpec, index: int, /) -> PyTreeSpec: +def treespec_child(treespec: PyTreeSpec, index: SupportsInt | SupportsIndex, /) -> PyTreeSpec: """Return the treespec of the child of a treespec at the given index. See also :func:`treespec_children`, :func:`treespec_entries`, and :meth:`PyTreeSpec.child`. @@ -2974,10 +2994,68 @@ def treespec_is_prefix( *, strict: bool = False, ) -> bool: - """Return whether ``treespec`` is a prefix of ``other_treespec``. + r"""Return whether ``treespec`` is a prefix of ``other_treespec``. + + A treespec is a *prefix* of another when the latter can be built by replacing some of the + former's leaves with subtrees. This is the broadcasting-compatibility relation behind + multi-input :func:`tree_map` and :func:`tree_broadcast_prefix`: a prefix treespec broadcasts + over any tree it is a prefix of, pairing each of its leaves with the corresponding subtree of + the fuller structure. See also :func:`treespec_is_suffix` and :meth:`PyTreeSpec.is_prefix`. + The prefix relation, together with :func:`treespec_is_suffix` and the ``<``, ``<=``, ``>``, + ``>=`` operators (``<`` and ``>`` are the ``strict=True`` variants), is a **preorder**: + reflexive (every treespec is a non-strict prefix of itself) and transitive, but **neither a + partial order nor a total order**. + + - Not antisymmetric (so not a partial order): ``a <= b and b <= a`` does **not** imply + ``a == b``. Metadata that does not change how children are partitioned, such as a + :class:`collections.deque` ``maxlen`` or a :class:`dict` key order, is transparent to the + prefix relation but significant to equality, so they are mutual prefixes yet unequal. + - Not total: two treespecs can be incomparable, with neither a prefix of the other. Examples + are a :class:`tuple` and a same-arity :class:`list`, or :class:`dict`\s with different key + sets. + + >>> treespec_is_prefix(tree_structure(1), tree_structure([1, 2, 3])) + True + >>> treespec_is_prefix(tree_structure([1, 2, 3]), tree_structure(1)) + False + + Mutual prefixes need not be equal, since the relation is a preorder, not a partial order. A + :class:`collections.deque` ``maxlen`` is transparent to the prefix relation but significant to + equality: + + >>> from collections import deque + >>> a = tree_structure(deque([1, 2], maxlen=2)) + >>> b = tree_structure(deque([1, 2], maxlen=5)) + >>> treespec_is_prefix(a, b) and treespec_is_prefix(b, a) + True + >>> a <= b and b <= a + True + >>> a == b + False + + Two treespecs can be incomparable, since the relation is not a total order. A :class:`tuple` + and a same-arity :class:`list` are neither a prefix of the other: + + >>> a = tree_structure((1, 2)) + >>> b = tree_structure([1, 2]) + >>> treespec_is_prefix(a, b) + False + >>> treespec_is_prefix(b, a) + False + >>> a <= b + False + >>> b <= a + False + >>> a < b + False + >>> b < a + False + >>> a == b + False + Args: treespec (PyTreeSpec): A treespec. other_treespec (PyTreeSpec): Another treespec to compare against. @@ -2999,8 +3077,19 @@ def treespec_is_suffix( ) -> bool: """Return whether ``treespec`` is a suffix of ``other_treespec``. + This is the reverse of :func:`treespec_is_prefix`: ``treespec_is_suffix(a, b)`` is equivalent to + ``treespec_is_prefix(b, a)``. Like the prefix relation, it is a **preorder**: reflexive and + transitive, but neither a partial order (mutual suffixes need not be equal) nor a total order + (two treespecs can be incomparable). See :func:`treespec_is_prefix` for the full discussion and + examples. + See also :func:`treespec_is_prefix` and :meth:`PyTreeSpec.is_suffix`. + >>> treespec_is_suffix(tree_structure([1, 2, 3]), tree_structure(1)) + True + >>> treespec_is_suffix(tree_structure(1), tree_structure([1, 2, 3])) + False + Args: treespec (PyTreeSpec): A treespec. other_treespec (PyTreeSpec): Another treespec to compare against. @@ -3615,12 +3704,7 @@ def helper( # pylint: disable=too-many-locals none_is_leaf=none_is_leaf, namespace=namespace, ) - full_tree_one_level_output = ( - full_tree_children, - full_tree_metadata, - full_tree_entries, - _, - ) = tree_flatten_one_level( + full_tree_children, full_tree_metadata, _, _ = tree_flatten_one_level( full_subtree, none_is_leaf=none_is_leaf, namespace=namespace, @@ -3640,13 +3724,13 @@ def helper( # pylint: disable=too-many-locals prefix_tree_keys_set = set(prefix_tree_keys) full_tree_keys_set = set(full_tree_keys) if prefix_tree_keys_set != full_tree_keys_set: - missing_keys = sorted(prefix_tree_keys_set.difference(full_tree_keys_set)) - extra_keys = sorted(full_tree_keys_set.difference(prefix_tree_keys_set)) + missing_keys = prefix_tree_keys_set.difference(full_tree_keys_set) + extra_keys = full_tree_keys_set.difference(prefix_tree_keys_set) key_difference = '' if missing_keys: - key_difference += f'\nmissing key(s):\n {missing_keys}' + key_difference += f'\nmissing key(s):\n {total_order_sorted(missing_keys)}' if extra_keys: - key_difference += f'\nextra key(s):\n {extra_keys}' + key_difference += f'\nextra key(s):\n {total_order_sorted(extra_keys)}' yield lambda name: ValueError( f'pytree structure error: different pytree keys at key path\n' f' {accessor.codify(name) if accessor else name + " tree root"}\n' @@ -3712,7 +3796,8 @@ def helper( # pylint: disable=too-many-locals return # don't look for more errors in this subtree # If the root types and numbers of children agree, there must be an error in a subtree, - # so recurse: + # so recurse. A custom node type may report per-instance entries that differ from the full + # tree's while its metadata does not; `broadcast_prefix` accepts that, so it is not an error. entries = [ prefix_tree_one_level_output.path_entry_type( e, @@ -3721,18 +3806,6 @@ def helper( # pylint: disable=too-many-locals ) for e in prefix_tree_entries ] - entries_ = [ - full_tree_one_level_output.path_entry_type( - e, - full_tree_type, - full_tree_one_level_output.kind, - ) - for e in full_tree_entries - ] - assert ( - both_standard_dict # special handling for dictionary types already done in the keys check above - or entries == entries_ - ), f'equal pytree nodes gave different keys: {entries} and {entries_}' # pylint: disable-next=invalid-name for e, t1, t2 in zip(entries, prefix_tree_children, full_tree_children): yield from helper(accessor + e, t1, t2) # type: ignore[arg-type] diff --git a/optree/typing.py b/optree/typing.py index 77dc4ce4..eb2cec14 100644 --- a/optree/typing.py +++ b/optree/typing.py @@ -418,7 +418,7 @@ def is_namedtuple_instance(obj: object, /) -> bool: @_override_with_(_C.is_namedtuple_class) -def is_namedtuple_class(cls: type, /) -> bool: +def is_namedtuple_class(cls: object, /) -> bool: """Return whether the class is a subclass of namedtuple.""" return ( isinstance(cls, type) @@ -525,7 +525,7 @@ def is_structseq_instance(obj: object, /) -> bool: @_override_with_(_C.is_structseq_class) -def is_structseq_class(cls: type, /) -> bool: +def is_structseq_class(cls: object, /) -> bool: """Return whether the class is a class of PyStructSequence.""" if ( isinstance(cls, type) @@ -550,6 +550,11 @@ def is_structseq_class(cls: type, /) -> bool: # pylint: disable-next=line-too-long StructSequenceFieldType: type[types.MemberDescriptorType] = type(type(sys.version_info).major) # type: ignore[assignment] +# The name reported for an unnamed PyStructSequence slot; CPython's C-level marker (not a valid +# identifier, so accessors fall back to index access for such slots). +# pylint: disable-next=invalid-name +PyStructSequence_UnnamedField: str = _C.PyStructSequence_UnnamedField # 'unnamed field' + @_override_with_(_C.structseq_fields) def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: @@ -563,20 +568,45 @@ def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: if not is_structseq_class(cls): raise TypeError(f'Expected an instance of PyStructSequence type, got {obj!r}.') + n_sequence_fields: int = cls.n_sequence_fields # type: ignore[attr-defined] + n_unnamed_fields: int = cls.n_unnamed_fields # type: ignore[attr-defined] + if platform.python_implementation() == 'PyPy': # pragma: pypy cover - indices_by_name = { - name: member.index # type: ignore[attr-defined] + # PyPy has no unnamed sequence fields: a field descriptor exposes `.index` as its sequence + # position, and hidden fields have an index >= n_sequence_fields. (`n_unnamed_fields == 0`, + # see PyPy's `lib_pypy/_structseq.py`) Map index -> name and defensively fill any missing + # (i.e. unnamed) sequence slot with the marker, should that invariant ever change. + names_by_index: dict[int, str] = { + member.index: name # type: ignore[attr-defined] for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) } - fields = sorted(indices_by_name, key=indices_by_name.get) # type: ignore[arg-type] else: # pragma: pypy no cover - fields = [ + # CPython's `member_descriptor` hides the field offset, so positions are unreadable in pure + # Python. `vars()` yields members in field order, and only sequence slots can be unnamed. + named = [ name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) - ] - return tuple(fields[: cls.n_sequence_fields]) # type: ignore[attr-defined] + ][: n_sequence_fields - n_unnamed_fields] + + positions: Iterable[int] = range(n_sequence_fields) + if n_unnamed_fields > 0: + # Unnamed slots may sit anywhere, not only at the tail, so recover each named field's + # position from a probe whose slots hold distinct sentinels and match them by identity. + sentinels = [object() for _ in range(n_sequence_fields)] + try: + probe = cls(sentinels) + index_of = {id(sentinel): index for index, sentinel in enumerate(sentinels)} + positions = [index_of[id(getattr(probe, name))] for name in named] + except (TypeError, ValueError, KeyError, AttributeError): # pragma: no cover + pass # the type rejects placeholder values, fall back to assuming a trailing layout + names_by_index = dict(zip(positions, named)) + + return tuple( + names_by_index.get(index, PyStructSequence_UnnamedField) + for index in range(n_sequence_fields) + ) del _tp_cache diff --git a/optree/utils.py b/optree/utils.py index 7622f5be..1878757e 100644 --- a/optree/utils.py +++ b/optree/utils.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Iterable, Sequence +from operator import itemgetter from typing import TYPE_CHECKING, Any, Callable, overload @@ -37,30 +38,30 @@ def total_order_sorted( types of keys. """ sequence = list(iterable) + # Apply `key` up front: it runs exactly once per element, and a `TypeError` from the callback + # propagates instead of being mistaken for a comparison failure and swallowed below. + keys: list[Any] = sequence if key is None else [key(x) for x in sequence] + decorated = list(zip(keys, sequence)) + def by_type_and_key(pair: tuple[Any, T], /) -> tuple[str, Any]: + y = pair[0] + return (f'{y.__class__.__module__}.{y.__class__.__qualname__}', y) + + # `list.sort()` leaves the list partially reordered when a comparison raises, so each attempt + # sorts a copy (`sorted()`) and only a fully sorted one is committed. try: # Sort directly if possible - return sorted(sequence, key=key, reverse=reverse) # type: ignore[type-var,arg-type] + ordered = sorted(decorated, key=itemgetter(0), reverse=reverse) except TypeError: - if key is None: - - def key_fn(x: T) -> tuple[str, Any]: - return (f'{x.__class__.__module__}.{x.__class__.__qualname__}', x) - - else: - - def key_fn(x: T) -> tuple[str, Any]: - y = key(x) - return (f'{y.__class__.__module__}.{y.__class__.__qualname__}', y) - try: # Add `{obj.__class__.__module__}.{obj.__class__.__qualname__}` to the key order to make # it sortable between different types (e.g., `int` vs. `str`) - return sorted(sequence, key=key_fn, reverse=reverse) + ordered = sorted(decorated, key=by_type_and_key, reverse=reverse) except TypeError: # cannot sort the keys (e.g., user-defined types) - # Fall back to the original order. `reverse` is intentionally not applied: like a stable - # sort, incomparable elements (treated as "all equal") keep their original order. + # Keep the insertion order. `reverse` is intentionally not applied: like a stable sort, + # incomparable elements (treated as "all equal") keep their input order. return sequence + return [x for _, x in ordered] @overload diff --git a/src/optree.cpp b/src/optree.cpp index f1dd9fd0..e724a7ec 100644 --- a/src/optree.cpp +++ b/src/optree.cpp @@ -51,6 +51,7 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] mod.doc() = "Optimized PyTree Utilities. (C extension module built from " + std::string(__FILE_RELPATH_FROM_PROJECT_ROOT__) + ")"; mod.attr("Py_TPFLAGS_BASETYPE") = py::int_(Py_TPFLAGS_BASETYPE); + mod.attr("PyStructSequence_UnnamedField") = py::str(PyStructSequenceUnnamedField()); // Meta information during build py::dict BUILDTIME_METADATA{}; @@ -504,13 +505,32 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] "Return a string representation of the treespec.") .def_method_pos_only("__hash__", &PyTreeSpec::HashValue, "Return the hash of the treespec.") .def_method_pos_only("__len__", &PyTreeSpec::GetNumLeaves, "Number of leaves in the tree.") + // Known limitation: pybind11's `tp_new` only allocates the wrapper, so between + // `cls.__new__(cls)` and `__setstate__` a method call reads uninitialized memory (undefined + // behavior, and `PYTREESPEC_SANITY_CHECK` cannot catch it). Only the GC is covered, by the + // `is_holder_constructed` guards in `PyTpTraverse` / `PyTpClear`. .def(py::pickle([](const PyTreeSpec &t) -> py::object { return t.ToPicklable(); }, [](const py::object &o) -> std::unique_ptr { return PyTreeSpec::FromPicklable(o); }), "Serialization support for PyTreeSpec.", py::arg("state"), - py::pos_only()); + py::pos_only()) + .def( + "__reduce__", + [](const py::handle &self) -> py::object { + // pybind11's pickle support (`__getstate__`/`__setstate__`) reconstructs via the + // protocol >= 2 `copyreg.__newobj__` reduction. At protocol 0/1 the default + // reduction goes through `object.__new__`, which pybind11 rejects with an + // untranslated C++ exception that aborts the interpreter. Return the `__newobj__` + // reduction explicitly so every protocol reconstructs via `cls.__new__(cls)`. + const py::object newobj = py::module_::import("copyreg").attr("__newobj__"); + return py::make_tuple(newobj, + py::make_tuple(py::type::handle_of(self)), + self.attr("__getstate__")()); + }, + "Reduce the treespec to a picklable form supporting all pickle protocols.", + py::pos_only()); auto PyTreeIterTypeObject = #if defined(PYBIND11_HAS_INTERNALS_WITH_SMART_HOLDER_SUPPORT) @@ -543,6 +563,8 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] py::arg("leaf_predicate") = std::nullopt, py::arg("none_is_leaf") = false, py::arg("namespace") = "") + // Known limitation: `PyTreeIter.__new__(PyTreeIter)` leaves the C++ iterator unbuilt, so + // the methods below would read uninitialized memory (see the `PyTreeSpec` note above). .def_method_pos_only("__iter__", &PyTreeIter::Iter, "Return the iterator object itself.") .def_method_pos_only("__next__", &PyTreeIter::Next, "Return the next leaf in the pytree."); diff --git a/src/registry.cpp b/src/registry.cpp index 5c12296c..34fef8d6 100644 --- a/src/registry.cpp +++ b/src/registry.cpp @@ -22,8 +22,6 @@ limitations under the License. #include // std::remove_const_t #include // std::move, std::make_pair -#include - #include "optree/optree.h" namespace optree { @@ -69,9 +67,8 @@ template template PyTreeTypeRegistry &PyTreeTypeRegistry::GetSingleton(); template PyTreeTypeRegistry &PyTreeTypeRegistry::GetSingleton(); -ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_namespace) const { - const scoped_read_lock lock{sm_mutex}; - +ssize_t PyTreeTypeRegistry::SizeImpl(const std::optional ®istry_namespace) const { + // The caller must hold `sm_mutex`. ssize_t count = py::ssize_t_cast(m_registrations.size()); for (const auto &[named_type, _] : m_named_registrations) { if (!registry_namespace || named_type.first == *registry_namespace) [[likely]] { @@ -81,17 +78,20 @@ ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_name return count; } -template -/*static*/ void PyTreeTypeRegistry::RegisterImpl(const py::object &cls, - const py::function &flatten_func, - const py::function &unflatten_func, - const py::object &path_entry_type, - const std::string ®istry_namespace) { - auto ®istry = GetSingleton(); - - if (registry.m_builtins_types.find(cls) != registry.m_builtins_types.end()) [[unlikely]] { - throw py::value_error("PyTree type " + PyRepr(cls) + - " is a built-in type and cannot be re-registered."); +ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_namespace) const { + const scoped_read_lock lock{sm_mutex}; + return SizeImpl(registry_namespace); +} + +// The caller must hold `sm_mutex` in write mode. No Python may run here; see `RegistryStatus`. +PyTreeTypeRegistry::RegistryStatus PyTreeTypeRegistry::RegisterImpl( + const py::object &cls, + const py::function &flatten_func, + const py::function &unflatten_func, + const py::object &path_entry_type, + const std::string ®istry_namespace) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { + return RegistryStatus::BuiltinType; } auto registration = std::make_shared>(); @@ -100,59 +100,20 @@ template registration->flatten_func = py::reinterpret_borrow(flatten_func); registration->unflatten_func = py::reinterpret_borrow(unflatten_func); registration->path_entry_type = py::reinterpret_borrow(path_entry_type); + // The registration only ever borrows objects the caller keeps alive, so the drop below when the + // insert fails cannot reach zero and cannot run Python. if (registry_namespace.empty()) [[unlikely]] { - if (!registry.m_registrations.emplace(cls, std::move(registration)).second) [[unlikely]] { - throw py::value_error("PyTree type " + PyRepr(cls) + - " is already registered in the global namespace."); - } - if (IsStructSequenceClass(cls)) [[unlikely]] { - PyErr_WarnEx(PyExc_UserWarning, - ("PyTree type " + PyRepr(cls) + - " is a class of `PyStructSequence`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions.") - .c_str(), - /*stack_level=*/2); - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - PyErr_WarnEx(PyExc_UserWarning, - ("PyTree type " + PyRepr(cls) + - " is a subclass of `collections.namedtuple`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions.") - .c_str(), - /*stack_level=*/2); + if (!m_registrations.emplace(cls, std::move(registration)).second) [[unlikely]] { + return RegistryStatus::AlreadyRegistered; } } else [[likely]] { - if (!registry.m_named_registrations + if (!m_named_registrations .emplace(std::make_pair(registry_namespace, cls), std::move(registration)) .second) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) << " is already registered in namespace " - << PyRepr(registry_namespace) << "."; - throw py::value_error(oss.str()); - } - if (IsStructSequenceClass(cls)) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) - << " is a class of `PyStructSequence`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions in namespace " - << PyRepr(registry_namespace) << "."; - PyErr_WarnEx(PyExc_UserWarning, - oss.str().c_str(), - /*stack_level=*/2); - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) - << " is a subclass of `collections.namedtuple`, " - "which is already registered in the global namespace. " - "Override it with custom flatten/unflatten functions in namespace " - << PyRepr(registry_namespace) << "."; - PyErr_WarnEx(PyExc_UserWarning, - oss.str().c_str(), - /*stack_level=*/2); + return RegistryStatus::AlreadyRegistered; } } + return RegistryStatus::Ok; } /*static*/ void PyTreeTypeRegistry::Register(const py::object &cls, @@ -160,88 +121,174 @@ template const py::function &unflatten_func, const py::object &path_entry_type, const std::string ®istry_namespace) { - const scoped_write_lock lock{sm_mutex}; - - RegisterImpl(cls, - flatten_func, - unflatten_func, - path_entry_type, - registry_namespace); - RegisterImpl(cls, - flatten_func, - unflatten_func, - path_entry_type, - registry_namespace); - cls.inc_ref(); - flatten_func.inc_ref(); - unflatten_func.inc_ref(); - path_entry_type.inc_ref(); + // Classify the type BEFORE taking `sm_mutex`: `IsStructSequenceClass` / `IsNamedTupleClass` run + // Python and release the GIL, and doing that under the write lock inverts the GIL <-> + // `sm_mutex` lock order against a concurrent flatten (mirrors `Unregister`). + const char *overridden_kind = nullptr; + if (IsStructSequenceClass(cls)) [[unlikely]] { + overridden_kind = " is a class of `PyStructSequence`, "; + } else if (IsNamedTupleClass(cls)) [[unlikely]] { + overridden_kind = " is a subclass of `collections.namedtuple`, "; + } + + // Acquire both singletons BEFORE `sm_mutex`, mirroring `Init`/`Clear`. Under + // `per_interpreter_gil`, `GetSingleton()` releases the GIL on every call once a subinterpreter + // has existed; doing that while holding `sm_mutex` inverts the GIL <-> `sm_mutex` lock order + // against a concurrent flatten (read lock) and deadlocks. + auto ®istry1 = GetSingleton(); + auto ®istry2 = GetSingleton(); + + RegistryStatus status = RegistryStatus::Ok; + { + const scoped_write_lock lock{sm_mutex}; + + status = registry1.RegisterImpl(cls, + flatten_func, + unflatten_func, + path_entry_type, + registry_namespace); + if (status == RegistryStatus::Ok) [[likely]] { + status = registry2.RegisterImpl(cls, + flatten_func, + unflatten_func, + path_entry_type, + registry_namespace); + } + if (status == RegistryStatus::Ok) [[likely]] { + cls.inc_ref(); + flatten_func.inc_ref(); + unflatten_func.inc_ref(); + path_entry_type.inc_ref(); + } + } + + // Format the error only after the lock is released: `PyRepr` runs the (meta)class `__repr__` as + // Python bytecode, which can hand off the GIL to a thread blocking on `sm_mutex` in read mode. + if (status != RegistryStatus::Ok) [[unlikely]] { + if (status == RegistryStatus::BuiltinType) [[unlikely]] { + throw py::value_error("PyTree type " + PyRepr(cls) + + " is a built-in type and cannot be re-registered."); + } + std::ostringstream oss{}; + oss << "PyTree type " << PyRepr(cls) << " is already registered in "; + if (registry_namespace.empty()) [[unlikely]] { + oss << "the global namespace."; + } else [[likely]] { + oss << "namespace " << PyRepr(registry_namespace) << "."; + } + throw py::value_error(oss.str()); + } + + // Warn only once the registration succeeded: a rejected one overrides nothing. `PyErr_WarnEx` + // runs Python, so it must not run under `sm_mutex` either. Under warnings-as-errors it raises, + // so undo the registration to keep `Register` atomic. + if (overridden_kind != nullptr) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTree type " << PyRepr(cls) << overridden_kind + << "which is already registered in the global namespace. " + "Override it with custom flatten/unflatten functions"; + if (!registry_namespace.empty()) [[likely]] { + oss << " in namespace " << PyRepr(registry_namespace); + } + oss << "."; + try { + if (PyErr_WarnEx(PyExc_UserWarning, oss.str().c_str(), /*stack_level=*/2) < 0) + [[unlikely]] { + throw py::error_already_set(); + } + } catch (...) { + Unregister(cls, registry_namespace); + throw; + } + } } -template -/*static*/ PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::UnregisterImpl( +// The caller must hold `sm_mutex` in write mode. No Python may run here; see `RegistryStatus`. +PyTreeTypeRegistry::RegistryStatus PyTreeTypeRegistry::UnregisterImpl( const py::object &cls, - const std::string ®istry_namespace) { - auto ®istry = GetSingleton(); - - if (registry.m_builtins_types.find(cls) != registry.m_builtins_types.end()) [[unlikely]] { - throw py::value_error("PyTree type " + PyRepr(cls) + - " is a built-in type and cannot be unregistered."); + const std::string ®istry_namespace, + RegistrationPtr ®istration) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { + return RegistryStatus::BuiltinType; } if (registry_namespace.empty()) [[unlikely]] { - const auto it = registry.m_registrations.find(cls); - if (it == registry.m_registrations.end()) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { - oss << "is a class of `PyStructSequence`, " - << "which is not explicitly registered in the global namespace."; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - oss << "is a subclass of `collections.namedtuple`, " - << "which is not explicitly registered in the global namespace."; - } else [[likely]] { - oss << "is not registered in the global namespace."; - } - throw py::value_error(oss.str()); + const auto it = m_registrations.find(cls); + if (it == m_registrations.end()) [[unlikely]] { + return RegistryStatus::NotRegistered; } - RegistrationPtr registration = it->second; - registry.m_registrations.erase(it); - return registration; + registration = it->second; + m_registrations.erase(it); } else [[likely]] { - const auto named_it = - registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); - if (named_it == registry.m_named_registrations.end()) [[unlikely]] { - std::ostringstream oss{}; - oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { - oss << "is a class of `PyStructSequence`, " - << "which is not explicitly registered "; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { - oss << "is a subclass of `collections.namedtuple`, " - << "which is not explicitly registered "; - } else [[likely]] { - oss << "is not registered "; - } - oss << "in namespace " << PyRepr(registry_namespace) << "."; - throw py::value_error(oss.str()); + const auto named_it = m_named_registrations.find(std::make_pair(registry_namespace, cls)); + if (named_it == m_named_registrations.end()) [[unlikely]] { + return RegistryStatus::NotRegistered; } - RegistrationPtr registration = named_it->second; - registry.m_named_registrations.erase(named_it); - return registration; + registration = named_it->second; + m_named_registrations.erase(named_it); } + return RegistryStatus::Ok; } /*static*/ void PyTreeTypeRegistry::Unregister(const py::object &cls, const std::string ®istry_namespace) { - const scoped_write_lock lock{sm_mutex}; - - const auto registration1 = UnregisterImpl(cls, registry_namespace); - const auto registration2 = UnregisterImpl(cls, registry_namespace); - EXPECT_TRUE(registration1->type.is(registration2->type)); - EXPECT_TRUE(registration1->flatten_func.is(registration2->flatten_func)); - EXPECT_TRUE(registration1->unflatten_func.is(registration2->unflatten_func)); - EXPECT_TRUE(registration1->path_entry_type.is(registration2->path_entry_type)); + // Classify the type BEFORE taking `sm_mutex`. On the not-found path `UnregisterImpl` builds its + // error message from `IsStructSequenceClass` / `IsNamedTupleClass`, which run Python and + // release the GIL; calling them while holding `sm_mutex` in write mode inverts the GIL <-> + // `sm_mutex` lock order and deadlocks a concurrent flatten that holds the GIL while waiting on + // `sm_mutex` in read mode (mirrors `Register`). + const bool is_structsequence_class = IsStructSequenceClass(cls); + const bool is_namedtuple_class = IsNamedTupleClass(cls); + + // Acquire both singletons BEFORE `sm_mutex`, mirroring `Init`/`Clear` (see `Lookup`/`Register` + // for the lock-order rationale). + auto ®istry1 = GetSingleton(); + auto ®istry2 = GetSingleton(); + + // These outlive the locked scope: dropping the last reference to a member runs arbitrary Python + // (`__del__`, weakref callbacks) that can re-enter optree and deadlock on `sm_mutex`. + RegistrationPtr registration1{nullptr}; + RegistrationPtr registration2{nullptr}; + RegistryStatus status = RegistryStatus::Ok; + { + const scoped_write_lock lock{sm_mutex}; + + status = registry1.UnregisterImpl(cls, registry_namespace, registration1); + if (status == RegistryStatus::Ok) [[likely]] { + status = registry2.UnregisterImpl(cls, registry_namespace, registration2); + } + if (status == RegistryStatus::Ok) [[likely]] { + EXPECT_TRUE(registration1->type.is(registration2->type)); + EXPECT_TRUE(registration1->flatten_func.is(registration2->flatten_func)); + EXPECT_TRUE(registration1->unflatten_func.is(registration2->unflatten_func)); + EXPECT_TRUE(registration1->path_entry_type.is(registration2->path_entry_type)); + } + } + + // Format the error only after the lock is released (mirrors `Register`). + if (status != RegistryStatus::Ok) [[unlikely]] { + if (status == RegistryStatus::BuiltinType) [[unlikely]] { + throw py::value_error("PyTree type " + PyRepr(cls) + + " is a built-in type and cannot be unregistered."); + } + std::ostringstream oss{}; + oss << "PyTree type " << PyRepr(cls) << " "; + if (is_structsequence_class) [[unlikely]] { + oss << "is a class of `PyStructSequence`, which is not explicitly registered "; + } else if (is_namedtuple_class) [[unlikely]] { + oss << "is a subclass of `collections.namedtuple`, which is not explicitly registered "; + } else [[likely]] { + oss << "is not registered "; + } + if (registry_namespace.empty()) [[unlikely]] { + oss << "in the global namespace."; + } else [[likely]] { + oss << "in namespace " << PyRepr(registry_namespace) << "."; + } + throw py::value_error(oss.str()); + } + + // Drop the registry's references with the lock released; the registrations die at scope exit. registration1->type.dec_ref(); registration1->flatten_func.dec_ref(); registration1->unflatten_func.dec_ref(); @@ -252,18 +299,24 @@ template /*static*/ PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::Lookup( const py::object &cls, const std::string ®istry_namespace) { - const scoped_read_lock lock{sm_mutex}; - + // Acquire the singleton BEFORE `sm_mutex`, mirroring `Init`/`Clear`. Under + // `per_interpreter_gil`, `GetSingleton()` releases the GIL on every call once a subinterpreter + // has existed; doing that while holding `sm_mutex` inverts the GIL <-> `sm_mutex` lock order + // against a concurrent registration (write lock) and deadlocks. const auto ®istry = GetSingleton(); - if (!registry_namespace.empty()) [[unlikely]] { - const auto named_it = - registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); - if (named_it != registry.m_named_registrations.end()) [[likely]] { - return named_it->second; + + { + const scoped_read_lock lock{sm_mutex}; + if (!registry_namespace.empty()) [[unlikely]] { + const auto named_it = + registry.m_named_registrations.find(std::make_pair(registry_namespace, cls)); + if (named_it != registry.m_named_registrations.end()) [[likely]] { + return named_it->second; + } } + const auto it = registry.m_registrations.find(cls); + return it != registry.m_registrations.end() ? it->second : nullptr; } - const auto it = registry.m_registrations.find(cls); - return it != registry.m_registrations.end() ? it->second : nullptr; } template PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::Lookup( @@ -327,8 +380,19 @@ template PyTreeKind PyTreeTypeRegistry::GetKind( EXPECT_EQ(registry1.m_named_registrations.size(), registry2.m_named_registrations.size()); } - auto atexit_register = py::getattr(py::module_::import("atexit"), "register"); - atexit_register(py::cpp_function(&Clear)); + // `atexit.register` runs Python, so it must not run under `sm_mutex`, and it can raise: without + // the rollback a failed import would leave an ID that no callback can ever remove (mirrors + // `WeakKeyCache::LookupOrInsert`). The rollback locks with the GIL held, as `Clear` does. + try { + auto atexit_register = py::getattr(py::module_::import("atexit"), "register"); + atexit_register(py::cpp_function(&Clear)); + } catch (...) { + const scoped_write_lock lock{sm_mutex}; + + sm_alive_interpids.erase(interpid); + --sm_num_interpreters_seen; + throw; + } } // NOLINTNEXTLINE[readability-function-cognitive-complexity] @@ -337,6 +401,11 @@ template PyTreeKind PyTreeTypeRegistry::GetKind( auto ®istry2 = GetSingleton(); const auto interpid = GetCurrentPyInterpreterID(); + // Detached under the lock and destroyed after it, for the reason given in `Unregister`. + RegistrationsMap registrations1{}; + NamedRegistrationsMap named_registrations1{}; + RegistrationsMap registrations2{}; + NamedRegistrationsMap named_registrations2{}; { const scoped_write_lock lock{sm_mutex}; @@ -404,25 +473,25 @@ template PyTreeKind PyTreeTypeRegistry::GetKind( } #endif - for (const auto &[_, registration1] : registry1.m_registrations) { - registration1->type.dec_ref(); - registration1->flatten_func.dec_ref(); - registration1->unflatten_func.dec_ref(); - registration1->path_entry_type.dec_ref(); - } - for (const auto &[_, registration1] : registry1.m_named_registrations) { - registration1->type.dec_ref(); - registration1->flatten_func.dec_ref(); - registration1->unflatten_func.dec_ref(); - registration1->path_entry_type.dec_ref(); - } - registry1.m_builtins_types.clear(); - registry1.m_registrations.clear(); - registry1.m_named_registrations.clear(); + registry1.m_registrations.swap(registrations1); + registry1.m_named_registrations.swap(named_registrations1); registry2.m_builtins_types.clear(); - registry2.m_registrations.clear(); - registry2.m_named_registrations.clear(); + registry2.m_registrations.swap(registrations2); + registry2.m_named_registrations.swap(named_registrations2); + } + + for (const auto &[_, registration1] : registrations1) { + registration1->type.dec_ref(); + registration1->flatten_func.dec_ref(); + registration1->unflatten_func.dec_ref(); + registration1->path_entry_type.dec_ref(); + } + for (const auto &[_, registration1] : named_registrations1) { + registration1->type.dec_ref(); + registration1->flatten_func.dec_ref(); + registration1->unflatten_func.dec_ref(); + registration1->path_entry_type.dec_ref(); } } diff --git a/src/treespec/constructors.cpp b/src/treespec/constructors.cpp index 380f6b84..db36f570 100644 --- a/src/treespec/constructors.cpp +++ b/src/treespec/constructors.cpp @@ -72,9 +72,17 @@ template Node node; node.kind = PyTreeTypeRegistry::GetKind(handle, node.custom, registry_namespace); - const auto verify_children = [&handle, &node, ®istry_namespace]( - const std::vector &children, - std::vector &treespecs) -> void { + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; + + const auto verify_children = + // NOLINTNEXTLINE[readability-function-cognitive-complexity] + [&handle, &node, ®istry_namespace, &is_dict_insertion_ordered_in_current_namespace]( + const std::vector &children, + std::vector &treespecs) -> void { for (const py::object &child : children) { if (!py::isinstance(child)) [[unlikely]] { std::ostringstream oss{}; @@ -105,6 +113,12 @@ template } } } + // A node depends on the namespace if it resolves a custom type, or if it is a dict whose + // keys are kept in insertion order (see the sort at the Dict case). + const bool depends_on_namespace = + node.kind == PyTreeKind::Custom || + ((node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) && + is_dict_insertion_ordered_in_current_namespace); if (!common_registry_namespace.empty()) [[likely]] { if (registry_namespace.empty()) [[likely]] { registry_namespace = common_registry_namespace; @@ -114,22 +128,28 @@ template << ", got " << PyRepr(common_registry_namespace) << "."; throw py::value_error(oss.str()); } - } else if (node.kind != PyTreeKind::Custom) [[likely]] { - registry_namespace = ""; + } else if (!depends_on_namespace) [[likely]] { + registry_namespace = ""; // mirrors `Flatten` } }; switch (node.kind) { case PyTreeKind::Leaf: { node.arity = 0; - PyErr_WarnEx(PyExc_UserWarning, - "PyTreeSpec::MakeFromCollection() is called on a leaf.", - /*stack_level=*/2); + // A childless node resolves no custom type, so it does not depend on the namespace. + // Keeping the caller's would make otherwise-identical treespecs compare unequal. + registry_namespace = ""; + if (PyErr_WarnEx(PyExc_UserWarning, + "PyTreeSpec::MakeFromCollection() is called on a leaf.", + /*stack_level=*/2) < 0) [[unlikely]] { + throw py::error_already_set(); + } break; } case PyTreeKind::None: { node.arity = 0; + registry_namespace = ""; if constexpr (!NoneIsLeaf) { break; } @@ -170,8 +190,7 @@ template keys = DictKeys(dict); if (node.kind != PyTreeKind::OrderedDict) [[likely]] { node.original_keys = DictFromKeys(dict); - if (!PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace)) - [[likely]] { + if (!is_dict_insertion_ordered) [[likely]] { TotalOrderSort(keys); } } @@ -272,6 +291,22 @@ template out->m_traversal.emplace_back(std::move(node)); out->m_none_is_leaf = NoneIsLeaf; out->m_namespace = registry_namespace; + // Reject a namespace promotion (an empty caller namespace adopting a child spec's namespace) + // that would rebind a custom node to a different registration than the one it holds, e.g. the + // root node, or a globally-resolved child, resolved in the global registry while the promoted + // namespace registers the type differently. The result keeps each node's original registration, + // mirroring the compose / transform / broadcast merge guards. Skipped for an empty namespace, + // which resolves every custom node globally. + if (!registry_namespace.empty()) [[unlikely]] { + if (const auto stale_type = out->FindStaleCustomType(registry_namespace)) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be composed into a collection: custom PyTree type " + << PyRepr(*stale_type) + << " no longer resolves to its original registration in namespace " + << PyRepr(registry_namespace) << "."; + throw py::value_error(oss.str()); + } + } out->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*out); return out; diff --git a/src/treespec/flatten.cpp b/src/treespec/flatten.cpp index 0894824d..0aff13dc 100644 --- a/src/treespec/flatten.cpp +++ b/src/treespec/flatten.cpp @@ -204,17 +204,11 @@ bool PyTreeSpec::FlattenInto(const py::handle &handle, const bool &none_is_leaf, const std::string ®istry_namespace) { bool found_custom = false; - bool is_dict_insertion_ordered = false; - bool is_dict_insertion_ordered_in_current_namespace = false; - { -#if defined(OPTREE_HAS_READ_WRITE_LOCK) - const scoped_read_lock lock{PyTreeTypeRegistry::sm_dict_order_mutex}; -#endif - is_dict_insertion_ordered = PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace); - is_dict_insertion_ordered_in_current_namespace = - PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace, - /*inherit_global_namespace=*/false); - } + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; if (none_is_leaf) [[unlikely]] { if (!is_dict_insertion_ordered) [[likely]] { @@ -481,17 +475,11 @@ bool PyTreeSpec::FlattenIntoWithPath(const py::handle &handle, const bool &none_is_leaf, const std::string ®istry_namespace) { bool found_custom = false; - bool is_dict_insertion_ordered = false; - bool is_dict_insertion_ordered_in_current_namespace = false; - { -#if defined(OPTREE_HAS_READ_WRITE_LOCK) - const scoped_read_lock lock{PyTreeTypeRegistry::sm_dict_order_mutex}; -#endif - is_dict_insertion_ordered = PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace); - is_dict_insertion_ordered_in_current_namespace = - PyTreeTypeRegistry::IsDictInsertionOrdered(registry_namespace, - /*inherit_global_namespace=*/false); - } + const auto dict_order_flags = + PyTreeTypeRegistry::GetDictInsertionOrderedFlags(registry_namespace); + const bool is_dict_insertion_ordered = dict_order_flags.with_inherited_global_namespace; + const bool is_dict_insertion_ordered_in_current_namespace = + dict_order_flags.in_current_namespace; auto stack = reserved_vector(4); if (none_is_leaf) [[unlikely]] { @@ -569,9 +557,11 @@ py::list PyTreeSpec::FlattenUpTo(const py::object &tree) const { auto it = m_traversal.crbegin(); const ssize_t num_leaves = GetNumLeaves(); - py::list leaves{num_leaves}; - ssize_t leaf = num_leaves - 1; - while (!agenda.empty()) { + // A pre-sized `py::list` is GC-tracked with NULL slots from creation, and the walk below runs + // arbitrary user code (custom `flatten_func`, dict key `__hash__`/`__eq__`, `__repr__`) that + // can reach it via `gc.get_objects()`. Collect into a reversed vector instead. + auto leaves = reserved_vector(num_leaves); + while (!agenda.empty()) [[likely]] { if (it == m_traversal.crend()) [[unlikely]] { std::ostringstream oss{}; oss << "Tree structures did not match; expected: " << ToString() @@ -585,9 +575,8 @@ py::list PyTreeSpec::FlattenUpTo(const py::object &tree) const { switch (node.kind) { case PyTreeKind::Leaf: { - EXPECT_GE(leaf, 0, "Leaf count mismatch."); - ListSetItem(leaves, leaf, object); - --leaf; + EXPECT_LT(py::ssize_t_cast(leaves.size()), num_leaves, "Leaf count mismatch."); + leaves.emplace_back(object); break; } @@ -796,13 +785,18 @@ py::list PyTreeSpec::FlattenUpTo(const py::object &tree) const { INTERNAL_ERROR(); } } - if (it != m_traversal.crend() || leaf != -1) [[unlikely]] { + if (it != m_traversal.crend() || py::ssize_t_cast(leaves.size()) != num_leaves) [[unlikely]] { std::ostringstream oss{}; oss << "Tree structures did not match; expected: " << ToString() << ", got: " << PyRepr(tree) << "."; throw py::value_error(oss.str()); } - return leaves; + py::list result{num_leaves}; + ssize_t index = num_leaves; + for (const py::object &leaf : leaves) { + ListSetItem(result, --index, leaf); + } + return result; } template diff --git a/src/treespec/gc.cpp b/src/treespec/gc.cpp index f6c3700d..4de736b1 100644 --- a/src/treespec/gc.cpp +++ b/src/treespec/gc.cpp @@ -30,6 +30,10 @@ using pybind11::detail::is_holder_constructed; namespace optree { +// No exception may escape a `tp_traverse` / `tp_clear` slot: unwinding across the `extern "C"` +// boundary calls `std::terminate`. In particular, no `PYTREESPEC_SANITY_CHECK` below: `PyTpClear` +// empties the traversal, so a cleared but still-alive treespec would abort on the next collection. + // NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ int PyTreeSpec::PyTpTraverse(PyObject *self_base, visitproc visit, void *arg) { Py_VISIT(Py_TYPE(self_base)); @@ -38,11 +42,49 @@ namespace optree { return 0; } auto &self = thread_safe_cast(py::handle{self_base}); - PYTREESPEC_SANITY_CHECK(self); - for (const auto &node : self.m_traversal) { + + // Report a registration's members once, and only when this treespec owns every reference to it. + // The registration holds one reference to each member however many nodes point at it, so + // reporting per node would decrement the same object once per node and underflow its shadow + // refcount. While the registry still holds the registration it also keeps the members alive, so + // skipping then leaks nothing. + // + // The holders are counted by rescanning the traversal rather than through a map: a treespec + // references very few distinct registrations, and every container that could hold them + // allocates, which `tp_traverse` cannot afford (see above). + // + // Known limitation: a treespec can only count its own nodes, so when several treespecs each + // hold part of the references none of them reports the members and a cycle through them + // survives. Fixing that needs the registration to be a garbage-collected object with its own + // `tp_traverse`, so each edge is reported by its owner and no counting is needed. + const ssize_t num_nodes = py::ssize_t_cast(self.m_traversal.size()); + for (ssize_t i = 0; i < num_nodes; ++i) { + const auto &node = self.m_traversal[i]; Py_VISIT(node.node_data.ptr()); Py_VISIT(node.node_entries.ptr()); Py_VISIT(node.original_keys.ptr()); + if (node.custom == nullptr) [[likely]] { + continue; + } + // Scanning from the start, the first match decides: before `i` an earlier node already + // reported this registration, at `i` this node is the first holder and counts the rest. + ssize_t num_holders = 0; + for (ssize_t j = 0; j < num_nodes; ++j) { + if (self.m_traversal[j].custom != node.custom) [[likely]] { + continue; + } + if (j < i) [[likely]] { + num_holders = 0; // not the first holder, skip reporting + break; + } + ++num_holders; + } + if (num_holders > 0 && node.custom.use_count() == num_holders) [[unlikely]] { + Py_VISIT(node.custom->type.ptr()); + Py_VISIT(node.custom->flatten_func.ptr()); + Py_VISIT(node.custom->unflatten_func.ptr()); + Py_VISIT(node.custom->path_entry_type.ptr()); + } } return 0; } @@ -53,16 +95,17 @@ namespace optree { return 0; } auto &self = thread_safe_cast(py::handle{self_base}); - PYTREESPEC_SANITY_CHECK(self); for (auto &node : self.m_traversal) { Py_CLEAR(node.node_data.ptr()); Py_CLEAR(node.node_entries.ptr()); Py_CLEAR(node.original_keys.ptr()); + node.custom.reset(); } self.m_traversal.clear(); return 0; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ int PyTreeIter::PyTpTraverse(PyObject *self_base, visitproc visit, void *arg) { Py_VISIT(Py_TYPE(self_base)); if (!::is_holder_constructed(self_base)) [[unlikely]] { @@ -74,6 +117,11 @@ namespace optree { Py_VISIT(obj.ptr()); } Py_VISIT(self.m_root.ptr()); + if (self.m_leaf_predicate) [[likely]] { + // The leaf predicate is an owned Python callback; it must be visited so the cyclic GC can + // see reference cycles that pass through it (otherwise such cycles leak). + Py_VISIT(self.m_leaf_predicate->ptr()); + } return 0; } @@ -88,6 +136,10 @@ namespace optree { } self.m_agenda.clear(); Py_CLEAR(self.m_root.ptr()); + // Reset the optional rather than clearing the held function: `NextImpl` tests the optional for + // presence, so leaving it engaged around a null callable would call through it after a + // collection. + self.m_leaf_predicate.reset(); return 0; } diff --git a/src/treespec/richcomparison.cpp b/src/treespec/richcomparison.cpp index 96b61798..42b60d2d 100644 --- a/src/treespec/richcomparison.cpp +++ b/src/treespec/richcomparison.cpp @@ -16,6 +16,7 @@ limitations under the License. */ #include // std::copy, std::reverse +#include // std::make_reverse_iterator #include // std::unordered_map #include // std::vector @@ -129,7 +130,13 @@ bool PyTreeSpec::IsPrefix(const PyTreeSpec &other, const bool &strict) const { EXPECT_EQ(reordered_other_offsets.front(), b->num_nodes, "PyTreeSpec traversal out of range."); - auto original_b = other.m_traversal.crbegin() + (b - other_traversal.crbegin()); + // Snapshot `b`'s subtree from the working copy before permuting its children + // below. The permutation copies overlapping child ranges, so the source must be + // a stable snapshot. It must be taken from the working copy (not the pristine + // `other.m_traversal`): a previously-processed ancestor dict may have relocated + // this subtree, so `b`'s offset no longer matches the pristine traversal. + const std::vector b_subtree(b.base() - b->num_nodes, b.base()); + const auto original_b = std::make_reverse_iterator(b_subtree.cend()); for (const auto &[i, j] : reordered_index_to_index) { std::copy(original_b + other_offsets[j + 1], original_b + other_offsets[j], diff --git a/src/treespec/serialization.cpp b/src/treespec/serialization.cpp index 04a36052..29123565 100644 --- a/src/treespec/serialization.cpp +++ b/src/treespec/serialization.cpp @@ -22,6 +22,7 @@ limitations under the License. #include // std::string #include // std::this_thread::get_id #include // std::unordered_set +#include // std::pair, std::move #include "optree/optree.h" @@ -137,9 +138,17 @@ std::string PyTreeSpec::ToStringImpl() const { case PyTreeKind::NamedTuple: { const py::object type = node.node_data; const auto fields = NamedTupleGetFields(type); - EXPECT_EQ(TupleGetSize(fields), - node.arity, - "Number of fields and entries does not match."); + // The field names are read from the (mutable) `_fields` attribute at repr time, so + // a caller may have changed them after the treespec was built. Report the mismatch + // as a `ValueError`, not an internal error, since the cause is external. + if (TupleGetSize(fields) != node.arity) [[unlikely]] { + std::ostringstream oss{}; + oss << "Number of fields (" << TupleGetSize(fields) << ") of namedtuple type " + << PyRepr(type) << " does not match the arity (" << node.arity + << ") of the treespec node. The `_fields` attribute may have been modified " + "after the treespec was created."; + throw py::value_error(oss.str()); + } const std::string kind = PyStr(EVALUATE_WITH_LOCK_HELD(py::getattr(type, "__name__"), type)); sstream << kind << "("; @@ -213,15 +222,24 @@ std::string PyTreeSpec::ToStringImpl() const { const py::object qualname = EVALUATE_WITH_LOCK_HELD(py::getattr(type, "__qualname__"), type); sstream << PyStr(qualname) << "("; - bool first = true; + ssize_t index = 0; auto child_it = agenda.cend() - node.arity; for (const py::handle &field : fields) { - if (!first) [[likely]] { + if (index > 0) [[likely]] { sstream << ", "; } - sstream << PyStr(field) << "=" << *child_it; + const std::string name = PyStr(field); + // An unnamed slot has no valid identifier, so render it angle-bracketed like + // CPython's `` rather than as the bare marker, which would read as an + // invalid keyword argument. The index disambiguates multiple unnamed slots. + if (name == PyStructSequenceUnnamedField()) [[unlikely]] { + sstream << ""; + } else [[likely]] { + sstream << name; + } + sstream << "=" << *child_it; ++child_it; - first = false; + ++index; } sstream << ")"; break; @@ -304,16 +322,32 @@ py::object PyTreeSpec::ToPicklable() const { const scoped_critical_section2 cs{ node.custom != nullptr ? py::handle{node.custom->type} : py::handle{}, node.node_data}; + + // Copy the node's mutable containers so the pickled state cannot alias (and, if the caller + // mutates it, corrupt) the immutable spec. Only dict-like keys are mutable and internal; a + // namedtuple/PyStructSequence type or custom metadata is left as-is. + py::object node_data = + node.node_data ? py::reinterpret_borrow(node.node_data) : py::none(); + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::OrderedDict) [[unlikely]] { + node_data = ListCopy(node.node_data); + } else if (node.kind == PyTreeKind::DefaultDict) [[unlikely]] { + const auto metadata = py::reinterpret_borrow(node.node_data); + node_data = + py::make_tuple(TupleGetItem(metadata, 0), ListCopy(TupleGetItem(metadata, 1))); + } + TupleSetItem(node_states, i++, py::make_tuple(py::int_(static_cast(node.kind)), py::int_(node.arity), - node.node_data ? node.node_data : py::none(), + std::move(node_data), node.node_entries ? node.node_entries : py::none(), node.custom != nullptr ? node.custom->type : py::none(), py::int_(node.num_leaves), py::int_(node.num_nodes), - node.original_keys ? node.original_keys : py::none())); + node.original_keys + ? static_cast(DictCopy(node.original_keys)) + : static_cast(py::none()))); } return py::make_tuple(node_states, py::bool_(m_none_is_leaf), py::str(m_namespace)); } @@ -321,9 +355,25 @@ py::object PyTreeSpec::ToPicklable() const { // NOLINTBEGIN[cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers] // NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ std::unique_ptr PyTreeSpec::FromPicklable(const py::object &picklable) { + const auto malformed = [](const std::string &reason) -> std::runtime_error { + return std::runtime_error("Malformed pickled PyTreeSpec: " + reason + "."); + }; + // `DistinctCount` hashes the keys, so an unhashable one raises `TypeError`. Report it as a + // malformed pickle like every other structural defect instead of letting it escape. + const auto distinct_count = [&malformed](const py::handle &keys) -> ssize_t { + try { + return DistinctCount(keys); + } catch (py::error_already_set &ex) { + if (!ex.matches(PyExc_TypeError)) [[unlikely]] { + throw; + } + throw malformed("the keys are not hashable"); + } + }; + const auto state = thread_safe_cast(picklable); if (state.size() != 3) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + throw malformed("the state is not a 3-tuple"); } bool none_is_leaf = false; std::string registry_namespace{}; @@ -332,56 +382,158 @@ py::object PyTreeSpec::ToPicklable() const { out->m_namespace = registry_namespace = thread_safe_cast(state[2]); const auto node_states = thread_safe_cast(state[0]); for (const auto &item : node_states) { - const auto t = thread_safe_cast(item); + const auto node_state = thread_safe_cast(item); + const auto node_state_size = node_state.size(); + if (node_state_size != 7 && node_state_size != 8) [[unlikely]] { + throw malformed("a node state is not a 7- or 8-tuple"); + } + Node &node = out->m_traversal.emplace_back(); - node.kind = static_cast(thread_safe_cast(t[0])); - node.arity = thread_safe_cast(t[1]); - if (t.size() != 7) [[unlikely]] { - if (t.size() == 8) [[likely]] { - if (t[7].is_none()) [[likely]] { - if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) - [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); - } - } else [[unlikely]] { - if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) - [[likely]] { - node.original_keys = DictFromKeys(t[7]); - } else [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); - } + const auto kind_value = thread_safe_cast(node_state[0]); + if (kind_value < 0 || kind_value >= static_cast(PyTreeKind::NumKinds)) + [[unlikely]] { + throw malformed("the node kind is out of range"); + } + node.kind = static_cast(kind_value); + node.arity = thread_safe_cast(node_state[1]); + if (node.arity < 0) [[unlikely]] { + throw malformed("the node arity is negative"); + } + if (node_state_size == 8) [[likely]] { + const auto &original_keys = node_state[7]; + if (original_keys.is_none()) [[likely]] { + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) + [[unlikely]] { + throw malformed("a dict node is missing its original keys"); } } else [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + if (node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) + [[likely]] { + // `DictFromKeys` builds a new dict, so the caller cannot mutate it afterwards. + node.original_keys = DictFromKeys(original_keys); + } else [[unlikely]] { + throw malformed("a non-dict node must not have original keys"); + } } } + + node.num_leaves = thread_safe_cast(node_state[5]); + node.num_nodes = thread_safe_cast(node_state[6]); + if (node.num_leaves < 0 || node.num_nodes < 1) [[unlikely]] { + throw malformed("a node has a negative or invalid size"); + } + + const auto &node_data = node_state[2]; + const auto &node_entries = node_state[3]; + const auto &custom_type = node_state[4]; switch (node.kind) { case PyTreeKind::Leaf: - case PyTreeKind::None: + case PyTreeKind::None: { + if (!node_data.is_none()) [[unlikely]] { + throw malformed("a leaf or none node must not have node data"); + } + // A leaf or none node is childless; a nonzero arity would let it absorb preceding + // subtrees (folding consistently) and silently drop leaves on unflatten. + if (node.arity != 0) [[unlikely]] { + throw malformed("a leaf or none node must have arity 0"); + } + // With `none_is_leaf`, None is flattened as a leaf, so a flattened tree never + // contains a None-kind node; a reconstructed one would later trip an InternalError + // in `FlattenUpTo` (`GetKind` never returns `None` under `NoneIsLeaf`). + if (node.kind == PyTreeKind::None && none_is_leaf) [[unlikely]] { + throw malformed("a none node cannot appear when none_is_leaf is set"); + } + break; + } + case PyTreeKind::Tuple: case PyTreeKind::List: { - if (!t[2].is_none()) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + if (!node_data.is_none()) [[unlikely]] { + throw malformed("a tuple or list node must not have node data"); } break; } case PyTreeKind::Dict: case PyTreeKind::OrderedDict: { - node.node_data = thread_safe_cast(t[2]); + // Copy the keys instead of borrowing them. + node.node_data = ListCopy(thread_safe_cast(node_data)); + if (ListGetSize(node.node_data) != node.arity) [[unlikely]] { + throw malformed("the number of keys does not match the arity"); + } + // The keys must be hashable and distinct; a duplicate or unhashable key would + // collapse or fail when the dict is rebuilt, desyncing the keys from the children. + if (distinct_count(node.node_data) != node.arity) [[unlikely]] { + throw malformed("the keys are not distinct"); + } + break; + } + + case PyTreeKind::NamedTuple: { + node.node_data = thread_safe_cast(node_data); + if (!IsNamedTupleClass(node.node_data)) [[unlikely]] { + throw malformed("the node data is not a namedtuple type"); + } + if (TupleGetSize(NamedTupleGetFields(node.node_data)) != node.arity) [[unlikely]] { + throw malformed("the number of fields does not match the arity"); + } break; } - case PyTreeKind::NamedTuple: case PyTreeKind::StructSequence: { - node.node_data = thread_safe_cast(t[2]); + node.node_data = thread_safe_cast(node_data); + if (!IsStructSequenceClass(node.node_data)) [[unlikely]] { + throw malformed("the node data is not a PyStructSequence type"); + } + if (TupleGetSize(StructSequenceGetFields(node.node_data)) != node.arity) + [[unlikely]] { + throw malformed("the number of fields does not match the arity"); + } + break; + } + + case PyTreeKind::DefaultDict: { + // A default dict stores its metadata as a 2-tuple `(default_factory, sorted_keys)`. + // `MakeNode` reads it with raw tuple/list accessors, so validate the shape here to + // avoid type-confusion on malformed input. + const auto metadata = thread_safe_cast(node_data); + if (metadata.size() != 2) [[unlikely]] { + throw malformed("the defaultdict metadata is not a 2-tuple"); + } + // `default_factory` is passed to `defaultdict(...)`, which requires None or + // callable. + const auto default_factory = TupleGetItem(metadata, 0); + if (!(default_factory.is_none() || + static_cast(PyCallable_Check(default_factory.ptr())))) [[unlikely]] { + throw malformed("the `default_factory` is not callable"); + } + // Copy the keys and rebuild the metadata tuple around the copy (see the dict case). + const auto keys = ListCopy(thread_safe_cast(metadata[1])); + if (ListGetSize(keys) != node.arity) [[unlikely]] { + throw malformed("the number of keys does not match the arity"); + } + if (distinct_count(keys) != node.arity) [[unlikely]] { + throw malformed("the keys are not distinct"); + } + node.node_data = py::make_tuple(default_factory, keys); + break; + } + + case PyTreeKind::Deque: { + // A deque's `maxlen` is None (unbounded) or a non-negative int bounding its length, + // so it must be at least the node's arity. + if (!node_data.is_none()) [[likely]] { + if (PyLong_Check(node_data.ptr()) == 0 || + thread_safe_cast(node_data) < node.arity) [[unlikely]] { + throw malformed("the deque maxlen is invalid"); + } + } + node.node_data = node_data; break; } - case PyTreeKind::DefaultDict: - case PyTreeKind::Deque: case PyTreeKind::Custom: { - node.node_data = t[2]; + node.node_data = node_data; break; } @@ -389,24 +541,24 @@ py::object PyTreeSpec::ToPicklable() const { default: INTERNAL_ERROR(); } - if (node.kind == PyTreeKind::Custom) [[unlikely]] { // NOLINT - if (!t[3].is_none()) [[unlikely]] { - node.node_entries = thread_safe_cast(t[3]); + if (node.kind == PyTreeKind::Custom) [[unlikely]] { + if (!node_entries.is_none()) [[unlikely]] { + node.node_entries = thread_safe_cast(node_entries); } - if (t[4].is_none()) [[unlikely]] { + if (custom_type.is_none()) [[unlikely]] { node.custom = nullptr; } else [[likely]] { if (none_is_leaf) [[unlikely]] { node.custom = - PyTreeTypeRegistry::Lookup(t[4], registry_namespace); + PyTreeTypeRegistry::Lookup(custom_type, registry_namespace); } else [[likely]] { node.custom = - PyTreeTypeRegistry::Lookup(t[4], registry_namespace); + PyTreeTypeRegistry::Lookup(custom_type, registry_namespace); } } if (node.custom == nullptr) [[unlikely]] { std::ostringstream oss{}; - oss << "Unknown custom type in pickled PyTreeSpec: " << PyRepr(t[4]); + oss << "Unknown custom type in pickled PyTreeSpec: " << PyRepr(custom_type); if (!registry_namespace.empty()) [[likely]] { oss << " in namespace " << PyRepr(registry_namespace); } else [[unlikely]] { @@ -415,21 +567,69 @@ py::object PyTreeSpec::ToPicklable() const { oss << "."; throw std::runtime_error(oss.str()); } - } else if (!t[3].is_none() || !t[4].is_none()) [[unlikely]] { - throw std::runtime_error("Malformed pickled PyTreeSpec."); + // A built-in registration (NoneType/tuple/list/dict/...) lives in the same map but has + // empty flatten/unflatten callables, which `MakeNode`'s custom branch would call. + if (node.custom->kind != PyTreeKind::Custom) [[unlikely]] { + throw malformed("the custom type is a built-in type"); + } + } else if (!node_entries.is_none() || !custom_type.is_none()) [[unlikely]] { + throw malformed("a non-custom node must not have node entries or a custom type"); } - if (node.original_keys && DictGetSize(node.original_keys) != node.arity) [[unlikely]] { - throw std::runtime_error("Number of keys does not match arity in pickled PyTreeSpec."); + if (node.original_keys) [[unlikely]] { + if (DictGetSize(node.original_keys) != node.arity) [[unlikely]] { + throw malformed("the number of original keys does not match the arity"); + } + // `original_keys` records the insertion order of the same keys stored (sorted) in + // node_data; its key set must match, or unflatten would map children onto keys the dict + // never had. + const auto keys = (node.kind == PyTreeKind::DefaultDict + ? TupleGetItemAs(node.node_data, 1) + : py::reinterpret_borrow(node.node_data)); + if (!DictKeysEqual(keys, py::reinterpret_borrow(node.original_keys))) + [[unlikely]] { + throw malformed("the keys do not match the original keys"); + } } if (node.node_entries && !node.node_entries.is_none() && TupleGetSize(node.node_entries) != node.arity) [[unlikely]] { - throw std::runtime_error( - "Number of node entries does not match arity in pickled PyTreeSpec."); + throw malformed("the number of node entries does not match the arity"); } + } - node.num_leaves = thread_safe_cast(t[5]); - node.num_nodes = thread_safe_cast(t[6]); + // Validate that the reconstructed traversal is structurally consistent. + // `PYTREESPEC_SANITY_CHECK` only checks the final node, so a malformed pickle could otherwise + // smuggle in inconsistent arity / num_nodes / num_leaves that cause out-of-bounds access when + // the spec is later used. Walk the post-order traversal, folding each node's children off a + // stack of subtree sizes. + { + auto subtree_sizes = + reserved_vector>( + out->m_traversal.size()); + for (const Node &node : out->m_traversal) { + if (static_cast(subtree_sizes.size()) < node.arity) [[unlikely]] { + throw malformed("a node has more children than available subtrees"); + } + ssize_t children_num_nodes = 0; + ssize_t children_num_leaves = 0; + for (ssize_t i = 0; i < node.arity; ++i) { + children_num_nodes += subtree_sizes.back().first; + children_num_leaves += subtree_sizes.back().second; + subtree_sizes.pop_back(); + } + const ssize_t expected_num_nodes = children_num_nodes + 1; + const ssize_t expected_num_leaves = + (node.kind == PyTreeKind::Leaf ? ssize_t{1} : children_num_leaves); + if (node.num_nodes != expected_num_nodes || node.num_leaves != expected_num_leaves) + [[unlikely]] { + throw malformed("a node's size is inconsistent with its children"); + } + subtree_sizes.emplace_back(node.num_nodes, node.num_leaves); + } + if (subtree_sizes.size() != 1) [[unlikely]] { + throw malformed("the traversal does not yield a single tree"); + } } + out->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*out); return out; diff --git a/src/treespec/traversal.cpp b/src/treespec/traversal.cpp index 5815d8be..8422b780 100644 --- a/src/treespec/traversal.cpp +++ b/src/treespec/traversal.cpp @@ -15,9 +15,12 @@ limitations under the License. ================================================================================ */ +#include // std::memory_order_acquire, std::memory_order_release +#include // std::rethrow_exception, std::current_exception #include // std::optional #include // std::ostringstream #include // std::runtime_error +#include // std::this_thread::get_id, std::thread::id #include "optree/optree.h" @@ -161,6 +164,14 @@ py::object PyTreeIter::NextImpl() { } py::object PyTreeIter::Next() { + // `NextImpl` runs user code (the `is_leaf` predicate, custom `flatten_func`) that may call + // `next()` on this same iterator. `m_mutex` is not recursive and the GIL is released while + // waiting on it, so that would hang; reject it as CPython's "generator already executing" does. + const auto ident = std::this_thread::get_id(); + if (m_running_thread_id.load(std::memory_order_acquire) == ident) [[unlikely]] { + throw std::runtime_error("PyTreeIter is already iterating."); + } + #if !defined(Py_GIL_DISABLED) const py::gil_scoped_release_simple gil_release{}; #endif @@ -169,10 +180,19 @@ py::object PyTreeIter::Next() { #if !defined(Py_GIL_DISABLED) const py::gil_scoped_acquire_simple gil_acquire{}; #endif - if (m_none_is_leaf) [[unlikely]] { - return NextImpl(); - } else [[likely]] { - return NextImpl(); + m_running_thread_id.store(ident, std::memory_order_release); + try { + py::object leaf{}; + if (m_none_is_leaf) [[unlikely]] { + leaf = NextImpl(); + } else [[likely]] { + leaf = NextImpl(); + } + m_running_thread_id.store(std::thread::id{}, std::memory_order_release); + return leaf; + } catch (...) { + m_running_thread_id.store(std::thread::id{}, std::memory_order_release); + std::rethrow_exception(std::current_exception()); } } } diff --git a/src/treespec/treespec.cpp b/src/treespec/treespec.cpp index 44305dd0..d884dfda 100644 --- a/src/treespec/treespec.cpp +++ b/src/treespec/treespec.cpp @@ -185,13 +185,36 @@ namespace optree { } } +std::optional PyTreeSpec::FindStaleCustomType( + const std::string &target_namespace) const { + for (const Node &node : m_traversal) { + if (node.kind == PyTreeKind::Custom) [[unlikely]] { + const auto registration = + (m_none_is_leaf + ? PyTreeTypeRegistry::Lookup(node.custom->type, target_namespace) + : PyTreeTypeRegistry::Lookup(node.custom->type, + target_namespace)); + if (registration != node.custom) [[unlikely]] { + return node.custom->type; + } + } + } + return {}; +} + // NOLINTNEXTLINE[readability-function-cognitive-complexity] /*static*/ std::tuple PyTreeSpec::BroadcastToCommonSuffixImpl( std::vector &nodes, const std::vector &traversal, const ssize_t &pos, const std::vector &other_traversal, - const ssize_t &other_pos) { + const ssize_t &other_pos, + const ssize_t &depth) { + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during broadcasting the treespecs."); + throw py::error_already_set(); + } const Node &root = traversal.at(pos); const Node &other_root = other_traversal.at(other_pos); EXPECT_GE(pos + 1, @@ -276,15 +299,19 @@ namespace optree { const auto expected_keys = (root.kind != PyTreeKind::DefaultDict ? py::reinterpret_borrow(root.node_data) : TupleGetItemAs(root.node_data, 1)); - auto other_keys = (other_root.kind != PyTreeKind::DefaultDict - ? py::reinterpret_borrow(other_root.node_data) - : TupleGetItemAs(other_root.node_data, 1)); + const auto other_keys = (other_root.kind != PyTreeKind::DefaultDict + ? py::reinterpret_borrow(other_root.node_data) + : TupleGetItemAs(other_root.node_data, 1)); const py::dict dict{}; for (ssize_t i = 0; i < other_root.arity; ++i) { DictSetItem(dict, ListGetItem(other_keys, i), py::int_(i)); } if (!DictKeysEqual(expected_keys, dict)) [[unlikely]] { - TotalOrderSort(other_keys); + // Build the message from a sorted COPY of the keys. `other_keys` is a borrow of the + // argument spec's live `node_data`; sorting it in place would permute the keys + // while the child subtrees stay put, silently corrupting a spec the caller still + // holds. + const py::list sorted_other_keys = SortedDictKeys(dict); const auto [missing_keys, extra_keys] = DictKeysDifference(expected_keys, dict); std::ostringstream key_difference_sstream{}; if (ListGetSize(missing_keys) != 0) [[likely]] { @@ -295,7 +322,7 @@ namespace optree { } std::ostringstream oss{}; oss << "dictionary key mismatch; expected key(s): " << PyRepr(expected_keys) - << ", got key(s): " << PyRepr(other_keys) << key_difference_sstream.str() + << ", got key(s): " << PyRepr(sorted_other_keys) << key_difference_sstream.str() << "."; throw py::value_error(oss.str()); } @@ -314,7 +341,12 @@ namespace optree { other_cur = other_curs[py::cast(DictGetItem(dict, key))]; const auto [num_nodes, other_num_nodes, new_num_nodes, new_num_leaves] = // NOLINTNEXTLINE[misc-no-recursion] - BroadcastToCommonSuffixImpl(nodes, traversal, cur, other_traversal, other_cur); + BroadcastToCommonSuffixImpl(nodes, + traversal, + cur, + other_traversal, + other_cur, + depth + 1); cur -= num_nodes; nodes[start_num_nodes].num_nodes += new_num_nodes; nodes[start_num_nodes].num_leaves += new_num_leaves; @@ -393,7 +425,12 @@ namespace optree { for (ssize_t i = root.arity - 1; i >= 0; --i) { const auto [num_nodes, other_num_nodes, new_num_nodes, new_num_leaves] = // NOLINTNEXTLINE[misc-no-recursion] - BroadcastToCommonSuffixImpl(nodes, traversal, cur, other_traversal, other_cur); + BroadcastToCommonSuffixImpl(nodes, + traversal, + cur, + other_traversal, + other_cur, + depth + 1); cur -= num_nodes; other_cur -= other_num_nodes; nodes[start_num_nodes].num_nodes += new_num_nodes; @@ -405,6 +442,7 @@ namespace optree { nodes[start_num_nodes].num_leaves}; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec &other) const { PYTREESPEC_SANITY_CHECK(*this); PYTREESPEC_SANITY_CHECK(other); @@ -420,12 +458,36 @@ std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec throw py::value_error(oss.str()); } + const std::string &target_namespace = m_namespace.empty() ? other.m_namespace : m_namespace; auto treespec = std::make_unique(); treespec->m_none_is_leaf = m_none_is_leaf; - if (other.m_namespace.empty()) [[likely]] { - treespec->m_namespace = m_namespace; + treespec->m_namespace = target_namespace; + + if (!target_namespace.empty()) [[likely]] { + // The compatibility check above rejects two distinct non-empty namespaces, so the adopted + // namespace equals one side's namespace and at most the other (empty) side differs from it. + // Re-check only that side: adopting a namespace under which its custom nodes resolve to a + // different registration would silently rebind them (the result keeps the original ones). + std::optional stale_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, other.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + stale_type = FindStaleCustomType(target_namespace); + } else if (target_namespace != other.m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, m_namespace, "Namespace mismatch."); + EXPECT_TRUE(other.m_namespace.empty(), "Namespace mismatch."); + stale_type = other.FindStaleCustomType(target_namespace); + } + if (stale_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*stale_type) + << " no longer resolves to its original registration in namespace " + << PyRepr(target_namespace) << "."; + throw py::value_error(oss.str()); + } } else [[unlikely]] { - treespec->m_namespace = other.m_namespace; + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + EXPECT_TRUE(other.m_namespace.empty(), "Namespace mismatch."); } const ssize_t num_nodes = GetNumNodes(); @@ -436,7 +498,8 @@ std::unique_ptr PyTreeSpec::BroadcastToCommonSuffix(const PyTreeSpec m_traversal, num_nodes - 1, other.m_traversal, - other_num_nodes - 1); + other_num_nodes - 1, + 0); std::reverse(treespec->m_traversal.begin(), treespec->m_traversal.end()); EXPECT_EQ(num_nodes_walked, num_nodes, @@ -573,11 +636,28 @@ std::unique_ptr PyTreeSpec::Transform(const std::optionalm_none_is_leaf = m_none_is_leaf; treespec->m_namespace = common_registry_namespace; + + // Reject a transform whose unified namespace would rebind a custom node to a different + // registration than the one it holds (the result keeps each node's original registration; e.g. + // a globally-resolved custom node from the input under a non-empty unified namespace). Only + // relevant for a non-empty namespace: an empty one resolves every custom node globally. + if (!common_registry_namespace.empty()) [[unlikely]] { + if (const auto &stale_type = treespec->FindStaleCustomType(common_registry_namespace)) + [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be transformed: custom PyTree type " << PyRepr(*stale_type) + << " no longer resolves to its original registration in namespace " + << PyRepr(common_registry_namespace) << "."; + throw py::value_error(oss.str()); + } + } + treespec->m_traversal.shrink_to_fit(); PYTREESPEC_SANITY_CHECK(*treespec); return treespec; } +// NOLINTNEXTLINE[readability-function-cognitive-complexity] std::unique_ptr PyTreeSpec::Compose(const PyTreeSpec &inner) const { PYTREESPEC_SANITY_CHECK(*this); PYTREESPEC_SANITY_CHECK(inner); @@ -593,12 +673,36 @@ std::unique_ptr PyTreeSpec::Compose(const PyTreeSpec &inner) const { throw py::value_error(oss.str()); } + const std::string &target_namespace = m_namespace.empty() ? inner.m_namespace : m_namespace; auto treespec = std::make_unique(); treespec->m_none_is_leaf = m_none_is_leaf; - if (inner.m_namespace.empty()) [[likely]] { - treespec->m_namespace = m_namespace; + treespec->m_namespace = target_namespace; + + if (!target_namespace.empty()) [[likely]] { + // The compatibility check above rejects two distinct non-empty namespaces, so the adopted + // namespace equals one side's namespace and at most the other (empty) side differs from it. + // Re-check only that side: adopting a namespace under which its custom nodes resolve to a + // different registration would silently rebind them (the result keeps the original ones). + std::optional stale_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, inner.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + stale_type = FindStaleCustomType(target_namespace); + } else if (target_namespace != inner.m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, m_namespace, "Namespace mismatch."); + EXPECT_TRUE(inner.m_namespace.empty(), "Namespace mismatch."); + stale_type = inner.FindStaleCustomType(target_namespace); + } + if (stale_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*stale_type) + << " no longer resolves to its original registration in namespace " + << PyRepr(target_namespace) << "."; + throw py::value_error(oss.str()); + } } else [[unlikely]] { - treespec->m_namespace = inner.m_namespace; + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + EXPECT_TRUE(inner.m_namespace.empty(), "Namespace mismatch."); } const ssize_t num_outer_leaves = GetNumLeaves(); @@ -638,6 +742,11 @@ ssize_t PyTreeSpec::PathsImpl(PathVector &paths, // NOLINT[misc-no-recursion] const ssize_t &depth) const { const Node &root = m_traversal.at(pos); EXPECT_GE(pos + 1, root.num_nodes, "PyTreeSpec::Paths() walked off start of array."); + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during walking the tree."); + throw py::error_already_set(); + } ssize_t cur = pos - 1; // NOLINTNEXTLINE[misc-no-recursion] @@ -736,6 +845,11 @@ ssize_t PyTreeSpec::AccessorsImpl(Span &accessors, // NOLINT[misc-no-recursion] const Node &root = m_traversal.at(pos); EXPECT_GE(pos + 1, root.num_nodes, "PyTreeSpec::TypedPaths() walked off start of array."); + if (depth > MAX_RECURSION_DEPTH) [[unlikely]] { + PyErr_SetString(PyExc_RecursionError, + "Maximum recursion depth exceeded during walking the tree."); + throw py::error_already_set(); + } ssize_t cur = pos - 1; const py::object node_type = GetType(root); @@ -860,11 +974,11 @@ py::list PyTreeSpec::Entries() const { case PyTreeKind::Dict: case PyTreeKind::OrderedDict: { const scoped_critical_section cs{root.node_data}; - return py::getattr(root.node_data, "copy")(); + return ListCopy(root.node_data); } case PyTreeKind::DefaultDict: { const scoped_critical_section cs{root.node_data}; - return py::getattr(TupleGetItem(root.node_data, 1), "copy")(); + return ListCopy(TupleGetItemAs(root.node_data, 1)); } case PyTreeKind::NumKinds: diff --git a/tests/concurrent/test_subinterpreters.py b/tests/concurrent/test_subinterpreters.py index aec86e9a..6a109168 100644 --- a/tests/concurrent/test_subinterpreters.py +++ b/tests/concurrent/test_subinterpreters.py @@ -343,3 +343,175 @@ def check_import(): output='', rerun=NUM_FLAKY_RERUNS, ) + + +def test_type_cache_cleanup_across_subinterpreters(): + # Exercise the process-global type caches across subinterpreter churn. The caches key on the + # type's memory ADDRESS, which the allocator may recycle after a type is freed. Each interpreter + # flattens a RANDOMLY chosen type and reads its repr, firing both the classification and the + # field-name cache; distinct field names make a stale entry on a recycled address observable. + # The subprocess makes finalization real. + check_script_in_subprocess( + f""" + import contextlib + import textwrap + from concurrent import interpreters + + resolve = textwrap.dedent( + ''' + import keyword + import os + import random + import string + import time + from collections import namedtuple + + import optree + + # Distinct PyStructSequence types, each with a distinct first sequence field. + cases = [ + (os.stat_result, 'st_mode'), + (time.struct_time, 'tm_year'), + (os.times_result, 'user'), + ] + case = random.choice([*cases, None]) + if case is None: + # Exercise the classification cache with a namedtuple type. + field_names = [] + while not field_names: + field_names = [ + ''.join(random.choices(string.ascii_lowercase, k=random.randint(1, 8))) + for _ in range(random.randint(2, 5)) + ] + field_names = [n for n in field_names if not keyword.iskeyword(n)] # avoid reserved names + field_names = list(dict.fromkeys(field_names)) # deduplicate + NamedTupleType = namedtuple('NamedTupleType', field_names) + leaves, treespec = optree.tree_flatten(NamedTupleType(*range(len(field_names)))) + assert len(leaves) == len(field_names), (leaves, treespec) + assert (field_names[0] + '=*') in repr(treespec), repr(treespec) + else: + # Exercise the classification cache and the field-name cache with a PyStructSequence type. + PyStructSequenceType, first_field = case + obj = PyStructSequenceType(range(PyStructSequenceType.n_sequence_fields)) + assert (first_field + '=*') in repr(optree.tree_structure(obj)), first_field + ''' + ).strip() + + exec(resolve) # the main interpreter resolves a random type + for _ in range({NUM_FUTURES}): + with contextlib.closing(interpreters.create()) as subinterpreter: + subinterpreter.exec(resolve) + exec(resolve) # the main interpreter must still resolve correctly after the churn + """, + output='', + rerun=NUM_FLAKY_RERUNS, + ) + + +def test_registry_init_failure_does_not_leak_the_interpreter_id(): + # Regression: `PyTreeTypeRegistry::Init` inserted the interpreter into the alive set and bumped + # the seen counter BEFORE `atexit.register(&Clear)` succeeded, with no rollback. A failed import + # then left behind an ID that no callback can ever remove, so the registry keeps believing a + # dead interpreter is alive and its shutdown invariants never hold again. `Init` must mark the + # interpreter only once the cleanup is registered, the way `WeakKeyCache::LookupOrInsert` does. + check_script_in_subprocess( + """ + import contextlib + import faulthandler + import textwrap + from concurrent import interpreters + + import optree + + faulthandler.dump_traceback_later(60, exit=True) # watchdog: abort the process on a hang + + before_ids = optree._C.get_alive_interpreter_ids() + before_seen = optree._C.get_num_interpreters_seen() + assert before_ids == {optree._C.get_current_interpreter_id()}, before_ids + + with contextlib.closing(interpreters.create()) as subinterpreter: + subinterpreter.exec( + textwrap.dedent( + ''' + import atexit + + def failing_register(*args, **kwargs): + raise RuntimeError('injected atexit failure') + + atexit.register = failing_register + try: + import optree + except ImportError: + pass + else: + raise AssertionError('the injected failure did not propagate') + ''', + ).strip(), + ) + + after_ids = optree._C.get_alive_interpreter_ids() + after_seen = optree._C.get_num_interpreters_seen() + assert after_ids == before_ids, (before_ids, after_ids) + assert after_seen == before_seen, (before_seen, after_seen) + """, + output='', + ) + + +def test_registry_lookup_no_gil_lock_order_deadlock(): + # Regression: `Lookup`/`Register`/`Unregister` called `GetSingleton()` while holding `sm_mutex`. + # Under `per_interpreter_gil` that call releases the GIL once any subinterpreter has imported + # `optree._C`, so a flatten thread drops the GIL holding the read lock while a registration + # thread holds the GIL waiting on the write lock: a lock-order inversion that hangs the process. + # The watchdog turns the hang into a non-zero exit; the fix acquires the singleton before + # locking. + check_script_in_subprocess( + f""" + import contextlib + import faulthandler + import threading + import time + from concurrent import interpreters + + import optree + + # Latch pybind11's `has_seen_non_main_interpreter` so `GetSingleton()` releases the GIL on + # every subsequent call (the single-interpreter fast path is disabled process-wide). + with contextlib.closing(interpreters.create()) as subinterpreter: + subinterpreter.exec('import optree._C') + + faulthandler.dump_traceback_later(20, exit=True) # a deadlock becomes a non-zero exit + + stop = threading.Event() + tree = {{'a': 1, 'b': [2, 3], 'c': (4, 5)}} + + def flatten_worker(): + while not stop.is_set(): + optree.tree_flatten(tree) + + def register_worker(index): + cls = type(f'DeadlockNode_{{index}}', (), {{}}) + while not stop.is_set(): + optree.register_pytree_node( + cls, + lambda x: ((), None), + lambda metadata, children: cls(), + namespace='deadlock' + ) + optree.unregister_pytree_node(cls, namespace='deadlock') + + workers = [threading.Thread(target=flatten_worker) for _ in range({NUM_WORKERS})] + workers += [threading.Thread(target=register_worker, args=(index,)) for index in range(2)] + for worker in workers: + worker.start() + time.sleep(0.5) # let readers and writers contend on `sm_mutex` + stop.set() + for worker in workers: + worker.join(timeout=5) + assert not worker.is_alive(), 'worker thread did not terminate (deadlock)' + + faulthandler.cancel_dump_traceback_later() + """, + output=None, + rerun=NUM_FLAKY_RERUNS, + ) diff --git a/tests/helpers.py b/tests/helpers.py index 5440bb11..dde27fa0 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -75,6 +75,14 @@ reason='Py_GIL_DISABLED is set', ) +# Free-threaded builds before 3.14 hold deferred references to type objects in per-thread caches, so +# `gc_collect()` cannot force a heap type to be reclaimed there. +HAS_DEFERRED_TYPE_REFS = Py_GIL_DISABLED and sys.version_info < (3, 14) +skipif_deferred_type_refs = pytest.mark.skipif( + HAS_DEFERRED_TYPE_REFS, + reason='free-threaded builds before 3.14 keep deferred references to type objects', +) + PYPY = platform.python_implementation() == 'PyPy' skipif_pypy = pytest.mark.skipif( PYPY, diff --git a/tests/integrations/test_attrs.py b/tests/integrations/test_attrs.py index d8b7ac5c..bfb5024a 100644 --- a/tests/integrations/test_attrs.py +++ b/tests/integrations/test_attrs.py @@ -97,6 +97,33 @@ def __attrs_post_init__(self): assert foo == optree.tree_unflatten(treespec, leaves) +def test_define_reconstruction_reapplies_field_converters(): + # Characterization: the generated unflatten rebuilds via `cls(**fields)`, which re-runs attrs + # field converters. Since `flatten` reads the already-converted value, an idempotent converter + # round-trips exactly while a non-idempotent one does not. A class needing exact reconstruction + # should register explicitly with custom flatten/unflatten (see `register_node`). + + # Idempotent converter: `int(int(x)) == int(x)`, so the round-trip is exact. + @optree.integrations.attrs.define(namespace='test-attrs-converter-idempotent') + class Cast: + x: int = optree.integrations.attrs.field(converter=int) + + cast = Cast(5.0) # __init__ converter: x = int(5.0) = 5 + leaves, treespec = optree.tree_flatten(cast, namespace='test-attrs-converter-idempotent') + assert leaves == [5] + assert optree.tree_unflatten(treespec, leaves).x == 5 + + # Non-idempotent converter: unflatten re-applies it, so the stored `6` becomes `7`. + @optree.integrations.attrs.define(namespace='test-attrs-converter-shift') + class Shift: + x: int = optree.integrations.attrs.field(converter=lambda v: v + 1) + + shift = Shift(5) # __init__ converter: x = 5 + 1 = 6 + leaves, treespec = optree.tree_flatten(shift, namespace='test-attrs-converter-shift') + assert leaves == [6] + assert optree.tree_unflatten(treespec, leaves).x == 7 # re-applied: 6 + 1 = 7, not 6 + + def test_define_frozen(): @optree.integrations.attrs.frozen(namespace='test-attrs-frozen') class FrozenPoint: @@ -143,7 +170,12 @@ def foo(): def test_define_with_duplicate_registrations(): with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.integrations\.attrs\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): @optree.integrations.attrs.define(namespace='error') @@ -292,10 +324,83 @@ class Double: with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.integrations\.attrs\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): optree.integrations.attrs.register_node(Double, namespace='test-attrs-double-2') + # As the error suggests, the class can still be registered in another namespace via the generic + # API with explicit flatten/unflatten functions. + optree.register_pytree_node( + Double, + lambda d: ((d.x,), None, None), + lambda _, children: Double(*children), + namespace='test-attrs-double-2', + ) + optree.unregister_pytree_node(Double, namespace='test-attrs-double-2') + + +def test_register_node_failure_does_not_leak_fields_guard(): + @attrs.define + class Leak: + x: int + + namespace = 'test-attrs-register-leak' + # Occupy the (class, namespace) slot directly so the attrs `register_node()`'s internal + # `register_pytree_node()` call fails, after `register_node()` would set its `_FIELDS` guard. + optree.register_pytree_node( + Leak, + lambda leak: ((leak.x,), None, None), + lambda _, children: Leak(*children), + namespace=namespace, + ) + try: + with pytest.raises(ValueError, match='already registered'): + optree.integrations.attrs.register_node(Leak, namespace=namespace) + finally: + optree.unregister_pytree_node(Leak, namespace=namespace) + + # A failed registration must not leave the `_FIELDS` guard behind, or the class becomes + # impossible to register ever again: every retry would raise "... more than once". + # Re-registration after clearing the conflicting entry must succeed. + optree.integrations.attrs.register_node(Leak, namespace=namespace) + optree.unregister_pytree_node(Leak, namespace=namespace) + + +def test_register_node_rolls_back_when_the_guard_cannot_be_set(): + # `register_node()` commits the registry entry first and sets the `__optree_attrs_fields__` + # guard afterwards. If the class refuses the guard, both must be rolled back: a class left + # registered but unguarded is just as unrecoverable, because the retry then fails with + # "already registered" instead. + class RejectGuardOnce(type): + rejected = False + + def __setattr__(cls, name, value, /): + if name == '__optree_attrs_fields__' and not RejectGuardOnce.rejected: + RejectGuardOnce.rejected = True # a transient failure: reject only the first time + raise RuntimeError('cannot set the guard attribute') + super().__setattr__(name, value) + + @attrs.define + class Rejecting(metaclass=RejectGuardOnce): + x: int + + namespace = 'test-attrs-register-rollback' + with pytest.raises(RuntimeError, match=r'cannot set the guard attribute'): + optree.integrations.attrs.register_node(Rejecting, namespace=namespace) + + assert '__optree_attrs_fields__' not in Rejecting.__dict__ + with pytest.raises(ValueError, match=r'is not registered'): + optree.unregister_pytree_node(Rejecting, namespace=namespace) + + # The class is left exactly as it was, so the retry succeeds. + optree.integrations.attrs.register_node(Rejecting, namespace=namespace) + optree.unregister_pytree_node(Rejecting, namespace=namespace) + def test_register_init_false_class_warns(): @attrs.define(init=False) @@ -356,9 +461,64 @@ class EntryTest: assert 'AttrsEntry' in repr(entry_str) assert "'x'" in repr(entry_str) - entry_int = optree.integrations.attrs.AttrsEntry(1, EntryTest, optree.PyTreeKind.CUSTOM) - assert entry_int.field == 'y' - assert entry_int.name == 'y' + +def test_attrs_entry_integer_indexes_children(): + # For a class registered with `optree.integrations.attrs.register_node()`, an integer entry + # indexes the children that registration emits (the fields that are BOTH `pytree_node=True` and + # `init`). A non-child field interleaved between children must not shift the mapping. + @optree.integrations.attrs.register_node(namespace='test-attrs-entry-integer') + @attrs.define + class Foo: + a: int + b: int = optree.integrations.attrs.field(default=0, pytree_node=False) # not a child + # non-init -> not a tree child + d: int = optree.integrations.attrs.field(init=False, default=0, pytree_node=False) + c: int = 0 + + foo = Foo(1, 2, 3) # a=1, b=2, c=3 (d defaults to 0, not an init parameter) + assert tuple(a.name for a in attrs.fields(Foo)) == ('a', 'b', 'd', 'c') + + entry_int = optree.integrations.attrs.AttrsEntry(1, Foo, optree.PyTreeKind.CUSTOM) + assert entry_int.init_fields == ('a', 'b', 'c') # `d` is not an init field + assert entry_int.children_fields == ('a', 'c') # `b` is metadata and `d` is non-init + # The 2nd child is `c` (not the metadata field `b` nor the non-init field `d`). + assert entry_int.field == 'c' + assert entry_int.name == 'c' + assert entry_int.codify('x') == 'x.c' + assert entry_int(foo) == foo.c == 3 + + +def test_attrs_entry_integer_follows_generic_flatten_func(): + # Regression: an integer entry was resolved with the `optree.integrations.attrs` field + # predicate, which does not apply to a class registered with the generic + # `register_pytree_node()`. There the flatten function alone decides the children, so + # `pytree_node=False` must not drop a field and leave the entry indexing past the end of an + # internal list. + @attrs.define + class Foo: + a: int + b: int = optree.integrations.attrs.field(default=0, pytree_node=False) + + optree.register_pytree_node( + Foo, + lambda foo: ((foo.a, foo.b), None, None), # both fields are children, entries default + lambda _, children: Foo(*children), + path_entry_type=optree.integrations.attrs.AttrsEntry, + namespace='test-attrs-entry-generic', + ) + + foo = Foo(1, 2) + paths, leaves, treespec = optree.tree_flatten_with_path( + foo, + namespace='test-attrs-entry-generic', + ) + assert paths == [(0,), (1,)] + assert leaves == [1, 2] + + accessors = treespec.accessors() + assert [accessor(foo) for accessor in accessors] == [1, 2] + assert [accessor[0].field for accessor in accessors] == ['a', 'b'] + assert [accessor.codify('x') for accessor in accessors] == ['x.a', 'x.b'] def test_accessor_codify(): diff --git a/tests/test_accessors.py b/tests/test_accessors.py index bd50bc10..b69f743b 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -17,14 +17,22 @@ import dataclasses import itertools +import os import re from collections import OrderedDict, UserDict, UserList, defaultdict, deque +from operator import itemgetter from typing import Any, NamedTuple import pytest import optree -from helpers import TREE_ACCESSORS, SysFloatInfoType, assert_equal_type_and_value, parametrize +from helpers import ( + TREE_ACCESSORS, + SysFloatInfoType, + assert_equal_type_and_value, + parametrize, + skipif_pypy, +) def test_pytree_accessor_new(): @@ -228,6 +236,111 @@ def test_pytree_accessor_equal_hash(none_is_leaf): assert hash(accessor1) != hash(accessor2) +def test_pytree_accessor_not_equal_is_the_negation_of_equal(): + # Regression: `PyTreeAccessor` overrode `__eq__` but inherited `tuple.__ne__`, which compares + # element-wise and ignores the type check. Against a plain tuple holding the same entries, both + # `==` and `!=` were false, so a guard written with `!=` never fired. + entries = ( + optree.SequenceEntry(0, tuple, optree.PyTreeKind.TUPLE), + optree.MappingEntry('c', dict, optree.PyTreeKind.DICT), + ) + accessor = optree.PyTreeAccessor(entries) + + assert accessor == optree.PyTreeAccessor(entries) + assert not accessor != optree.PyTreeAccessor(entries) + assert accessor != optree.PyTreeAccessor(entries[:1]) + for other in (entries, list(entries), entries[:1], 'not-an-accessor'): + assert not accessor == other + assert accessor != other + assert not other == accessor + assert other != accessor + + # `PyTreeEntry` subclasses `object`, whose `__ne__` already negates `__eq__`. + for other in (entries, entries[0].entry, 'not-an-entry'): + assert (entries[0] == other) != (entries[0] != other) + assert (other == entries[0]) != (other != entries[0]) + + +def test_pytree_entry_equal_hash_with_non_function_call(): + # Regression: equality and hashing read `self.__class__.__call__.__code__`, which does not exist + # for a valid non-function implementation such as `itemgetter`, so both raised `AttributeError` + # on an entry that works perfectly well when called. + class ItemGetterEntry(optree.SequenceEntry): + __call__ = itemgetter(0) + + entry = ItemGetterEntry(0, list, optree.PyTreeKind.LIST) + assert entry([11, 22]) == 11 + assert entry == entry + assert isinstance(hash(entry), int) + assert len({entry, ItemGetterEntry(0, list, optree.PyTreeKind.LIST)}) == 1 + + +def test_pytree_entry_equal_hash_distinguishes_implementations(): + # Regression: comparing the bytecode alone collapsed entries whose implementations differ only + # in their defaults, so two entries that access different children compared equal and shared a + # hash bucket. + class FirstEntry(optree.SequenceEntry): + def __call__(self, obj, /, index=0): + return obj[index] + + class SecondEntry(optree.SequenceEntry): + def __call__(self, obj, /, index=1): + return obj[index] + + first = FirstEntry(0, list, optree.PyTreeKind.LIST) + second = SecondEntry(0, list, optree.PyTreeKind.LIST) + assert first([11, 22]) != second([11, 22]) + assert first != second + assert len({first, second}) == 2 + + # A subclass that does not override anything shares the very same implementations, so it stays + # equal to its base. + class Inherited(optree.SequenceEntry): + pass + + base = optree.SequenceEntry(0, list, optree.PyTreeKind.LIST) + assert Inherited(0, list, optree.PyTreeKind.LIST) == base + + +def test_getattr_entry_codify_non_identifier_name(): + # Regression: `codify()` emitted `node.x-y` for an attribute name that is not an identifier, + # which is not executable. `setattr` accepts any string, so fall back to `getattr` calls. + class Node: + pass + + node = Node() + setattr(node, 'x-y', 3) + entry = optree.GetAttrEntry('x-y', Node, optree.PyTreeKind.CUSTOM) + assert entry(node) == 3 + code = entry.codify('node') + assert code == "getattr(node, 'x-y')" + assert eval(code) == entry(node) + + # An identifier name, dotted or not, keeps the readable attribute form. + assert optree.GetAttrEntry('a', Node, optree.PyTreeKind.CUSTOM).codify('node') == 'node.a' + assert optree.GetAttrEntry('a.b', Node, optree.PyTreeKind.CUSTOM).codify('node') == 'node.a.b' + + # A mixed path only escapes the segments that need it. + mixed = optree.GetAttrEntry('a.x-y', Node, optree.PyTreeKind.CUSTOM) + assert mixed.codify('node') == "getattr(node.a, 'x-y')" + + +@skipif_pypy # PyPy names every sequence slot, so there is no unnamed slot to report +def test_structsequence_entry_repr_unnamed_slot(): + # An unnamed slot has no field name, so the repr shows the index it stands for rather than + # printing the marker as if it were a keyword argument. + st = os.stat_result(range(os.stat_result.n_fields)) + entries = [accessor[-1] for accessor in optree.tree_accessors(st)] + unnamed = [entry for entry in entries if entry.is_unnamed] + named = [entry for entry in entries if not entry.is_unnamed] + assert unnamed, 'expected at least one unnamed sequence slot' + + for entry in unnamed: + assert repr(entry) == f'StructSequenceEntry(entry={entry.entry!r}, type={os.stat_result!r})' + for entry in named: + assert repr(entry) == f'StructSequenceEntry(field={entry.field!r}, type={os.stat_result!r})' + + def test_pytree_entry_init(): for path_entry_type in ( optree.PyTreeEntry, diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py index 8a4bae98..5d53802a 100644 --- a/tests/test_dataclasses.py +++ b/tests/test_dataclasses.py @@ -932,6 +932,25 @@ def __post_init__(self): assert obj == optree.tree_unflatten(treespec, leaves) +def test_dataclass_post_init_rewriting_a_field_does_not_round_trip(): + # Characterization: the generated unflatten rebuilds via `cls(**fields)`, re-running `__init__` + # and `__post_init__`. A `__post_init__` that only derives non-field state round-trips (see + # `test_register_existing_class_with_metadata` above), but one that rewrites a stored field does + # not: flatten reads the rewritten value and unflatten rewrites it again. Such a class should + # register explicitly with custom flatten/unflatten (see `register_node`). + @optree.dataclasses.dataclass(namespace='test-dc-post-init-rewrite') + class Shift: + x: int + + def __post_init__(self): + self.x = self.x + 1 + + obj = Shift(5) # __post_init__: x = 5 + 1 = 6 + leaves, treespec = optree.tree_flatten(obj, namespace='test-dc-post-init-rewrite') + assert leaves == [6] + assert optree.tree_unflatten(treespec, leaves).x == 7 # re-applied: 6 + 1 = 7, not 6 + + def test_register_non_class(): with pytest.raises(TypeError, match='Expected a class'): optree.dataclasses.register_node(42, namespace='error') @@ -954,10 +973,125 @@ class Double: with pytest.raises( TypeError, - match=r'Cannot register .* as a pytree node more than once\.', + match=( + r'Cannot register .* as a pytree node more than once with ' + r'`optree\.dataclasses\.register_node\(\)`\. ' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)` ' + r'with explicit flatten/unflatten functions' + ), ): optree.dataclasses.register_node(Double, namespace='test-dc-double-2') + # As the error suggests, the class can still be registered in another namespace via the generic + # API with explicit flatten/unflatten functions. + optree.register_pytree_node( + Double, + lambda d: ((d.x,), None, None), + lambda _, children: Double(*children), + namespace='test-dc-double-2', + ) + optree.unregister_pytree_node(Double, namespace='test-dc-double-2') + + +def test_register_dataclass_with_initvar_rejected(): + # A dataclass with an `InitVar` cannot round-trip via the auto-generated flatten/unflatten: + # `InitVar`s are excluded from `dataclasses.fields()` (so they are neither flattened nor stored) + # yet remain required `__init__` parameters that `cls(**kwargs)` cannot restore. `register_node()` + # rejects it and points to the generic API with explicit flatten/unflatten functions. + @dataclasses.dataclass + class WithInitVar: + x: float + scale: dataclasses.InitVar[float] + scaled: float = optree.dataclasses.field(init=False, pytree_node=False, default=0.0) + + def __post_init__(self, scale): + self.scaled = self.x * scale + + with pytest.raises( + TypeError, + match=( + r'has `InitVar` field\(s\) .*cannot round-trip.*' + r'Use `optree\.register_pytree_node\(\)` or `optree\.register_pytree_node_class\(\)`' + ), + ): + optree.dataclasses.register_node(WithInitVar, namespace='test-dc-initvar') + + # The generic API with explicit flatten/unflatten works: keep the derived value as metadata and + # rebuild the instance without re-running `__init__`/`__post_init__`. + def flatten(obj): + return (obj.x,), obj.scaled + + def unflatten(scaled, children): + rebuilt = WithInitVar.__new__(WithInitVar) + rebuilt.x = children[0] + rebuilt.scaled = scaled + return rebuilt + + optree.register_pytree_node(WithInitVar, flatten, unflatten, namespace='test-dc-initvar') + obj = WithInitVar(3.0, 10.0) + leaves, treespec = optree.tree_flatten(obj, namespace='test-dc-initvar') + assert leaves == [3.0] + assert optree.tree_unflatten(treespec, leaves) == obj + optree.unregister_pytree_node(WithInitVar, namespace='test-dc-initvar') + + +def test_register_node_failure_does_not_leak_fields_guard(): + @dataclasses.dataclass + class Leak: + x: int + + namespace = 'test-dc-register-leak' + # Occupy the (class, namespace) slot directly so the dataclass `register_node()`'s internal + # `register_pytree_node()` call fails, after `register_node()` would set its `_FIELDS` guard. + optree.register_pytree_node( + Leak, + lambda leak: ((leak.x,), None, None), + lambda _, children: Leak(*children), + namespace=namespace, + ) + try: + with pytest.raises(ValueError, match='already registered'): + optree.dataclasses.register_node(Leak, namespace=namespace) + finally: + optree.unregister_pytree_node(Leak, namespace=namespace) + + # A failed registration must not leave the `_FIELDS` guard behind, or the class becomes + # impossible to register ever again: every retry would raise "... more than once". + # Re-registration after clearing the conflicting entry must succeed. + optree.dataclasses.register_node(Leak, namespace=namespace) + optree.unregister_pytree_node(Leak, namespace=namespace) + + +def test_register_node_rolls_back_when_the_guard_cannot_be_set(): + # `register_node()` commits the registry entry first and sets the `__optree_dataclass_fields__` + # guard afterwards. If the class refuses the guard, both must be rolled back: a class left + # registered but unguarded is just as unrecoverable, because the retry then fails with + # "already registered" instead. + class RejectGuardOnce(type): + rejected = False + + def __setattr__(cls, name, value, /): + if name == '__optree_dataclass_fields__' and not RejectGuardOnce.rejected: + RejectGuardOnce.rejected = True # a transient failure: reject only the first time + raise RuntimeError('cannot set the guard attribute') + super().__setattr__(name, value) + + @dataclasses.dataclass + class Rejecting(metaclass=RejectGuardOnce): + x: int + + namespace = 'test-dc-register-rollback' + with pytest.raises(RuntimeError, match=r'cannot set the guard attribute'): + optree.dataclasses.register_node(Rejecting, namespace=namespace) + + assert '__optree_dataclass_fields__' not in Rejecting.__dict__ + with pytest.raises(ValueError, match=r'is not registered'): + optree.unregister_pytree_node(Rejecting, namespace=namespace) + + # The class is left exactly as it was, so the retry succeeds. + optree.dataclasses.register_node(Rejecting, namespace=namespace) + optree.unregister_pytree_node(Rejecting, namespace=namespace) + def test_register_init_false_class_warns(): @dataclasses.dataclass(init=False) @@ -1131,6 +1265,60 @@ class EntryTest: assert 'DataclassEntry' in repr(entry_str) assert "'x'" in repr(entry_str) - entry_int = optree.DataclassEntry(1, EntryTest, optree.PyTreeKind.CUSTOM) - assert entry_int.field == 'y' - assert entry_int.name == 'y' + +def test_dataclass_entry_integer_indexes_children(): + # For a class registered with `optree.dataclasses.register_node()`, an integer entry indexes the + # children that registration emits (the fields that are BOTH `pytree_node=True` and `init`). A + # non-child field interleaved between children must not shift the mapping. + @optree.dataclasses.register_node(namespace='test-dataclass-entry-integer') + @dataclasses.dataclass + class Foo: + a: int + # metadata (init, not a child) + b: int = optree.dataclasses.field(default=0, pytree_node=False) + # non-init -> not a tree child + d: int = optree.dataclasses.field(init=False, default=0, pytree_node=False) + c: int = 0 + + foo = Foo(1, 2, 3) # a=1, b=2, c=3 (d defaults to 0, not an init parameter) + assert tuple(f.name for f in dataclasses.fields(Foo)) == ('a', 'b', 'd', 'c') + + entry_int = optree.DataclassEntry(1, Foo, optree.PyTreeKind.CUSTOM) + assert entry_int.init_fields == ('a', 'b', 'c') # `d` is not an init field + assert entry_int.children_fields == ('a', 'c') # `b` is metadata and `d` is non-init + # The 2nd child is `c` (not the metadata field `b` nor the non-init field `d`). + assert entry_int.field == 'c' + assert entry_int.name == 'c' + assert entry_int.codify('x') == 'x.c' + assert entry_int(foo) == foo.c == 3 + + +def test_dataclass_entry_integer_follows_generic_flatten_func(): + # Regression: an integer entry was resolved with the `optree.dataclasses` field predicate, which + # does not apply to a class registered with the generic `register_pytree_node()`. There the + # flatten function alone decides the children, so `pytree_node=False` must not drop a field and + # leave the entry indexing past the end of an internal list. + @dataclasses.dataclass + class Foo: + a: int + b: int = optree.dataclasses.field(default=0, pytree_node=False) + + optree.register_pytree_node( + Foo, + lambda foo: ((foo.a, foo.b), None, None), # both fields are children, entries default + lambda _, children: Foo(*children), + namespace='test-dataclass-entry-generic', + ) + + foo = Foo(1, 2) + paths, leaves, treespec = optree.tree_flatten_with_path( + foo, + namespace='test-dataclass-entry-generic', + ) + assert paths == [(0,), (1,)] + assert leaves == [1, 2] + + accessors = treespec.accessors() + assert [accessor(foo) for accessor in accessors] == [1, 2] + assert [accessor[0].field for accessor in accessors] == ['a', 'b'] + assert [accessor.codify('x') for accessor in accessors] == ['x.a', 'x.b'] diff --git a/tests/test_ops.py b/tests/test_ops.py index 04fd8e89..08ddcbc9 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -52,6 +52,7 @@ skipif_ios, skipif_wasm, ) +from optree.utils import total_order_sorted @skipif_wasm @@ -187,6 +188,48 @@ def test_flatten_dict_order(): assert list(restored_tree) == ['b', 'a', 'c'] +def test_flatten_dict_order_keys_that_fail_to_compare_midway(): + # Regression: the C++ `TotalOrderSort` sorted the key list IN PLACE, caught the `TypeError`, and + # then ran the total-order fallback on the same, already partially reordered list. `list.sort` + # leaves the list in an undefined order when a comparison raises, so a second failure committed + # that partial order instead of the documented insertion order. Each attempt must sort a copy. + # + # `Key.__lt__` raises only for the single pair that `list.sort` reaches after it has already + # moved elements, so both the direct sort and the total-order fallback fail on the same list. + class Key: + def __init__(self, value): + self.value = value + + def __repr__(self): + return f'Key({self.value})' + + def __eq__(self, other): + return isinstance(other, Key) and self.value == other.value + + def __hash__(self): + return hash(self.value) + + def __lt__(self, other): + if not isinstance(other, Key): + return NotImplemented + if (self.value, other.value) == (2, 1): + raise TypeError('Key(2) < Key(1) is undefined') + return self.value < other.value + + tree = {Key(1): 1, Key(0): 0, Key(2): 2} + insertion_order = [Key(1), Key(0), Key(2)] + + assert total_order_sorted(tree) == insertion_order + _, metadata, entries, _ = optree.tree_flatten_one_level(tree) + assert metadata == insertion_order + assert list(entries) == insertion_order + + leaves, treespec = optree.tree_flatten(tree) + assert treespec.entries() == insertion_order # the recursive walk must agree with one level + assert leaves == [1, 0, 2] + assert optree.tree_unflatten(treespec, leaves) == tree + + @parametrize( tree=list(TREES + LEAVES), none_is_leaf=[False, True], @@ -227,6 +270,44 @@ def test_tree_iter( next(it) +@skipif_wasm +@skipif_android +@skipif_ios +def test_tree_iter_reentrant_next(): + # Regression: `PyTreeIter::Next()` takes a non-recursive mutex and then runs user code + # (`is_leaf` and any custom `flatten_func`). Calling `next()` on the same iterator from there + # blocked on a mutex the thread already held, hanging the interpreter uninterruptibly. It must + # raise instead, mirroring CPython's "generator already executing". Run under a watchdog in a + # subprocess so a regression fails instead of wedging the test session. + check_script_in_subprocess( + """ + import faulthandler + + import optree + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + it = None + + def is_leaf(x): + if x == 1: + try: + next(it) + except RuntimeError as ex: + assert str(ex) == 'PyTreeIter is already iterating.', ex + else: + raise AssertionError('re-entrant next() did not raise') + return False + + it = optree.tree_iter([1, 2, 3], is_leaf=is_leaf) + assert list(it) == [1, 2, 3] + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + def test_traverse(): tree = {'b': 2, 'a': 1, 'c': {'f': None, 'e': 3, 'g': 4}} # tree @@ -480,6 +561,54 @@ def test_flatten_up_to_none_is_leaf(): assert subtrees == [accessor(tree) for accessor in optree.treespec_accessors(treespec)] +@skipif_wasm +@skipif_android +@skipif_ios +def test_flatten_up_to_does_not_expose_a_partially_filled_leaves_list(): + # Regression: `flatten_up_to` filled a pre-sized list, which is GC-tracked with NULL slots from + # the moment it is created, while running user code (a custom `flatten_func`, a dict key's + # `__hash__`/`__eq__`, `__repr__` on the error paths). Such code can reach the list through + # `gc.get_objects()`, and reading an unwritten slot crashes the interpreter. Run in a subprocess: + # a regression aborts rather than fails. + check_script_in_subprocess( + """ + import gc + + import optree + + NUM_LEAVES = 1237 + captured = [] + + class Key: + def __hash__(self): + captured.extend( + obj for obj in gc.get_objects() + if type(obj) is list and len(obj) == NUM_LEAVES + ) + return 0 + + def __eq__(self, other): + return isinstance(other, Key) + + treespec = optree.tree_structure(({Key(): (1, 2, 3)}, tuple(range(NUM_LEAVES - 3)))) + assert treespec.num_leaves == NUM_LEAVES + # An arity mismatch deep in the tree: the walk runs `Key.__hash__` well before the last leaf. + try: + treespec.flatten_up_to(({Key(): (1, 2)}, tuple(range(NUM_LEAVES - 3)))) + except ValueError: + pass + else: + raise AssertionError('flatten_up_to did not report the arity mismatch') + + for lst in captured: + list(lst) # a NULL slot here crashes or raises `SystemError` + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + @parametrize( leaves_fn=[ optree.tree_leaves, @@ -509,6 +638,23 @@ def test_flatten_is_leaf(leaves_fn): assert leaves == [1, 2, 3, 4, None, 5, None] +def test_flatten_is_leaf_shrinking_the_list(): + # Regression: the flattener reads the list length once and then indexes it while running + # `is_leaf` in between. Shrinking the list from the predicate used to read past the end on + # Python < 3.13, which uses the unchecked `PyList_GET_ITEM`. It must raise `IndexError` on every + # supported version. + lst = [0, 1, 2, 3, 4, 5, 6, 7] + + def is_leaf(x): + if x == 0: + del lst[1:] # shrink while the flattener is looping over indices 0..7 + return False + + with pytest.raises(IndexError, match=re.escape('list index out of range')): + optree.tree_flatten(lst, is_leaf=is_leaf) + assert lst == [0] + + @parametrize( structure_fn=[ optree.tree_structure, @@ -2824,6 +2970,45 @@ def test_tree_partition(fillvalue, none_is_leaf): assert right == {'x': 7, 'y': (fillvalue, fillvalue)} +def test_tree_partition_with_no_leaves(): + # Regression: the zero-leaf guard belongs to inferring the inner structure from the first + # output. `tree_partition` passes an explicit inner treespec, so a leafless tree partitions into + # two copies of itself and the predicate is never called. The guard also hid a second bug: + # transposing zero outputs yielded no subtrees at all rather than one empty subtree per inner + # leaf, so removing it alone raised `Too few leaves for PyTreeSpec`. + for tree in [(), [], {}, {'a': [], 'b': ()}, [(), {}]]: + calls = [] + left, right = optree.tree_partition( + lambda x, calls=calls: calls.append(x) or True, + tree, + ) + assert left == tree, (tree, left) + assert right == tree, (tree, right) + assert calls == [], (tree, calls) + + # A tree whose only leaf is `None` has no leaves at all unless `none_is_leaf` is set, so the + # predicate is never called and both halves keep the `None`. + calls = [] + left, right = optree.tree_partition(lambda x: calls.append(x) or True, {'a': None}) + assert (left, right) == ({'a': None}, {'a': None}) + assert calls == [] + # With `none_is_leaf`, `None` is a real leaf: the predicate runs and the fill value is `None` + # either way, so both halves still read as `{'a': None}`. + calls = [] + left, right = optree.tree_partition( + lambda x: calls.append(x) or True, + {'a': None}, + none_is_leaf=True, + ) + assert (left, right) == ({'a': None}, {'a': None}) + assert calls == [None] + + # Non-empty trees are unaffected. + left, right = optree.tree_partition(lambda x: x > 1, [1, 2, 3]) + assert left == [None, 2, 3] + assert right == [1, None, None] + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -3120,6 +3305,207 @@ def test_tree_broadcast_common(): ) +def test_tree_broadcast_common_does_not_apply_is_leaf_to_internal_sentinel(): + # `tree_broadcast_common` fills an internal `common_suffix_tree` with a private `object()` + # sentinel, then replicates each leaf across the matching subtree. The user's `is_leaf` must run + # only on values from the input trees, never the sentinel: a value-dependent predicate would + # otherwise mis-structure the placeholder subtree or crash while inspecting it. + + # Collapses a list unless every element is an int. When the predicate keeps the other's list + # (all ints), the fix must not re-apply the predicate to the sentinel subtree, which would + # collapse `[, ]` and under-replicate. When the predicate collapses the real other + # tree to a leaf (a list holding a tuple), there is simply nothing to broadcast and the inputs + # pass through unchanged. + def collapse_non_int_lists(x): + return isinstance(x, list) and not all(isinstance(e, int) for e in x) + + assert optree.tree_broadcast_common( + 1, + [2, 3], + is_leaf=collapse_non_int_lists, + ) == ([1, 1], [2, 3]) + assert optree.tree_broadcast_common( + 1, + [2, (3, 4)], + is_leaf=collapse_non_int_lists, + ) == (1, [2, (3, 4)]) + assert optree.tree_broadcast_common( + {'a': 1}, + {'a': [2, 3]}, + is_leaf=collapse_non_int_lists, + ) == ({'a': [1, 1]}, {'a': [2, 3]}) + assert optree.tree_broadcast_common( + {'a': 1}, + {'a': [2, (3, 4)]}, + is_leaf=collapse_non_int_lists, + ) == ({'a': 1}, {'a': [2, (3, 4)]}) + + # Inspects the value; must not crash on the sentinel (`object() < 0` raises `TypeError`), + # regardless of the container (list, tuple, nested) the sentinel subtree takes. + def negative_scalar(x): + return not isinstance(x, (list, tuple, dict)) and x < 0 + + assert optree.tree_broadcast_common( + 1, + [2, 3], + is_leaf=negative_scalar, + ) == ([1, 1], [2, 3]) + assert optree.tree_broadcast_common( + [1], + [(2, 3)], + is_leaf=negative_scalar, + ) == ([(1, 1)], [(2, 3)]) + assert optree.tree_broadcast_common( + [1, 1], + [[2], [3, 4]], + is_leaf=negative_scalar, + ) == ([[1], [1, 1]], [[2], [3, 4]]) + # `broadcast_common` delegates to `tree_broadcast_common`, so it must not crash either. + assert optree.broadcast_common(1, [2, 3], is_leaf=negative_scalar) == ([1, 1], [2, 3]) + + +def test_tree_broadcast_common_value_dependent_is_leaf_is_lossy(): + # Characterization, separate from the R18 sentinel fix: a `is_leaf` that classifies by leaf + # *value* rather than type/structure can be lossy under broadcasting. `is_shape_like` treats a + # tuple of ints as one leaf, so broadcasting the scalar `5` into a two-slot yields `(5, 5)`, + # which the predicate re-reads as a single shape. Such a result still shares the literal + # common-suffix structure, but re-flattens to a different leaf count under the predicate. The + # cases below show it is lossy for a number yet faithful for a shape or a list slot. This pins + # the known limitation; prefer type/structure predicates for broadcasting. + def is_shape_like(x): + return type(x) is tuple and all(isinstance(v, int) for v in x) + + # A tuple slot is lossy: `(5, 5)` re-reads as a shape leaf. + assert optree.tree_broadcast_common( + [5], + [(1, 'x')], + is_leaf=is_shape_like, + ) == ([(5, 5)], [(1, 'x')]) + # A list slot is not: `[5, 5]` is not a shape tuple, so it stays consistent. + assert optree.tree_broadcast_common( + [5], + [[1, 2]], + is_leaf=is_shape_like, + ) == ([[5, 5]], [[1, 2]]) + + # A mixed tree shows both outcomes at once. The shape `(5, 5)` broadcast into a two-slot becomes + # `((5, 5), (5, 5))`, two shape leaves, faithful; the number `3` broadcast the same way becomes + # `(3, 3)`, which the predicate re-reads as one shape leaf, lossy. + assert optree.tree_broadcast_common( + [(5, 5), (1.0, 2)], + [(2, 'x'), 3], + is_leaf=is_shape_like, + ) == ([((5, 5), (5, 5)), (1.0, 2)], [(2, 'x'), (3, 3)]) + + +def test_tree_broadcast_map_does_not_reapply_is_leaf_to_broadcast_subtrees(): + # Regression: the broadcast result was re-flattened with the caller's `is_leaf`. A predicate + # that classifies by leaf VALUE then treated a subtree broadcasting had just created as a leaf, + # so `func` was called once on the whole subtree instead of once per common-suffix leaf. The + # common structure now comes from the input trees, where `is_leaf` describes the caller's data. + def is_shape_like(x): + return type(x) is tuple and all(isinstance(v, int) for v in x) + + calls = [] + + def func(x, y): + calls.append((x, y)) + return f'{x}:{y}' + + assert optree.tree_broadcast_map( + func, + [5], + [(1, 'x')], + is_leaf=is_shape_like, + ) == [('5:1', '5:x')] + assert calls == [(5, 1), (5, 'x')] + + paths = [] + assert optree.tree_broadcast_map_with_path( + lambda p, x, y: (paths.append(p), f'{x}:{y}')[1], + [5], + [(1, 'x')], + is_leaf=is_shape_like, + ) == [('5:1', '5:x')] + assert paths == [(0, 0), (0, 1)] + + accessors = [] + assert optree.tree_broadcast_map_with_accessor( + lambda a, x, y: (accessors.append(len(a)), f'{x}:{y}')[1], + [5], + [(1, 'x')], + is_leaf=is_shape_like, + ) == [('5:1', '5:x')] + assert accessors == [2, 2] + + +def test_tree_broadcast_map_unchanged_for_structural_predicates(): + # The common structure still honors `is_leaf` where it describes the INPUT trees, and plain + # broadcasting is unaffected by the change. + assert optree.tree_broadcast_map( + lambda x, y: x + y, + [1, 2], + [[3, 4], [5, 6]], + ) == [[4, 5], [7, 8]] + assert optree.tree_broadcast_map(lambda x: x * 2, {'a': 1, 'b': 2}) == {'a': 2, 'b': 4} + assert optree.tree_broadcast_map( + lambda a, b, c: a + b + c, + [1], + [[2, 3]], + [[4, 5]], + ) == [[7, 9]] + # A type-based predicate keeps a marked subtree intact on both sides. + assert optree.tree_broadcast_map( + lambda x, y: (x, y), + [1], + [[2, 3]], + is_leaf=lambda x: isinstance(x, list) and len(x) == 2, + ) == [(1, [2, 3])] + + +def test_tree_broadcast_map_single_input_flattens_only_once(): + # Regression: with no `rests` to broadcast against, the tree was flattened twice (once to derive + # the common structure, once by `flatten_up_to`). The docstrings promise equivalence to + # `tree_map`, which flattens once, so a one-shot `flatten_func` broke here but not there. + counter = Counter() + + class OneShot: + def __init__(self, x): + self.x = x + + def __eq__(self, other): + return isinstance(other, OneShot) and self.x == other.x + + def flatten_func(node): + if counter.increment() > 1: + raise RuntimeError(f'`flatten_func` called {counter.count} times') + return (node.x,), None + + namespace = 'broadcast-map-one-shot' + optree.register_pytree_node( + OneShot, + flatten_func, + lambda _, children: OneShot(*children), + namespace=namespace, + ) + + for name, func in ( + ('tree_broadcast_map', lambda x: x + 1), + ('tree_broadcast_map_with_path', lambda p, x: (p, x + 1)), + ('tree_broadcast_map_with_accessor', lambda a, x: (a, x + 1)), + ): + broadcast_map = getattr(optree, name) + map_ = getattr(optree, name.replace('tree_broadcast_map', 'tree_map')) + + counter.count = 0 + expected = map_(func, OneShot(1), namespace=namespace) + assert counter.count == 1, name + + counter.count = 0 + assert broadcast_map(func, OneShot(1), namespace=namespace) == expected, name + assert counter.count == 1, name + + def test_broadcast_common(): assert optree.broadcast_common(1, [2, 3, 4]) == ([1, 1, 1], [2, 3, 4]) assert optree.broadcast_common([1, 2, 3], [4, 5, 6]) == ([1, 2, 3], [4, 5, 6]) diff --git a/tests/test_prefix_errors.py b/tests/test_prefix_errors.py index 3f4e1f17..9447c29b 100644 --- a/tests/test_prefix_errors.py +++ b/tests/test_prefix_errors.py @@ -390,6 +390,46 @@ def test_different_metadata(): raise e('in_axes') +def test_different_metadata_heterogeneous_keys(): + # Regression: formatting the key difference sorted the keys directly, so a dictionary with + # mutually incomparable keys raised a `TypeError` out of `prefix_errors()` instead of returning + # the error factories. Flattening supports such keys, so the message must too. + lhs, rhs = {1: 1, 'a': 2}, {2: 3, 'b': 4} + lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs) + with pytest.raises( + ValueError, + match=r'dictionary key mismatch; expected key\(s\): .*, got key\(s\): .*; dict: .*\.', + ): + optree.tree_map_(lambda x, y: None, lhs, rhs) + assert not lhs_treespec.is_prefix(rhs_treespec) + + (e,) = optree.prefix_errors(lhs, rhs) + # `total_order_sorted()` groups by `{module}.{qualname}`, so `builtins.int` precedes + # `builtins.str`. + expected = re.escape( + textwrap.dedent( + """ + pytree structure error: different pytree keys at key path + in_axes tree root + At that key path, the prefix pytree in_axes has a subtree of type + + with 2 key(s) + [1, 'a'] + but at the same key path the full pytree has a subtree of type + + but with 2 key(s) + [2, 'b'] + missing key(s): + [1, 'a'] + extra key(s): + [2, 'b'] + """, + ).strip(), + ) + with pytest.raises(ValueError, match=expected): + raise e('in_axes') + + def test_different_metadata_nested(): lhs, rhs = [{1: 2}], [{3: 4}] lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs) @@ -607,6 +647,37 @@ def test_no_errors(): () = optree.prefix_errors(lhs, rhs) +def test_no_errors_for_custom_node_with_value_dependent_entries(): + # Regression: a bare `assert` fired on a VALID prefix pair whose custom node reported different + # per-instance entries while its metadata agreed. `entries` is an independent element of the + # flatten output, so that registration is legal and `broadcast_prefix` accepts it; + # `prefix_errors` must return an empty list rather than raise (and it was elided under `-O`). + class Tagged: + def __init__(self, tags, values): + self.tags = list(tags) + self.values = list(values) + + optree.register_pytree_node( + Tagged, + lambda tagged: (tagged.values, None, tagged.tags), # metadata is always None, entries vary + lambda metadata, children: Tagged(range(len(list(children))), children), + namespace='value_dependent_entries', + ) + try: + kwargs = {'namespace': 'value_dependent_entries'} + lhs = Tagged(['x', 'y'], [1, 2]) + rhs = Tagged(['p', 'q'], [(1, 1), (2, 2)]) + assert optree.broadcast_prefix(lhs, rhs, **kwargs) == [1, 1, 2, 2] + () = optree.prefix_errors(lhs, rhs, **kwargs) + + # A genuine arity mismatch is still reported. + (e,) = optree.prefix_errors(lhs, Tagged(['x', 'y', 'z'], [1, 2, 3]), **kwargs) + with pytest.raises(ValueError, match=r'different numbers of pytree children'): + raise e('in_axes') + finally: + optree.unregister_pytree_node(Tagged, namespace='value_dependent_entries') + + def test_different_structure_no_children(): (e,) = optree.prefix_errors((), ([],)) expected = re.escape( diff --git a/tests/test_registry.py b/tests/test_registry.py index 3900b3e5..0a8a4e02 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -19,6 +19,7 @@ import pickle import re import sys +import warnings import weakref from collections import UserDict, UserList, namedtuple from dataclasses import dataclass @@ -32,10 +33,14 @@ NODETYPE_REGISTRY, PYPY, Py_GIL_DISABLED, + check_script_in_subprocess, disable_systrace, gc_collect, parametrize, + skipif_android, + skipif_ios, skipif_pypy, + skipif_wasm, ) from optree.registry import DictMetaData @@ -135,6 +140,36 @@ def __tree_unflatten__(cls, metadata, children): optree.register_pytree_node_class(namespace='mylist3')(MyList1) +def test_register_pytree_node_duplicate_registration_does_not_warn(): + # A rejected re-registration overrides nothing, so it must raise and emit no override warning. + mytuple = namedtuple('mytuple_duplicate', ['a', 'b']) # noqa: PYI024 + + with pytest.warns(UserWarning, match=re.escape('is a subclass of `collections.namedtuple`')): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_duplicate', + ) + + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + with pytest.raises( + ValueError, + match=r"PyTree type.*is already registered in namespace 'mytuple_duplicate'\.", + ): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_duplicate', + ) + assert [str(w.message) for w in caught] == [] + finally: + optree.unregister_pytree_node(mytuple, namespace='mytuple_duplicate') + + def test_register_pytree_node_with_invalid_namespace(): with pytest.raises(TypeError, match='The namespace must be a string'): @@ -401,6 +436,334 @@ def test_register_pytree_node_namedtuple(): assert treespec1 != treespec3 +def test_register_pytree_node_warning_as_error_does_not_corrupt_registry(): + # Registering a `namedtuple` / `PyStructSequence` subclass emits a `UserWarning`. Under + # warnings-as-errors that warning is raised as an exception; previously the C++ registry had + # already committed the registration and ignored the escalation, leaving corrupt state and a + # `SystemError`. The registration must instead be rolled back and the escalated warning raised. + mytuple = namedtuple('mytuple_warn_error', ['a', 'b']) # noqa: PYI024 + tree = mytuple(1, 2) + + with warnings.catch_warnings(): + warnings.simplefilter('error') + with pytest.raises(UserWarning): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_warn_error', + ) + + # The escalated registration must not have stuck (no custom node registered). + assert 'CustomTreeNode' not in str( + optree.tree_structure(tree, namespace='mytuple_warn_error'), + ) + + # A subsequent non-error registration must succeed cleanly, proving no half-committed state. + with pytest.warns(UserWarning, match=re.escape('is a subclass of `collections.namedtuple`')): + optree.register_pytree_node( + mytuple, + lambda t: (tuple(t), None, None), + lambda _, t: mytuple(*t), + namespace='mytuple_warn_error', + ) + assert 'CustomTreeNode' in str( + optree.tree_structure(tree, namespace='mytuple_warn_error'), + ) + optree.unregister_pytree_node(mytuple, namespace='mytuple_warn_error') + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_register_pytree_node_no_deadlock_with_concurrent_flatten(): + # Regression: `register_pytree_node` takes the registry write lock, and the namedtuple / + # PyStructSequence detection it runs releases the GIL. Holding the registry lock across that GIL + # release inverts the GIL <-> lock order and deadlocks a concurrent `tree_flatten` that holds + # the GIL while waiting on the registry read lock, wedging the whole interpreter. Run register / + # unregister concurrently with flatten in a subprocess under a watchdog; the process must finish. + check_script_in_subprocess( + """ + import faulthandler + import threading + + import optree + + class MyType: + def __init__(self, children): + self.children = children + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + def register_loop(): + for _ in range(5000): + try: + optree.register_pytree_node( + MyType, + lambda o: (tuple(o.children), None, None), + lambda _, c: MyType(list(c)), + namespace='deadlock', + ) + except ValueError: + pass + try: + optree.unregister_pytree_node(MyType, namespace='deadlock') + except ValueError: + pass + + def flatten_loop(): + tree = {'a': 1, 'b': 2, 'c': 3} + for _ in range(5000): + optree.tree_flatten(tree) + + threads = [ + threading.Thread(target=register_loop), + threading.Thread(target=flatten_loop), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_unregister_pytree_node_not_found_no_deadlock_with_concurrent_flatten(): + # Regression: `unregister_pytree_node` takes the registry write lock, and on the NOT-FOUND path + # it builds its error message using the namedtuple / PyStructSequence detection, which releases + # the GIL. Holding the registry lock across that GIL release inverts the GIL <-> lock order and + # deadlocks a concurrent `tree_flatten` that holds the GIL while waiting on the registry read + # lock (the symmetric twin of the `register_pytree_node` deadlock). Unregister a never-registered + # type (always not-found) concurrently with flatten in a subprocess under a watchdog; the process + # must finish. + check_script_in_subprocess( + """ + import faulthandler + import threading + + import optree + + class NeverRegistered: # never registered -> unregister always hits the not-found path + pass + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + def unregister_loop(): + for _ in range(5000): + try: + optree.unregister_pytree_node(NeverRegistered, namespace='deadlock') + except ValueError: + pass + + def flatten_loop(): + tree = {'a': 1, 'b': 2, 'c': 3} + for _ in range(5000): + optree.tree_flatten(tree) + + threads = [threading.Thread(target=unregister_loop) for _ in range(2)] + threads += [threading.Thread(target=flatten_loop) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_registry_error_paths_no_deadlock_with_reentrant_repr(): + # Regression: the duplicate-registration and not-found error messages call `repr()` on the type + # while the registry WRITE lock is held. A metaclass `__repr__` that re-enters optree then takes + # the registry READ lock on the same thread, and the lock is not recursive, so it self-deadlocked + # before the error was even raised. Run under a watchdog in a subprocess. + check_script_in_subprocess( + """ + import faulthandler + + import optree + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + class Meta(type): + def __repr__(cls): + optree.tree_flatten({'a': 1}) # -> registry lookup -> registry read lock + return f'<{cls.__name__}>' + + class MyType(metaclass=Meta): + pass + + class NeverRegistered(metaclass=Meta): + pass + + optree.register_pytree_node( + MyType, + lambda obj: ((), None, None), + lambda metadata, children: MyType(), + namespace='reentrant_repr', + ) + try: + optree.register_pytree_node( + MyType, + lambda obj: ((), None, None), + lambda metadata, children: MyType(), + namespace='reentrant_repr', + ) + except ValueError as ex: + assert '' in str(ex), ex + else: + raise AssertionError('duplicate registration did not raise') + + try: + optree.unregister_pytree_node(NeverRegistered, namespace='reentrant_repr') + except ValueError as ex: + assert '' in str(ex), ex + else: + raise AssertionError('unregistering an unregistered type did not raise') + + optree.unregister_pytree_node(MyType, namespace='reentrant_repr') + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +@skipif_pypy # `ref() is None` right after `unregister_node` needs CPython refcounting +def test_unregister_node_no_deadlock_with_reentrant_weakref_callback(): + # Regression: `Unregister` dropped the last references to the registration's Python objects + # while still holding the registry WRITE lock. That runs arbitrary user code (`__del__`, weakref + # callbacks), and a callback re-entering optree took the registry READ lock on the same + # non-recursive lock. Register through `_C` so the C++ registration holds the only references. + # Run under a watchdog in a subprocess. + check_script_in_subprocess( + """ + import faulthandler + import weakref + + import optree + import optree._C as _C + from optree.accessors import PyTreeEntry + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + class Foo: + pass + + def make_funcs(): + return lambda obj: ((), None, None), lambda metadata, children: Foo() + + flatten_func, unflatten_func = make_funcs() + _C.register_node(Foo, flatten_func, unflatten_func, PyTreeEntry, 'reentrant_decref') + + fired = [] + + def on_death(ref): + fired.append(optree.tree_flatten({'a': 1})) # -> registry lookup -> read lock + + ref = weakref.ref(flatten_func, on_death) + del flatten_func, unflatten_func + + _C.unregister_node(Foo, 'reentrant_decref') + assert ref() is None + assert len(fired) == 1, fired + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +def test_get_registry_size_concurrent_no_spurious_error(): + # Regression: `GetRegistrySize()` reads the two internal registries (NoneIsNode / NoneIsLeaf) + # and checks they differ by exactly one (the extra `None` node). Reading them in two separate + # locked `Size()` calls dropped the lock between the reads, so a concurrent (un)registration + # could slip in and spuriously trip that check, surfacing as a `SystemError`. This manifests on + # free-threading builds (the GIL otherwise serializes the two reads). Hammer `get_registry_size()` + # while a type churns in and out of the registry, in a subprocess under a watchdog; the process + # must finish with no error recorded. + check_script_in_subprocess( + """ + import faulthandler + import threading + import time + + import optree + import optree._C as _C + + class MyType: + def __init__(self, children): + self.children = children + + faulthandler.dump_traceback_later(30, exit=True) # watchdog: abort the process on a hang + + errors = [] + stop = threading.Event() + + def register_loop(): + for _ in range(5000): + try: + optree.register_pytree_node( + MyType, + lambda o: (tuple(o.children), None, None), + lambda _, c: MyType(list(c)), + namespace='size-race', + ) + except ValueError: + pass + try: + optree.unregister_pytree_node(MyType, namespace='size-race') + except ValueError: + pass + stop.set() + + def size_loop(): + while not stop.is_set(): + try: + _C.get_registry_size() + _C.get_registry_size('size-race') + except Exception as exc: # noqa: BLE001 + errors.append(repr(exc)) + return + # `get_registry_size()` takes the registry lock in read mode, so spinning here keeps + # overlapping shared holds on it and starves `register_loop`: `stop` is never set + # and the watchdog fires. The sleep must be non-zero, since `time.sleep(0)` maps to + # `Sleep(0)` on Windows and can return without ever dropping the lock. + time.sleep(0.001) + + threads = [ + threading.Thread(target=register_loop), + threading.Thread(target=size_loop), + threading.Thread(target=size_loop), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not errors, errors + print('COMPLETED') + """, + output='COMPLETED', + rstrip=True, + ) + + def test_flatten_with_wrong_number_of_returns(): @optree.register_pytree_node_class(namespace='error') class MyList1(UserList): @@ -579,6 +942,36 @@ def __tree_unflatten__(cls, metadata, children): assert handler is registry3[cls] +def test_pytree_node_registry_get_named_namespace_takes_precedence_over_global(): + # When a type is registered in BOTH the global namespace and a named namespace, the dict form of + # `.get(namespace=...)` must return the NAMED handler for that type, matching the single-class + # lookup, which checks the named namespace before falling back to the global one. Registering + # the global entry LAST would otherwise let it shadow the named one in the returned dict. + class Both: + def __init__(self, a): + self.a = a + + def flatten(both): + return (both.a,), None, None + + def unflatten(metadata, children): + return Both(*children) + + optree.register_pytree_node(Both, flatten, unflatten, namespace='both-ns') + optree.register_pytree_node(Both, flatten, unflatten, namespace=GLOBAL_NAMESPACE) + try: + # Single-class lookup already prefers the named namespace. + assert optree.register_pytree_node.get(Both, namespace='both-ns').namespace == 'both-ns' + # The dict form must agree: the named handler wins over the global one. + registry = optree.register_pytree_node.get(namespace='both-ns') + assert registry[Both].namespace == 'both-ns' + # The global-only view still shows the global handler. + assert optree.register_pytree_node.get(namespace=GLOBAL_NAMESPACE)[Both].namespace == '' + finally: + optree.unregister_pytree_node(Both, namespace='both-ns') + optree.unregister_pytree_node(Both, namespace=GLOBAL_NAMESPACE) + + def test_pytree_node_registry_get_with_invalid_arguments(): registry = optree.register_pytree_node.get() assert optree.register_pytree_node.get(None) == registry diff --git a/tests/test_treespec.py b/tests/test_treespec.py index 643176a0..f068edf3 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -16,7 +16,10 @@ # pylint: disable=missing-function-docstring,invalid-name import contextlib +import gc import itertools +import math +import os import pickle import platform import re @@ -24,8 +27,10 @@ import subprocess import sys import tempfile +import time +import warnings import weakref -from collections import OrderedDict, UserList, defaultdict, deque +from collections import OrderedDict, UserList, defaultdict, deque, namedtuple import pytest @@ -33,6 +38,7 @@ import optree from helpers import ( GLOBAL_NAMESPACE, + HAS_DEFERRED_TYPE_REFS, NAMESPACED_TREE, PYPY, STANDARD_DICT_TYPES, @@ -48,6 +54,7 @@ parametrize, recursionlimit, skipif_android, + skipif_deferred_type_refs, skipif_ios, skipif_pypy, skipif_wasm, @@ -260,6 +267,25 @@ def test_treespec_string_representation(data): assert new_tree == reconstructed_tree +@skipif_pypy # CPython-only: `os.stat_result` slots 7, 8, 9 are unnamed; PyPy names them +def test_treespec_structseq_unnamed_field_string_representation(): + # `os.stat_result` renders its UNNAMED sequence slots (7, 8, 9) with the synthetic `` + # placeholder, following CPython's `` convention for names that are not identifiers, + # rather than the bare `unnamed field` marker which reads as an invalid keyword. CPython's + # `stat_result_desc` has pinned the 7 named + 3 unnamed sequence fields for 16 years, so the + # repr is asserted exactly; the hidden float `st_atime` fields (indices >= 10) are not part of + # the sequence and must not leak into it. + assert os.stat_result.n_sequence_fields == 10 + assert os.stat_result.n_unnamed_fields == 3 + st = os.stat_result(range(os.stat_result.n_fields)) + representation = str(optree.tree_structure(st)) + assert representation == ( + 'PyTreeSpec(os.stat_result(' + 'st_mode=*, st_ino=*, st_dev=*, st_nlink=*, st_uid=*, st_gid=*, st_size=*, ' + '=*, =*, =*))' + ) + + def test_treespec_with_empty_tuple_string_representation(): assert str(optree.tree_structure(())) == r'PyTreeSpec(())' @@ -276,6 +302,50 @@ def test_treespec_with_empty_dict_string_representation(): assert str(optree.tree_structure({})) == r'PyTreeSpec({})' +def test_treespec_namedtuple_repr_with_divergent_fields_raises_value_error(): + # If a namedtuple's `_fields` is mutated after the treespec is built, the recorded arity and the + # now-divergent field count disagree. The repr must raise a clear `ValueError` attributing the + # cause, not an `InternalError` telling the user to file a bug report. + Point = namedtuple('Point', ('x', 'y')) # noqa: PYI024 + treespec = optree.tree_structure(Point(1, 2)) + assert str(treespec) == 'PyTreeSpec(Point(x=*, y=*))' + + Point._fields = ('x', 'y', 'z') # diverge: 3 fields vs the treespec's arity of 2 + with pytest.raises(ValueError, match=r'does not match the arity'): + repr(treespec) + + +def test_treespec_setstate_rejects_structseq_field_arity_mismatch(): + # A PyStructSequence type's sequence-field count is fixed in C, so a node's arity must equal it + # (unlike a namedtuple, whose `_fields` can be mutated after the fact). `FromPicklable` (via + # `__setstate__`/`pickle`) must reject a crafted state pairing a PyStructSequence type with a + # mismatched arity at load time, rather than build a corrupt treespec that later aborts (e.g. in + # repr with an `InternalError`). + spec = optree.tree_structure(time.gmtime()) # struct_time: 9 sequence fields + node_states, none_is_leaf, namespace = spec.__getstate__() + # Swap the type to os.stat_result (10 sequence fields) while keeping the arity of 9. + crafted = tuple( + (kind, arity, os.stat_result if data is time.struct_time else data, *remaining) + for (kind, arity, data, *remaining) in node_states + ) + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + with pytest.raises(RuntimeError, match=r'does not match the arity'): + obj.__setstate__((crafted, none_is_leaf, namespace)) + + +def test_treespec_setstate_rejects_namedtuple_field_arity_mismatch(): + # A namedtuple's `_fields` can be mutated, so a crafted state can pair the type with an arity + # that no longer matches its field count. `FromPicklable` must reject it at load, rather than + # build a corrupt spec (the repr guards the post-load mutation case separately). + Point = namedtuple('Point', ('x', 'y')) # noqa: PYI024 + state = optree.tree_structure(Point(1, 2)).__getstate__() # arity 2 + Point._fields = ('x', 'y', 'z') # diverge: 3 fields vs the pickled arity of 2 + + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + with pytest.raises(RuntimeError, match=r'does not match the arity'): + obj.__setstate__(state) + + @disable_systrace def test_treespec_self_referential(): class Holder: @@ -344,6 +414,174 @@ def __repr__(self): assert wr() is None +@skipif_pypy # relies on CPython's reference-cycle collector +@skipif_deferred_type_refs +def test_treespec_custom_node_reference_cycle_is_collectable(): + # A treespec reaches a registered custom type through the shared registration held by its custom + # node. Once the registry no longer pins the registration, the node is its sole owner, so + # `PyTreeSpec::PyTpTraverse` reports those objects and the cyclic GC can collect a cycle through + # them. The cycle keeps the heap type alive, and free-threaded builds before 3.14 hold deferred + # references to type objects in per-thread caches, so `gc_collect()` cannot reclaim it there. + # 3.14t does, so this keeps free-threaded coverage. + class Cyclic: + pass + + optree.register_pytree_node( + Cyclic, + lambda cyclic: ((), None), + lambda metadata, children: None, # never called; the test only needs the registration + namespace='cycle_gc', + ) + try: + treespec = optree.tree_structure(Cyclic(), namespace='cycle_gc') + # Cycle: Cyclic -> __dict__ -> treespec -> registration -> Cyclic + Cyclic.self_spec = treespec + finally: + optree.unregister_pytree_node(Cyclic, namespace='cycle_gc') + + wr = weakref.ref(Cyclic) + del Cyclic, treespec + gc_collect() + assert wr() is None + + +@skipif_pypy # relies on CPython's reference-cycle collector +@skipif_deferred_type_refs +def test_treespec_custom_node_reference_cycle_is_collectable_with_repeated_nodes(): + # Regression: the traverse reported a registration's members only when a single node held it, + # so a treespec containing the same registered type more than once never reported them and the + # cycle leaked. The treespec collectively owns the registration in that case too: what matters + # is that no one outside it holds a reference. + for num_nodes in (1, 2, 5): + + class Cyclic: + pass + + optree.register_pytree_node( + Cyclic, + lambda cyclic: ((), None), + lambda metadata, children: None, + namespace='cycle_gc_repeated', + ) + try: + tree = [Cyclic() for _ in range(num_nodes)] + treespec = optree.tree_structure(tree, namespace='cycle_gc_repeated') + Cyclic.self_spec = treespec + finally: + optree.unregister_pytree_node(Cyclic, namespace='cycle_gc_repeated') + + wr = weakref.ref(Cyclic) + del Cyclic, treespec, tree + gc_collect() + assert wr() is None, num_nodes + + +@skipif_pypy # relies on CPython's reference-cycle collector +def test_treespec_shared_registration_refs_are_not_reported(): + # `PyTreeSpec::PyTpTraverse` must report only references the treespec owns. + # A shared registration holds one reference to each member however many nodes point at it, so + # reporting per node would underflow the object's shadow refcount and abort on debug builds. + # `gc.get_referrers()` walks `tp_traverse`, so it shows what the traversal reports. + class Shared: + pass + + optree.register_pytree_node( + Shared, + lambda shared: ((), None), + lambda metadata, children: None, + namespace='shared_gc', + ) + try: + # The registry holds the registration, so no treespec is its sole owner. + treespecs = [optree.tree_structure(Shared(), namespace='shared_gc') for _ in range(4)] + gc_collect() + assert not any(treespec in gc.get_referrers(Shared) for treespec in treespecs) + # The registration is also shared between treespecs, so dropping the registry's hold while + # more than one treespec remains must not make them report it either. + optree.unregister_pytree_node(Shared, namespace='shared_gc') + gc_collect() + assert not any(treespec in gc.get_referrers(Shared) for treespec in treespecs) + finally: + del treespecs + + wr = weakref.ref(Shared) + del Shared + gc_collect() + if not HAS_DEFERRED_TYPE_REFS: + assert wr() is None + + +@skipif_pypy # relies on CPython's reference-cycle collector +def test_treespec_shared_registration_is_still_not_reported_with_repeated_nodes(): + # The counterpart: while anything outside the treespec holds the registration, its members must + # not be reported however many nodes reference it, or the collector's shadow refcount underflows. + class Shared: + pass + + optree.register_pytree_node( + Shared, + lambda shared: ((), None), + lambda metadata, children: None, + namespace='shared_gc_repeated', + ) + treespec = optree.tree_structure([Shared(), Shared()], namespace='shared_gc_repeated') + other = optree.tree_structure(Shared(), namespace='shared_gc_repeated') + gc_collect() + # The registry still holds it. + assert treespec not in gc.get_referrers(Shared) + optree.unregister_pytree_node(Shared, namespace='shared_gc_repeated') + gc_collect() + # `other` still holds it. + assert treespec not in gc.get_referrers(Shared) + del other + gc_collect() + # Now the treespec's two nodes are the only holders, so it reports the members once. + assert treespec in gc.get_referrers(Shared) + del treespec + + wr = weakref.ref(Shared) + del Shared + gc_collect() + if not HAS_DEFERRED_TYPE_REFS: + assert wr() is None + + +@skipif_pypy # relies on CPython's reference-cycle collector +@pytest.mark.xfail( + strict=True, + reason='known limitation: a treespec cannot see registration references held by another treespec', +) +def test_treespec_reference_cycle_across_treespecs_is_collectable(): + # Known limitation. A treespec reports a registration's members only when its own nodes hold + # every reference to it, because that is all it can count. When two treespecs each hold some of + # the references, neither sees the other's, so neither reports the members and a cycle through + # them survives even though the two treespecs jointly own the registration. + # + # Resolving this needs the registration itself to be a garbage-collected object with its own + # `tp_traverse`, so each edge is reported by whoever owns it and no counting is required. + # Until then this is strictly better than reporting nothing at all, which never collects any of + # these cycles. + class Cyclic: + pass + + optree.register_pytree_node( + Cyclic, + lambda cyclic: ((), None), + lambda metadata, children: None, + namespace='cycle_gc_across', + ) + try: + treespecs = [optree.tree_structure(Cyclic(), namespace='cycle_gc_across') for _ in range(2)] + Cyclic.self_specs = treespecs + finally: + optree.unregister_pytree_node(Cyclic, namespace='cycle_gc_across') + + wr = weakref.ref(Cyclic) + del Cyclic, treespecs + gc_collect() + assert wr() is None + + @disable_systrace def test_treeiter_self_referential(): sentinel = object() @@ -375,6 +613,24 @@ def test_treeiter_self_referential(): assert wr() is None +def test_treeiter_leaf_predicate_no_reference_leak(): + # A reference cycle that runs through the `leaf_predicate` callback must be collectable. + # Regression: `PyTreeIter` tp_traverse / tp_clear previously ignored `m_leaf_predicate`, so a + # cycle through the predicate was invisible to the cyclic garbage collector and leaked. + def is_leaf(x): + return False + + it = optree.tree_iter({'a': 1, 'b': {'c': 2}}, is_leaf) + wr = weakref.ref(it) + assert next(it) == 1 + is_leaf.self_ref = it # cycle: it -> m_leaf_predicate (is_leaf) -> is_leaf.self_ref -> it + + del it, is_leaf + gc_collect() + if not PYPY: + assert wr() is None + + def test_treespec_with_namespace(): tree = NAMESPACED_TREE @@ -541,6 +797,29 @@ def test_treespec_pickle_roundtrip( ) +@skipif_wasm +@skipif_android +@skipif_ios +def test_treespec_pickle_all_protocols_roundtrip(): + # pybind11's pickle support reconstructs cleanly only at protocol >= 2. Protocols 0 and 1 used + # to reconstruct via `object.__new__`, which pybind11 rejects with an untranslated C++ exception + # that aborts the interpreter (SIGABRT). Run in a subprocess so a regression fails this test + # rather than killing the whole suite. + check_script_in_subprocess( + r""" + import pickle + + import optree + + spec = optree.tree_structure({'a': [1, 2], 'b': (3, 4)}) + for protocol in range(pickle.HIGHEST_PROTOCOL + 1): + restored = pickle.loads(pickle.dumps(spec, protocol=protocol)) + assert restored == spec, (protocol, restored, spec) + """, + output=None, + ) + + class Foo: def __init__(self, x, y): self.x = x @@ -592,6 +871,402 @@ def test_treespec_pickle_missing_registration(): treespec = pickle.loads(serialized) +def test_treespec_getstate_does_not_alias_internal_node_data(): + # `__getstate__` (used by `pickle`) must return a snapshot, not aliases of the immutable spec's + # internal mutable containers: the keys of a dict/OrderedDict/defaultdict node and its + # insertion-order keys dict. Mutating the returned state otherwise reaches back into the spec, + # desyncing the keys from the arity (repr raises an InternalError) or adding a spurious + # original key (unflatten returns an extra entry). A custom node's entries are immutable. + class Custom: + def __init__(self, *values): + self.values = values + + optree.register_pytree_node( + Custom, + lambda custom: (custom.values, None, tuple(range(len(custom.values)))), + lambda metadata, children: Custom(*children), + namespace='getstate_snapshot', + ) + try: + tree = { + 'b': Custom(1, 2), + 'a': 3, + 'od': OrderedDict([('y', 4), ('x', 5)]), + 'dd': defaultdict(int, {'q': 6, 'p': 7}), + } + spec = optree.tree_structure(tree, namespace='getstate_snapshot') + node_states, _, _ = state = spec.__getstate__() + before = repr(state) + + for node in node_states: + kind, node_data, node_entries, original_keys = node[0], node[2], node[3], node[7] + if kind in {optree.PyTreeKind.DICT, optree.PyTreeKind.ORDEREDDICT}: + node_data.append('injected') # a dict/OrderedDict node's keys list + elif kind == optree.PyTreeKind.DEFAULTDICT: + node_data[1].append('injected') # a defaultdict's (default_factory, keys) tuple + if isinstance(original_keys, dict): + original_keys['injected'] = None + assert node_entries is None or isinstance(node_entries, tuple) # entries are immutable + + assert repr(spec.__getstate__()) == before, 'mutating the pickled state corrupted the spec' + finally: + optree.unregister_pytree_node(Custom, namespace='getstate_snapshot') + + +def test_treespec_getstate_aliases_custom_node_data(): + # Limitation (characterization test): a custom node's `node_data` is the user-provided metadata, + # which `__getstate__` passes through by reference. optree copies its own dict/defaultdict keys + # (see `test_treespec_getstate_does_not_alias_internal_node_data`) but cannot generically + # deep-copy arbitrary metadata, so mutating it via the pickled state reaches back into the spec. + # Protecting custom metadata is the caller's responsibility; this pins the behavior. + class Custom: + def __init__(self, *children, alpha, beta=None): + self.children = children + self.metadata = {'alpha': alpha, 'beta': beta} + + def __eq__(self, other): + return ( + isinstance(other, Custom) + and self.children == other.children + and self.metadata == other.metadata + ) + + __hash__ = None + + optree.register_pytree_node( + Custom, + lambda custom: (custom.children, custom.metadata), # mutable dict metadata + lambda metadata, children: Custom(*children, **metadata), + namespace='getstate_alias_custom', + ) + try: + leaves, treespec = optree.tree_flatten( + Custom(1, 2, alpha=3, beta=4), + namespace='getstate_alias_custom', + ) + before = repr(treespec) + custom_state = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert custom_state[2] == {'alpha': 3, 'beta': 4} + + # Mutate the aliased metadata in place via the pickled state. + custom_state[2]['gamma'] = 5 # mutate the aliased metadata in place + + aliased = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert aliased[2] is custom_state[2] + assert aliased[2] == {'alpha': 3, 'beta': 4, 'gamma': 5} # the mutation reached the spec + assert repr(treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 3, 'beta': 4, 'gamma': 5}), + ) + # The corruption even reaches what `tree_unflatten` rebuilds, not just repr/getstate. + with pytest.raises(TypeError, match=r'unexpected keyword argument'): + optree.tree_unflatten(treespec, leaves) + + # Replacing the metadata wholesale reaches the spec the same way, but here unflatten + # succeeds and rebuilds a different object: the corruption is silent, not an error. + custom_state[2].clear() + custom_state[2]['alpha'] = 42 + aliased = next( + node for node in treespec.__getstate__()[0] if node[0] == optree.PyTreeKind.CUSTOM + ) + assert aliased[2] is custom_state[2] + assert aliased[2] == {'alpha': 42} + assert repr(treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 42}), + ) + reconstructed = optree.tree_unflatten(treespec, leaves) + reconstructed_treespec = optree.tree_structure( + reconstructed, + namespace='getstate_alias_custom', + ) + assert reconstructed == Custom(1, 2, alpha=42) + assert reconstructed.metadata == {'alpha': 42, 'beta': None} + assert reconstructed_treespec != treespec + assert repr(reconstructed_treespec) == before.replace( + repr({'alpha': 3, 'beta': 4}), + repr({'alpha': 42, 'beta': None}), + ) + finally: + optree.unregister_pytree_node(Custom, namespace='getstate_alias_custom') + + +def test_treespec_setstate_does_not_alias_supplied_node_data(): + # The symmetric half of `test_treespec_getstate_does_not_alias_internal_node_data`: + # `__setstate__` validated the supplied key list and then BORROWED it into the node, so + # mutating the state afterwards silently corrupted an already-restored treespec (the keys + # desync from the children, and `unflatten` pairs them up wrongly). It must copy instead. + # `original_keys` is already rebuilt via `dict.fromkeys`, so it is covered too. + def setstate(state): + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + obj.__setstate__(state) + return obj + + trees = [ + {'a': 0, 'b': 0}, + OrderedDict([('a', 0), ('b', 0)]), + defaultdict(int, {'a': 0, 'b': 0}), + {'b': 0, 'a': 0}, # sorted keys differ from the insertion order recorded in original_keys + ] + for tree in trees: + spec = optree.tree_structure(tree) + state = spec.__getstate__() + restored = setstate(state) + before = repr(restored) + expected_entries = restored.entries() + expected_tree = restored.unflatten([10, 20]) + + for node in state[0]: + kind, node_data, original_keys = node[0], node[2], node[7] + if kind in {optree.PyTreeKind.DICT, optree.PyTreeKind.ORDEREDDICT}: + node_data.reverse() # a dict/OrderedDict node's keys list + node_data.append('injected') + elif kind == optree.PyTreeKind.DEFAULTDICT: + node_data[1].reverse() # a defaultdict's (default_factory, keys) tuple + node_data[1].append('injected') + if isinstance(original_keys, dict): + original_keys['injected'] = None + + assert repr(restored) == before, tree + assert restored.entries() == expected_entries, tree + assert restored.unflatten([10, 20]) == expected_tree, tree + assert restored == spec, tree + + +def test_treespec_setstate_rejects_malformed_state(): + # `PyTreeSpec.__setstate__` (used by `pickle`) must reject structurally malformed state rather + # than build a corrupt spec that triggers out-of-bounds reads / crashes when later used. The + # per-node tuple layout is (kind, arity, node_data, node_entries, custom, num_leaves, num_nodes, + # original_keys); see `PyTreeSpec::FromPicklable`. + def setstate(state): + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + obj.__setstate__(state) + return obj + + CUSTOM = int(optree.PyTreeKind.CUSTOM) # noqa: N806 + LEAF = int(optree.PyTreeKind.LEAF) # noqa: N806 + NONE = int(optree.PyTreeKind.NONE) # noqa: N806 + TUPLE = int(optree.PyTreeKind.TUPLE) # noqa: N806 + DICT = int(optree.PyTreeKind.DICT) # noqa: N806 + NAMEDTUPLE = int(optree.PyTreeKind.NAMEDTUPLE) # noqa: N806 + DEFAULTDICT = int(optree.PyTreeKind.DEFAULTDICT) # noqa: N806 + DEQUE = int(optree.PyTreeKind.DEQUE) # noqa: N806 + STRUCTSEQUENCE = int(optree.PyTreeKind.STRUCTSEQUENCE) # noqa: N806 + NUM_KINDS = int(optree.PyTreeKind.NUM_KINDS) # noqa: N806 + leaf_node = (LEAF, 0, None, None, None, 1, 1, None) # arity 0, 1 leaf, 1 node + keys_ab = {'a': None, 'b': None} # original_keys for a 2-key ('a', 'b') dict node + + # Sanity: well-formed states still round-trip. + for spec in [ + optree.tree_structure((0, 0)), + optree.tree_structure({'a': 0, 'b': 0}), + optree.tree_structure(defaultdict(int, {'a': 0, 'b': 0})), + ]: + assert setstate(spec.__getstate__()) == spec + + malformed_exceptions = (RuntimeError, ValueError, TypeError) + + # The rejection cases below follow the order of the checks in `PyTreeSpec::FromPicklable`. + + # A state that is not a 3-tuple. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node,), False)) + + # A node state that is not a 7- or 8-tuple. + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1),), False, '')) + + # Kind out of range: the raw integer is validated before the narrowing `uint8_t` enum cast, + # which would otherwise wrap a bogus value to a valid-looking kind. + with pytest.raises(malformed_exceptions): + setstate((((NUM_KINDS, 0, None, None, None, 0, 1, None),), False, '')) + + # Negative arity. + with pytest.raises(malformed_exceptions): + setstate((((TUPLE, -1, None, None, None, 0, 1, None),), False, '')) + + # A dict node missing its original keys, and a non-dict node carrying them. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (DICT, 2, ['a', 'b'], None, None, 2, 3, None)), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1, 1, keys_ab),), False, '')) + + # A negative leaf count, or a non-positive node count. + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, -1, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, None, None, None, 1, 0, None),), False, '')) + + # Node data on a leaf or none node (childless kinds that must not carry any). + with pytest.raises(malformed_exceptions): + setstate((((LEAF, 0, 'data', None, None, 1, 1, None),), False, '')) + + # Leaf or none nodes are childless; a nonzero arity absorbs the preceding subtrees while still + # folding consistently, so the reconstructed spec reports a leaf/None while its num_leaves counts + # the absorbed children and unflatten silently drops them. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, (NONE, 1, None, None, None, 1, 2, None)), False, '')) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, (LEAF, 1, None, None, None, 1, 2, None)), False, '')) + + # A None-kind node cannot appear when none_is_leaf is set (None is flattened as a leaf then, so + # a flattened tree never contains a None node); accepting one later raises an InternalError. + with pytest.raises(malformed_exceptions): + setstate((((NONE, 0, None, None, None, 0, 1, None),), True, '')) + + # Node data on a tuple or list node. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (TUPLE, 2, 'data', None, None, 2, 3, None)), False, '')) + + # Dict key list shorter than arity (MakeNode would index past the list end). + short_keys = (DICT, 2, ['a'], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, short_keys), False, '')) + + # Dict with duplicate keys (would collapse the rebuilt dict), and with an unhashable key. + dup_keys = (DICT, 2, ['a', 'a'], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, dup_keys), False, '')) + unhashable_key = (DICT, 2, [[], []], None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, unhashable_key), False, '')) + + # NamedTuple / StructSequence node_data that is not the expected kind of type. + with pytest.raises(malformed_exceptions): + setstate((((NAMEDTUPLE, 0, int, None, None, 0, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate((((STRUCTSEQUENCE, 0, int, None, None, 0, 1, None),), False, '')) + + # DefaultDict metadata as a list where a 2-tuple is expected previously caused a raw tuple-item + # read to segfault; it is now coerced to a tuple and used safely. + restored = setstate( + ( + ( + leaf_node, + leaf_node, + (DEFAULTDICT, 2, [int, ['a', 'b']], None, None, 2, 3, keys_ab), + ), + False, + '', + ), + ) + assert optree.tree_unflatten(restored, [10, 20]) == defaultdict(int, {'a': 10, 'b': 20}) + + # DefaultDict metadata with the wrong tuple size is rejected. + wrong_metadata = (DEFAULTDICT, 2, (int, ['a', 'b'], 'extra'), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, wrong_metadata), False, '')) + + # DefaultDict default_factory that is neither None nor callable. + bad_factory = (DEFAULTDICT, 2, (42, ['a', 'b']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, bad_factory), False, '')) + + # DefaultDict keys too few, and DefaultDict keys not distinct (the Dict variants are above). + defaultdict_short = (DEFAULTDICT, 2, (int, ['a']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, defaultdict_short), False, '')) + defaultdict_dup = (DEFAULTDICT, 2, (int, ['a', 'a']), None, None, 2, 3, keys_ab) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, defaultdict_dup), False, '')) + + # Deque maxlen that is neither None nor an int, and maxlen smaller than the arity (a deque holds + # at most maxlen items, so arity <= maxlen). + with pytest.raises(malformed_exceptions): + setstate((((DEQUE, 0, 'x', None, None, 0, 1, None),), False, '')) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, (DEQUE, 2, 1, None, None, 2, 3, None)), False, '')) + + # A non-custom node carrying node entries or a custom type. + with pytest.raises(malformed_exceptions): + setstate( + ((leaf_node, leaf_node, (TUPLE, 2, None, ('a', 'b'), None, 2, 3, None)), False, ''), + ) + + # Original keys whose count (not just key set) disagrees with the arity. + short_original = (DICT, 2, ['a', 'b'], None, None, 2, 3, {'a': None}) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, short_original), False, '')) + + # Dict original_keys whose key set differs from the sorted key list. + mismatched_original = (DICT, 2, ['a', 'b'], None, None, 2, 3, {'a': None, 'c': None}) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, mismatched_original), False, '')) + + # A custom node whose node-entries count disagrees with the arity (needs a registered type). + class MalformedCustomNode: + pass + + optree.register_pytree_node( + MalformedCustomNode, + lambda obj: ((), None), + lambda metadata, children: MalformedCustomNode(), + namespace='malformed', + ) + try: + custom_node = (CUSTOM, 2, None, ('one-entry',), MalformedCustomNode, 2, 3, None) + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node, custom_node), False, 'malformed')) + finally: + optree.unregister_pytree_node(MalformedCustomNode, namespace='malformed') + + # A node claiming more children than the traversal provides. + with pytest.raises(malformed_exceptions): + setstate((((TUPLE, 2, None, None, None, 2, 3, None),), False, '')) + + # Inconsistent intermediate num_nodes (previously only the last node was checked). + with pytest.raises(malformed_exceptions): + setstate( + ( + ( + (LEAF, 0, None, None, None, 1, 5, None), # leaf claims num_nodes == 5 + leaf_node, + (TUPLE, 2, None, None, None, 2, 3, None), + ), + False, + '', + ), + ) + + # A traversal that yields more than one tree. + with pytest.raises(malformed_exceptions): + setstate(((leaf_node, leaf_node), False, '')) + + +def test_treespec_setstate_rejects_builtin_custom_type(): + # Regression: `FromPicklable` accepted any registration found for a CUSTOM node's custom type. + # The built-in registrations (NoneType/tuple/list/dict/...) live in the same map but carry empty + # flatten/unflatten callables, so the reconstructed node later called a null function pointer + # and crashed the interpreter. Both `__setstate__` and `pickle.loads` must reject it. + CUSTOM = int(optree.PyTreeKind.CUSTOM) # noqa: N806 + LEAF = int(optree.PyTreeKind.LEAF) # noqa: N806 + leaf_node = (LEAF, 0, None, None, None, 1, 1, None) + + for builtin_type in (list, dict, tuple, deque, OrderedDict, defaultdict, type(None)): + state = ((leaf_node, (CUSTOM, 1, None, None, builtin_type, 1, 2, None)), False, '') + obj = optree.PyTreeSpec.__new__(optree.PyTreeSpec) + with pytest.raises(RuntimeError, match=r'the custom type is a built-in type'): + obj.__setstate__(state) + + # The same state shipped as a pickle payload, hand-assembled the way an attacker would: + # `NEWOBJ` an empty spec, then `BUILD` it from the crafted state. + blob = b''.join( + [ + pickle.PROTO + bytes([2]), + pickle.GLOBAL + b'optree\nPyTreeSpec\n', + pickle.EMPTY_TUPLE + pickle.NEWOBJ, + pickle.dumps(state, protocol=2)[2:-1], # strip the PROTO / STOP framing + pickle.BUILD + pickle.STOP, + ], + ) + with pytest.raises(RuntimeError, match=r'the custom type is a built-in type'): + pickle.loads(blob) # crafted payload, the point of the test + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -744,6 +1419,535 @@ def test_treespec_compose_children( assert optree.treespec_is_suffix(expected_treespec, treespec, strict=False) +def test_treespec_compose_rejects_incompatible_namespace_merge(): + # Regression: composing an empty-namespace spec (whose custom nodes are resolved globally) with + # a spec in another namespace adopted that namespace but kept the global registrations. When the + # same type is registered differently in the two namespaces, the composed spec silently used the + # wrong flatten/unflatten (spurious flatten_up_to errors; corrupt pickle). Reject the merge. + class Pair: + def __init__(self, a, b): + self.a, self.b = a, b + + class Single: + def __init__(self, x): + self.x = x + + optree.register_pytree_node( + Pair, + lambda t: ((t.a, t.b), None, None), + lambda m, c: Pair(c[0], c[1]), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( # behavior differs from the global registration + Pair, + lambda t: ((t.b, t.a), None, None), + lambda m, c: Pair(c[1], c[0]), + namespace='behavior_change', + ) + optree.register_pytree_node( + Single, + lambda t: ((t.x,), None, None), + lambda m, c: Single(c[0]), + namespace='behavior_change', + ) + try: + outer = optree.tree_structure(Pair(0, 0)) + inner = optree.tree_structure(Single(0), namespace='behavior_change') + assert outer.namespace == '' + assert inner.namespace == 'behavior_change' + with pytest.raises(ValueError, match='original registration'): + outer.compose(inner) + + # `tree_transpose` builds its expected structure with `compose`, so the rejection surfaces + # through the public API too (here via its structure-mismatch diagnostic path). + with pytest.raises(ValueError, match='original registration'): + optree.tree_transpose(outer, inner, [1, 2, 3]) + + # `broadcast_to_common_suffix` adopts the namespace the same way. + self_spec = optree.tree_structure({'k': Pair(0, 0)}) + other_spec = optree.tree_structure({'k': Single(0)}, namespace='behavior_change') + with pytest.raises(ValueError, match='original registration'): + self_spec.broadcast_to_common_suffix(other_spec) + finally: + optree.unregister_pytree_node(Pair, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Pair, namespace='behavior_change') + optree.unregister_pytree_node(Single, namespace='behavior_change') + + +def test_treespec_compose_allows_compatible_namespace_merge(): + # The namespace-merge rejection must not over-reject: a custom type registered only globally + # resolves identically under any namespace (via global fallback), so merging an empty-namespace + # spec that uses it into another namespace is allowed and the result stays consistent. + class GlobalOnly: + def __init__(self, a, b): + self.a, self.b = a, b + + optree.register_pytree_node( + GlobalOnly, + lambda t: ((t.a, t.b), None, None), + lambda m, c: GlobalOnly(c[0], c[1]), + namespace=GLOBAL_NAMESPACE, + ) + try: + outer = optree.tree_structure(GlobalOnly(0, 0)) + assert outer.namespace == '' + + # Both empty -> the merge stays in the global namespace. + assert outer.compose(optree.tree_structure(0)).namespace == '' + + # Empty side (global-only custom) merged into a namespace: allowed, adopts the namespace, + # and unflattens consistently (the global registration is used throughout). + with optree.dict_insertion_ordered(True, namespace='no_override'): + inner = optree.tree_structure({'x': 0}, namespace='no_override') + assert inner.namespace == 'no_override' + composed = outer.compose(inner) + assert composed.namespace == 'no_override' + result = optree.tree_unflatten(composed, [1, 2]) + assert isinstance(result, GlobalOnly) + assert result.a == {'x': 1} + assert result.b == {'x': 2} + + # The cross-namespace merge equals building the composed structure directly with `tree_map` + # in the adopted namespace, compose's defining identity. + expected = optree.tree_structure( + optree.tree_map(lambda _: {'x': 0}, GlobalOnly(0, 0), namespace='no_override'), + namespace='no_override', + ) + assert composed == expected + + # broadcast_to_common_suffix likewise allows the compatible merge. + broadcasted = outer.broadcast_to_common_suffix( + optree.tree_structure(GlobalOnly(0, 0), namespace='no_override'), + ) + assert broadcasted.namespace == 'no_override' + finally: + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_broadcast_to_common_suffix_does_not_mutate_argument_on_key_mismatch(): + # Regression: BroadcastToCommonSuffixImpl built the "got key(s)" part of its key-mismatch error + # message by sorting the ARGUMENT spec's live dict-node key list IN PLACE: `other_keys` was a + # borrow of `node_data`, not a copy. For an OrderedDict the child subtrees stay in insertion + # order while the keys get permuted, silently corrupting a spec the caller still holds: repr, + # equality, hash, and unflatten all go wrong. The message must be built from a sorted COPY. + other = optree.tree_structure(OrderedDict([('c', 1), ('b', 2)])) + before_repr = str(other) + before_hash = hash(other) + this = optree.tree_structure({'a': 1}) + with pytest.raises(ValueError, match='dictionary key mismatch'): + this.broadcast_to_common_suffix(other) + # The argument spec must be byte-for-byte unchanged by the failed call. + assert str(other) == before_repr + assert hash(other) == before_hash + # And it must still unflatten in its ORIGINAL insertion order (c, b), not a sorted (b, c) order. + assert other.unflatten([10, 20]) == OrderedDict([('c', 10), ('b', 20)]) + + +def test_treespec_broadcast_to_common_suffix_preserves_custom_node_entries(): + # Broadcasting rebuilds each non-leaf node, so `.node_entries` must be carried across: losing + # it breaks accessors, which fall back to `range(arity)` (`GetAttrEntry(entry=0)`). + class Vector: + def __init__(self, a, c): + self.a, self.c = a, c + + optree.register_pytree_node( + Vector, + lambda o: ((o.a, o.c), None, ('a', 'c')), # 3-tuple flatten -> node_entries = ('a', 'c') + lambda metadata, children: Vector(*children), + path_entry_type=optree.GetAttrEntry, + namespace=GLOBAL_NAMESPACE, + ) + try: + spec = optree.tree_structure(Vector(1, 2)) + other = optree.tree_structure(Vector(3, 4)) + assert spec.entries() == ['a', 'c'] + + # Both specs share the same custom structure, so the common suffix is that structure and the + # explicit string entries must survive unchanged, not degrade to the fallback [0, 1]. + broadcasted = spec.broadcast_to_common_suffix(other) + assert broadcasted.entries() == ['a', 'c'] + assert broadcasted.paths() == spec.paths() + assert broadcasted.accessors() == spec.accessors() + finally: + optree.unregister_pytree_node(Vector, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_deep_walk_raises_recursion_error_not_segfault(): + # Regression: `PathsImpl`, `AccessorsImpl`, and `BroadcastToCommonSuffixImpl` recurse once per + # tree level. Without a depth guard, a deeply-nested spec (trivially built via doubling + # `compose`) overflowed the native C++ stack and crashed the interpreter with a SIGSEGV instead + # of raising a catchable `RecursionError`. + # Each `compose` doubles the depth, so ceil(log2(limit)) + 1 composes push it above the limit. + num_composes = math.ceil(math.log2(optree.MAX_RECURSION_DEPTH)) + 1 + deep = optree.tree_structure([0]) + for _ in range(num_composes): + deep = deep.compose(deep) + assert 2**num_composes > optree.MAX_RECURSION_DEPTH + with pytest.raises(RecursionError): + deep.paths() + with pytest.raises(RecursionError): + deep.accessors() + with pytest.raises(RecursionError): + deep.broadcast_to_common_suffix(deep) + + # Broadcasting the deep spec against a shallower spec whose depth is still below the limit + # recurses only as far as the common suffix, so it must succeed (not raise RecursionError or + # crash) in either direction, returning the deeper spec. + shallower = optree.tree_structure([0]) + for _ in range(num_composes - 2): # depth 2 ** (num_composes - 2), safely below the limit + shallower = shallower.compose(shallower) + assert 2 ** (num_composes - 2) < optree.MAX_RECURSION_DEPTH + assert deep.broadcast_to_common_suffix(shallower) == deep + assert shallower.broadcast_to_common_suffix(deep) == deep + + +def test_treespec_compose_rejects_namespace_override_with_different_arity(): + # A type registered globally flattens both members as children (arity 2); a namespace override + # flattens one member as a child and stores the other as node metadata (arity 1). Both + # registrations round-trip, but merging an empty-namespace spec (global, arity 2) into that + # namespace must be rejected: the composed spec would claim the namespace while carrying an + # arity-2 node that the namespace's registration cannot unflatten. + class TwoMember: + def __init__(self, a, b): + self.a, self.b = a, b + + def __eq__(self, other): + return isinstance(other, TwoMember) and (self.a, self.b) == (other.a, other.b) + + __hash__ = None + + optree.register_pytree_node( + TwoMember, + lambda t: ((t.a, t.b), None, None), # global: both members are children + lambda metadata, children: TwoMember(children[0], children[1]), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + TwoMember, + lambda t: ((t.a,), t.b, None), # override: one child, the other is metadata + lambda metadata, children: TwoMember(children[0], metadata), + namespace='arity_change', + ) + try: + obj = TwoMember(1, 2) + + # Both registrations round-trip on their own. + global_leaves, global_spec = optree.tree_flatten(obj) + assert global_leaves == [1, 2] + assert optree.tree_unflatten(global_spec, global_leaves) == obj + custom_leaves, custom_spec = optree.tree_flatten(obj, namespace='arity_change') + assert custom_leaves == [1] + assert optree.tree_unflatten(custom_spec, custom_leaves) == obj + + assert global_spec.namespace == '' + assert global_spec.num_leaves == 2 + assert custom_spec.namespace == 'arity_change' + assert custom_spec.num_leaves == 1 + with pytest.raises(ValueError, match='original registration'): + global_spec.compose(custom_spec) + finally: + optree.unregister_pytree_node(TwoMember, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(TwoMember, namespace='arity_change') + + +def test_treespec_transform_rejects_incompatible_namespace_merge(): + # `transform` unifies the namespace across the input spec and the transform outputs. If that + # unified (non-empty) namespace rebinds a custom node (e.g. the input's globally-resolved + # custom node) to a different registration, the transform must be rejected (same class as the + # compose / broadcast merge rejection). A globally-only-registered type is still allowed via + # fallback. + class Diverge: # variable arity; registered differently in the global and named namespaces + def __init__(self, *children): + self.children = children + + class GlobalOnly: # variable arity; registered only globally -> resolves via fallback anywhere + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Diverge, + lambda d: (d.children, None, None), + lambda metadata, children: Diverge(*children), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + Diverge, + lambda d: (tuple(reversed(d.children)), None, None), # divergent from the global reg + lambda metadata, children: Diverge(*reversed(children)), + namespace='transform_change', + ) + optree.register_pytree_node( + GlobalOnly, + lambda g: (g.children, None, None), + lambda metadata, children: GlobalOnly(*children), + namespace=GLOBAL_NAMESPACE, + ) + + def to_namespaced_leaf(_): + # Replace a leaf with a namespaced Diverge to inject the namespace (leaves have no arity). + return optree.tree_structure(Diverge(0), namespace='transform_change') + + def to_global_node(spec): + # Replace a node with a same-arity globally-resolved Diverge; it rebinds under the promoted + # namespace (Diverge is registered differently there). Generic over the node's arity. + return optree.tree_structure(Diverge(*range(spec.num_children))) + + def to_namespaced_node(spec): + # Outer node -> global Diverge (rebinds); inner nodes -> namespaced Diverge (injects the + # namespace). Both same-arity, so `f_node` alone drives the (f_node, None) rejection. + if spec.type is tuple: + return to_global_node(spec) + return optree.tree_structure( + Diverge(*range(spec.num_children)), + namespace='transform_change', + ) + + try: + # The rejection must fire for every `(f_node, f_leaf)` combination that puts a + # globally-resolved custom node under the non-empty unified namespace. + + # (None, f_leaf): the input's global Diverge is kept, f_leaf injects the namespace. + outer = optree.tree_structure(Diverge(0, 0)) + assert outer.namespace == '' + with pytest.raises(ValueError, match='original registration'): + outer.transform(None, to_namespaced_leaf) + + # (f_node, None): f_node alone yields a global Diverge above namespaced children. + with pytest.raises(ValueError, match='original registration'): + optree.tree_structure(([0], [0])).transform(to_namespaced_node, None) + + # (f_node, f_leaf): f_node injects the global Diverge, f_leaf injects the namespace. + with pytest.raises(ValueError, match='original registration'): + optree.tree_structure([0, 0]).transform(to_global_node, to_namespaced_leaf) + + # Compatible: GlobalOnly resolves identically under any namespace via fallback. + global_outer = optree.tree_structure(GlobalOnly(0, 0)) + transformed = global_outer.transform(None, to_namespaced_leaf) + assert transformed.namespace == 'transform_change' + finally: + optree.unregister_pytree_node(Diverge, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Diverge, namespace='transform_change') + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_from_collection_rejects_incompatible_namespace_promotion(): + # `treespec_from_collection` promotes an empty caller namespace to a child spec's namespace. If + # that promoted namespace rebinds a custom node the collection resolved globally (the root node, + # or a globally-resolved child) to a different registration, the result would claim the + # namespace while carrying the wrong registration: it must be rejected, exactly like compose / + # transform / broadcast. A globally-only-registered type is still allowed via fallback. + class Diverge: # variable arity; registered differently in the global and named namespaces + def __init__(self, *children): + self.children = children + + class GlobalOnly: # variable arity; registered only globally -> resolves via fallback anywhere + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Diverge, + lambda d: (d.children, None, None), + lambda metadata, children: Diverge(*children), + namespace=GLOBAL_NAMESPACE, + ) + optree.register_pytree_node( + Diverge, + lambda d: (tuple(reversed(d.children)), None, None), # divergent from the global reg + lambda metadata, children: Diverge(*reversed(children)), + namespace='from_coll_change', + ) + optree.register_pytree_node( + GlobalOnly, + lambda g: (g.children, None, None), + lambda metadata, children: GlobalOnly(*children), + namespace=GLOBAL_NAMESPACE, + ) + try: + # Incompatible: the globally-resolved Diverge rebinds under the promoted namespace. + foo = optree.tree_structure(Diverge(0, 0)) + child = optree.tree_structure(Diverge(0), namespace='from_coll_change') + assert foo.namespace == '' + with pytest.raises(ValueError, match='original registration'): + optree.treespec_from_collection([foo, child], namespace='') + + # Compatible: GlobalOnly resolves identically under any namespace via fallback. + global_spec = optree.tree_structure(GlobalOnly(0, 0)) + promoted = optree.treespec_from_collection([global_spec, child], namespace='') + assert promoted.namespace == 'from_coll_change' + finally: + optree.unregister_pytree_node(Diverge, namespace=GLOBAL_NAMESPACE) + optree.unregister_pytree_node(Diverge, namespace='from_coll_change') + optree.unregister_pytree_node(GlobalOnly, namespace=GLOBAL_NAMESPACE) + + +def test_treespec_dict_key_order_survives_namespace_promotion(): + # A dict node's key order is fixed at BUILD time by the namespace passed then. Operations that + # merge/promote a spec's namespace (`treespec_from_collection`, `compose`, `transform`) only + # re-tag it for custom-node resolution; like `compose` they NEVER reorder an already-built dict. + # So a dict built under the global ('') namespace (sorted keys) keeps that order even after + # promotion to an insertion-ordered namespace, intentionally differing from the same dict built + # directly under that namespace, while a dict built directly under the namespace keeps its + # insertion order (matching). This test locks that behavior across all three operations. + class Wrap: # a variable-arity custom node, so `f_node` can build same-arity replacements + def __init__(self, *children): + self.children = children + + optree.register_pytree_node( + Wrap, + lambda w: (w.children, None, None), + lambda metadata, children: Wrap(*children), + namespace='promote_order', + ) + try: + with optree.dict_insertion_ordered(True, namespace='promote_order'): + child = optree.tree_structure(Wrap(0), namespace='promote_order') + # Built directly under the insertion-ordered namespace: keys in insertion order (b, a). + genuine = optree.tree_structure({'b': Wrap(0), 'a': Wrap(0)}, namespace='promote_order') + assert genuine.entries() == ['b', 'a'] + + def to_namespaced_node(spec): + # Rewrite every non-dict node into a same-arity `Wrap` in the namespace (generic + # over the node's arity rather than tied to this test's shapes) so `f_node` alone + # can promote the spec. `transform` promotes only when some output carries a + # namespace, and only a custom node can. The outer dict node is kept so its key + # order stays observable. + if spec.type is dict: + return spec + return optree.tree_structure( + Wrap(*range(spec.num_children)), + namespace='promote_order', + ) + + def transform_combos(outer): + return { + 'transform(None, f_leaf)': outer.transform(None, lambda _: child), + 'transform(f_node, None)': outer.transform(to_namespaced_node, None), + 'transform(f_node, f_leaf)': outer.transform( + to_namespaced_node, + lambda _: child, + ), + } + + # Dicts built under the GLOBAL ('') namespace, sorted keys (a, b), then promoted. + from_global = { + 'from_collection': optree.treespec_from_collection( + {'b': child, 'a': child}, + namespace='', + ), + 'compose': optree.tree_structure({'b': 0, 'a': 0}).compose(child), + **transform_combos(optree.tree_structure({'b': [0], 'a': [0]})), + } + # Dicts built directly under the namespace, insertion-order keys (b, a). + from_namespace = { + 'from_collection': optree.treespec_from_collection( + {'b': child, 'a': child}, + namespace='promote_order', + ), + 'compose': optree.tree_structure( + {'b': 0, 'a': 0}, + namespace='promote_order', + ).compose(child), + **transform_combos( + optree.tree_structure({'b': [0], 'a': [0]}, namespace='promote_order'), + ), + } + + for name, spec in from_global.items(): + assert spec.namespace == 'promote_order', name # promoted for custom resolution ... + assert spec.entries() == ['a', 'b'], name # ... but the dict keeps '' (sorted) order + + for name, spec in from_namespace.items(): + assert spec.namespace == 'promote_order', name + assert spec.entries() == ['b', 'a'], name # insertion order kept, matches direct build + + # from_collection / compose reproduce genuine's flat structure exactly, so the only + # difference is the dict key order: global-built differs, namespace-built matches. + assert from_global['from_collection'] != genuine + assert from_global['compose'] != genuine + assert from_namespace['from_collection'] == genuine + assert from_namespace['compose'] == genuine + finally: + optree.unregister_pytree_node(Wrap, namespace='promote_order') + + +def test_treespec_is_prefix_nested_dict_key_reorder(): + # Regression: `IsPrefix` reorders a dict node's children in a working copy of the traversal to + # make key order irrelevant. When a NESTED dict also needed reordering, it indexed the pristine + # traversal by an offset into the already-mutated working copy, corrupting it -> a spurious + # `optree._C.InternalError` or a wrong boolean. Two treespecs that describe the SAME tree + # (differing only in dict key insertion order, at nested levels) must be mutual non-strict + # prefixes / suffixes. + + # Top-level AND nested dict keys reordered; the top-level reorder relocates the nested dict. + tree_a = OrderedDict([('a', 0), ('b', OrderedDict([('e', 0), ('g', 0)])), ('d', 0)]) + tree_b = OrderedDict([('b', OrderedDict([('g', 0), ('e', 0)])), ('a', 0), ('d', 0)]) + a = optree.tree_structure(tree_a) + b = optree.tree_structure(tree_b) + assert optree.treespec_is_prefix(a, b, strict=False) + assert optree.treespec_is_prefix(b, a, strict=False) + assert optree.treespec_is_suffix(a, b, strict=False) + assert optree.treespec_is_suffix(b, a, strict=False) + assert a <= b + assert b <= a + assert a >= b + assert b >= a + + # A nested dict whose reorder relocates a subtree containing another out-of-order dict. + tree_a2 = OrderedDict([('a', 0), ('d', 0), ('b', OrderedDict([('e', 0), ('f', 0)]))]) + tree_b2 = OrderedDict([('b', OrderedDict([('f', 0), ('e', 0)])), ('d', 0), ('a', 0)]) + a2 = optree.tree_structure(tree_a2) + b2 = optree.tree_structure(tree_b2) + assert optree.treespec_is_prefix(a2, b2, strict=False) + assert optree.treespec_is_prefix(b2, a2, strict=False) + assert a2 <= b2 + assert b2 <= a2 + + +def test_treespec_is_prefix_deque_maxlen_agnostic(): + # A deque's treespec stores both its arity and its `maxlen`, but `is_prefix` is arity-based and + # deliberately `maxlen`-AGNOSTIC. A deque holds at most `maxlen` items, so `arity <= maxlen` + # always holds and two flatten-compatible deques necessarily share the same arity while carrying + # any `maxlen1`/`maxlen2`; `maxlen` does not affect how children are partitioned, so gating the + # prefix relation on it would wrongly reject valid `flatten_up_to`/`broadcast_prefix` operations. + # `EqualTo`, by contrast, IS `maxlen`-sensitive (`unflatten` restores the exact `maxlen`), so + # `a <= b and b <= a` does NOT imply `a == b`: `is_prefix` is a preorder, not a partial order. + a = optree.tree_structure(deque([1, 2, 3], maxlen=3)) + b = optree.tree_structure(deque([1, 2, 3], maxlen=5)) + unbounded = optree.tree_structure(deque([1, 2, 3])) # maxlen=None + # Equality distinguishes maxlen (bounded vs bounded, and bounded vs unbounded). + assert a != b + assert a != unbounded + assert b != unbounded + + # Same arity, any maxlen (bounded or unbounded): mutual non-strict prefixes and suffixes, + # even though the specs are unequal, so mutual prefixes do NOT imply equality (a preorder). + for x, y in itertools.permutations([a, b, unbounded], 2): + assert x != y + assert optree.treespec_is_prefix(x, y, strict=False) + assert optree.treespec_is_suffix(x, y, strict=False) + assert x <= y + assert x >= y + + # Practical consequence: a prefix deque flattens / broadcasts a full deque of a different maxlen. + prefix_spec = optree.tree_structure(deque([1, 2, 3], maxlen=None)) + assert prefix_spec.flatten_up_to(deque([[10], [20, 21], [30]], maxlen=5)) == [ + [10], + [20, 21], + [30], + ] + assert optree.broadcast_prefix( + deque([1, 2, 3], maxlen=3), + deque([[0], [0, 0], [0]], maxlen=7), + ) == [1, 2, 2, 3] + + # Arity still gates the relation: a different-arity deque is not a prefix. + assert not optree.treespec_is_prefix( + optree.tree_structure(deque([1, 2], maxlen=9)), + a, + strict=False, + ) + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -1000,6 +2204,26 @@ def test_treespec_child( ] +def test_treespec_entry_and_child_accept_int_like_indices(): + # The compiled signatures advertise `SupportsInt | SupportsIndex`, and the runtime honors both, + # so the stubs must not narrow them to `int`. + class OnlyIndex: + def __index__(self): + return 1 + + class OnlyInt: + def __int__(self): + return 1 + + treespec = optree.tree_structure({'a': 1, 'b': 2}) + for index in (1, OnlyIndex(), OnlyInt()): + assert treespec.entry(index) == treespec.entry(1), index + assert treespec.child(index) == treespec.child(1), index + # The public wrappers must not narrow what the methods they forward to accept. + assert optree.treespec_entry(treespec, index) == treespec.entry(1), index + assert optree.treespec_child(treespec, index) == treespec.child(1), index + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -1420,6 +2644,36 @@ def test_treespec_leaf_none(namespace): ) +def test_treespec_from_collection_on_leaf_propagates_escalated_warning(): + # `treespec_from_collection()` on a leaf issues a UserWarning via `PyErr_WarnEx()`. When + # warnings are escalated to errors (e.g. `-W error`), the escalation must propagate cleanly as + # that UserWarning: the C++ code must check `PyErr_WarnEx()`'s return value and raise, not + # ignore it and return a result with the exception left set (which pybind11 surfaces as a + # confusing `SystemError: ... returned a result with an exception set`). + with warnings.catch_warnings(): + warnings.simplefilter('error') + with pytest.raises( + UserWarning, + match=re.escape('PyTreeSpec::MakeFromCollection() is called on a leaf.'), + ): + optree.treespec_from_collection(1) + + +def test_treespec_from_collection_drops_namespace_for_childless_roots(): + # Regression: a leaf or `None` root skipped the namespace-dropping step, so the caller's + # namespace stuck to a treespec that resolves no custom type. Two otherwise-identical treespecs + # built under different namespaces then compared unequal. + with warnings.catch_warnings(): + warnings.simplefilter('ignore', UserWarning) + for collection in (None, 1, 'leaf'): + first = optree.treespec_from_collection(collection, namespace='first') + second = optree.treespec_from_collection(collection, namespace='second') + assert first.namespace == '', (collection, first.namespace) + assert second.namespace == '', (collection, second.namespace) + assert first == second, collection + assert hash(first) == hash(second), collection + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -1784,6 +3038,30 @@ def __tree_unflatten__(cls, metadata, children): assert treespec1 == treespec2 +def test_treespec_dict_constructor_preserves_insertion_ordered_namespace(): + # Regression: under `dict_insertion_ordered` mode the key order of a dict spec depends on the + # namespace, so `treespec_dict(..., namespace=...)` must keep that namespace (like + # `tree_flatten`) instead of resetting it to '': an empty-namespace spec with unsorted keys is + # otherwise unreachable via `tree_flatten` and breaks equality/consistency. + leaf = optree.tree_structure(0) + + with optree.dict_insertion_ordered(True, namespace='namespace'): + constructed = optree.treespec_dict({'b': leaf, 'a': leaf}, namespace='namespace') + _, flattened = optree.tree_flatten({'b': 1, 'a': 2}, namespace='namespace') + + assert constructed.entries() == ['b', 'a'] # insertion order preserved + assert flattened.namespace == 'namespace' + assert constructed.namespace == 'namespace' # was '' before the fix + assert constructed == flattened + + # Without the mode, keys are sorted and the namespace is dropped, same as `tree_flatten`. + outside = optree.treespec_dict({'b': leaf, 'a': leaf}, namespace='namespace') + _, flattened_outside = optree.tree_flatten({'b': 1, 'a': 2}, namespace='namespace') + assert outside.entries() == ['a', 'b'] + assert outside.namespace == '' + assert outside == flattened_outside + + def test_treespec_constructor_none_treespec_inputs(): with pytest.raises(ValueError, match=r'Expected a\(n\) list of PyTreeSpec\(s\), got .*\.'): optree.treespec_list([optree.treespec_leaf(), 1]) diff --git a/tests/test_typing.py b/tests/test_typing.py index 51cb95cc..02ceee22 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -16,6 +16,7 @@ # pylint: disable=missing-function-docstring import enum +import os import re import sys import time @@ -28,14 +29,19 @@ import optree from helpers import ( PYBIND11_HAS_NATIVE_ENUM, + PYPY, CustomNamedTupleSubclass, CustomTuple, Py_GIL_DISABLED, Vector2D, + check_script_in_subprocess, disable_systrace, gc_collect, getrefcount, + skipif_android, + skipif_ios, skipif_pypy, + skipif_wasm, ) @@ -547,6 +553,27 @@ def test_structseq_fields(): 'tm_yday', 'tm_isdst', ) + # On CPython, `os.stat_result` has UNNAMED sequence slots 7, 8, 9 (the integer + # atime/mtime/ctime); the st_atime/st_mtime/st_ctime attributes are hidden FLOAT fields at + # higher field indices. `tp_members` must be mapped by offset, not by position, or those + # slots get mislabeled with the trailing hidden float names. + stat_fields = structseq_fields(os.stat_result) + assert len(stat_fields) == os.stat_result.n_sequence_fields + assert stat_fields[:7] == ( + 'st_mode', + 'st_ino', + 'st_dev', + 'st_nlink', + 'st_uid', + 'st_gid', + 'st_size', + ) + if not PYPY: + # PyPy has no unnamed fields: it names slots 7-9 `_integer_atime`/etc. and puts the + # hidden float `st_atime` at a later index, so this CPython-only check does not apply. + for name in stat_fields[7:10]: + assert name not in {'st_atime', 'st_mtime', 'st_ctime'} + assert not name.isidentifier() # the PyStructSequence unnamed-field marker with pytest.raises( TypeError, @@ -597,6 +624,36 @@ def test_structseq_fields(): structseq_fields(FakeStructSequence) +def test_structseq_accessor_unnamed_fields_codify_by_index(): + # The accessor round-trip (the generated code evaluates to the accessed value) must hold for + # every slot on every implementation. It exercises both codify styles: CPython leaves + # `os.stat_result` slots 7, 8, 9 UNNAMED, so their accessors codify to index access (matching the + # index-based `__call__`); PyPy names those slots (`_integer_atime` etc.) and codifies them by + # attribute. Either way `accessor.codify(...)` and `accessor(...)` resolve to the same `st[i]`. + st = os.stat(os.curdir) # a real stat_result, valid on both CPython and PyPy + accessors = optree.tree_accessors(st) + assert len(accessors) == os.stat_result.n_sequence_fields + for i, accessor in enumerate(accessors): + assert eval(accessor.codify('__st'), {'__st': st}, {}) == accessor(st) == st[i] + assert accessors[6].codify('__st') == '__st.st_size' # a named slot -> attribute access + + # Repeat with DISTINCT per-field values (`st[i] == i`) so the round-trip reliably catches the + # unnamed-slot mislabel: a real stat's whole-second atime could coincide with integer slot 7. On + # CPython slots 7, 8, 9 are UNNAMED, so their accessors must codify to index access; codifying slot + # 7 as `.st_atime` (the hidden FLOAT field CPython's own repr mislabels it with) would eval to + # that wrong value. PyPy names those slots (`_integer_atime` etc.) and aliases `st_atime` back to + # `self[7]`, so the unnamed-slot specifics below are asserted CPython-only. + st = os.stat_result(range(os.stat_result.n_fields)) + accessors = optree.tree_accessors(st) + assert len(accessors) == os.stat_result.n_sequence_fields + for i, accessor in enumerate(accessors): + assert eval(accessor.codify('__st'), {'__st': st}, {}) == accessor(st) == i + if not PYPY: + assert st.st_atime != st[7] # a different (hidden) field, not sequence slot 7 + for i in (7, 8, 9): + assert accessors[i].codify('__st') == f'__st[{i}]' + + @skipif_pypy @disable_systrace def test_structseq_fields_cache(): @@ -660,3 +717,123 @@ class Foo(metaclass=FooMeta): if not Py_GIL_DISABLED: assert called_with == 'Foo' assert wr() is None + + +@skipif_wasm +@skipif_android +@skipif_ios +@skipif_pypy # CPython-only: uses `atexit._ncallbacks()` and CPython type caches +def test_type_caches_register_interpreter_cleanup(): + # optree keeps three process-global type caches: namedtuple classification, PyStructSequence + # classification, and PyStructSequence field names. Each registers one per-interpreter `atexit` + # cleanup on its first insert (the classification caches at import via the registry, the + # field-name cache on first use). Measuring in a clean subprocess before importing optree pins + # optree's whole footprint: one callback for the registry plus one per cache. + check_script_in_subprocess( + r""" + import atexit + import time + + n0 = atexit._ncallbacks() + import optree + n1 = atexit._ncallbacks() + optree.is_namedtuple(int) + n2 = atexit._ncallbacks() + optree.is_structseq(int) + n3 = atexit._ncallbacks() + optree.structseq_fields(time.struct_time) + n4 = atexit._ncallbacks() + + assert n0 < n1, (n0, n1) + assert n1 <= n2, (n1, n2) + assert n2 <= n3, (n2, n3) + assert n3 <= n4, (n3, n4) + assert n4 - n0 == 4, (n0, n1, n2, n3, n4) + """, + output=None, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +@skipif_pypy # CPython-only: uses the CPython type caches +def test_type_cache_insert_failure_does_not_leave_a_dangling_entry(): + # Regression: the caches published an entry before taking a reference to the value and before + # creating the weakref that evicts it. If registering the per-interpreter `atexit` cleanup + # raised in between, the entry survived owning nothing and with no eviction hook, so the next + # lookup read the freed value and segfaulted. Run in a subprocess so a crash is a non-zero exit + # rather than a lost test session. + check_script_in_subprocess( + r""" + import atexit + import time + + import optree + + real_register = atexit.register + + def failing_register(*args, **kwargs): + raise RuntimeError('injected atexit failure') + + atexit.register = failing_register + try: + optree.structseq_fields(time.struct_time) + except RuntimeError: + pass + else: + raise AssertionError('the injected failure did not propagate') + finally: + atexit.register = real_register + + # The interpreter must not be marked as cleaned-up-registered by the failed attempt, or the + # cleanup would never be retried. Sample before the retry, which is the call that registers. + before = atexit._ncallbacks() + + # The failed insert must not be observable: the value is recomputed, not read back from a + # dangling entry. + fields = optree.structseq_fields(time.struct_time) + after = atexit._ncallbacks() + assert fields[:2] == ('tm_year', 'tm_mon'), fields + assert after == before + 1, (before, after) + """, + output=None, + ) + + +@skipif_wasm +@skipif_android +@skipif_ios +@skipif_pypy # CPython-only: uses the CPython type caches +def test_type_cache_insert_failure_before_import_does_not_crash(): + # The same hazard on the import-time path: the registry and the classification caches take their + # first entries while `optree` is being imported, so break `atexit.register` before the import + # rather than after. The initialization must fail as a normal `ImportError` and leave nothing + # half-registered behind, rather than caching an entry it does not own and crashing later. + # Re-importing in the same process is not possible once initialization has failed part way + # through, which is a pybind11 module-init limitation rather than something optree controls. + check_script_in_subprocess( + r""" + import atexit + import sys + + real_register = atexit.register + + def failing_register(*args, **kwargs): + raise RuntimeError('injected atexit failure') + + atexit.register = failing_register + try: + import optree + except ImportError: + pass + else: + raise AssertionError('the injected failure did not propagate') + finally: + atexit.register = real_register + + assert 'optree' not in sys.modules + assert 'optree._C' not in sys.modules + """, + output=None, + ) diff --git a/tests/test_utils.py b/tests/test_utils.py index f1abdb04..e2b38919 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -73,6 +73,37 @@ def __hash__(self): ] +def test_total_order_sorted_key_callback(): + # Regression: the `key` callback was passed to `sorted()` inside the `try` block, so a + # `TypeError` raised by the callback itself was mistaken for a comparison failure. The callback + # ran again in the fallback and the sequence came back unsorted with the error swallowed. + calls = [] + + def boom(x): + calls.append(x) + raise TypeError('callback boom') + + with pytest.raises(TypeError, match='callback boom'): + total_order_sorted([3, 1, 2], key=boom) + assert calls == [3] # the callback is not retried + + # Only comparison failures fall back, and the callback runs exactly once per element in every + # branch: the direct sort, the type-qualified sort, and the original-order fallback. + class NonSortable: + def __init__(self, x): + self.x = x + + for sequence, expected in ( + ([3, 1, 2], [1, 2, 3]), # sorts directly + ([3, '1', 2], [2, 3, '1']), # falls back to the type-qualified key + ([3, NonSortable(1), NonSortable(2)], None), # keeps the original order + ): + calls.clear() + result = total_order_sorted(sequence, key=lambda x: (calls.append(x), x)[1]) + assert calls == sequence + assert result == (expected if expected is not None else sequence) + + def test_safe_zip(): assert list(safe_zip([])) == [] assert list(safe_zip([1])) == [(1,)]