From 70fdc109572084bba83d620335935d06c38e4829 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 27 Jul 2026 16:56:50 +0800 Subject: [PATCH 1/6] fix(registry): correct GIL and registry lock ordering Several registry entry points released the GIL while holding `sm_mutex`, which inverts the GIL to `sm_mutex` order against a concurrent flatten that holds the GIL and waits on the read lock, wedging the interpreter. - `Register` and `Unregister` ran the namedtuple / PyStructSequence classification and `PyErr_WarnEx` under the write lock. Both run Python and can release the GIL. Classify and warn before taking the lock; emitting the override warning before any insert also keeps registration atomic under warnings-as-errors. - `Lookup`, `Register`, `Unregister` and `GetRegistrySize` called `GetSingleton()` under the lock. Under `per_interpreter_gil` that releases the GIL once any subinterpreter has imported the module. Acquire the singletons first. - `FlattenInto` read-locked `sm_dict_order_mutex` and then called `IsDictInsertionOrdered`, which locks it again. `std::shared_mutex` is not recursive, so the nested shared acquisition is undefined behaviour and deadlocks under writer preference. Add `GetDictInsertionOrderedFlags`, which returns both flags from one lock, and make `IsDictInsertionOrdered` delegate to it. - `GetRegistrySize` took `Size()` twice, dropping the lock between the two reads, so a concurrent registration could slip in and trip the consistency check. Split out `SizeImpl` and read both registries under a single lock. `RegisterImpl` and `UnregisterImpl` no longer need to be templated on `NoneIsLeaf` now that the callers hold both singletons, so they become plain member functions. --- CHANGELOG.md | 3 + include/optree/registry.h | 57 +++++++-- src/registry.cpp | 225 ++++++++++++++++++--------------- src/treespec/flatten.cpp | 32 ++--- tests/test_registry.py | 255 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 440 insertions(+), 132 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 76b6bb6b..62156358 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,9 @@ 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 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). ### Removed diff --git a/include/optree/registry.h b/include/optree/registry.h index 5d2be589..2c755c7c 100644 --- a/include/optree/registry.h +++ b/include/optree/registry.h @@ -133,8 +133,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 +169,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 +222,19 @@ 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); + 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); - template - [[nodiscard]] static RegistrationPtr UnregisterImpl(const py::object &cls, - const std::string ®istry_namespace); + [[nodiscard]] RegistrationPtr UnregisterImpl(const py::object &cls, + const std::string ®istry_namespace, + const bool &is_structsequence_class, + const bool &is_namedtuple_class); + + // 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/src/registry.cpp b/src/registry.cpp index 5c12296c..932171de 100644 --- a/src/registry.cpp +++ b/src/registry.cpp @@ -69,9 +69,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,15 +80,17 @@ 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]] { +ssize_t PyTreeTypeRegistry::Size(const std::optional ®istry_namespace) const { + const scoped_read_lock lock{sm_mutex}; + return SizeImpl(registry_namespace); +} + +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) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { throw py::value_error("PyTree type " + PyRepr(cls) + " is a built-in type and cannot be re-registered."); } @@ -101,29 +102,12 @@ template registration->unflatten_func = py::reinterpret_borrow(unflatten_func); registration->path_entry_type = py::reinterpret_borrow(path_entry_type); if (registry_namespace.empty()) [[unlikely]] { - if (!registry.m_registrations.emplace(cls, std::move(registration)).second) [[unlikely]] { + if (!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); - } } 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{}; @@ -131,27 +115,6 @@ template << 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); - } } } @@ -160,44 +123,86 @@ template const py::function &unflatten_func, const py::object &path_entry_type, const std::string ®istry_namespace) { - const scoped_write_lock lock{sm_mutex}; + // Emit the override warning for namedtuple / PyStructSequence subclasses BEFORE taking + // `sm_mutex`: the classification and `PyErr_WarnEx` run Python and release the GIL, and doing + // that under the write lock inverts the GIL to `sm_mutex` order against a concurrent flatten. + // Emitting before any insert also keeps registration atomic under warnings-as-errors. + 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"; + if (!registry_namespace.empty()) [[likely]] { + oss << " in namespace " << PyRepr(registry_namespace); + } + oss << "."; + if (PyErr_WarnEx(PyExc_UserWarning, oss.str().c_str(), /*stack_level=*/2) < 0) + [[unlikely]] { + throw py::error_already_set(); + } + } 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"; + if (!registry_namespace.empty()) [[likely]] { + oss << " in namespace " << PyRepr(registry_namespace); + } + oss << "."; + if (PyErr_WarnEx(PyExc_UserWarning, oss.str().c_str(), /*stack_level=*/2) < 0) + [[unlikely]] { + throw py::error_already_set(); + } + } + + // 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(); - RegisterImpl(cls, + { + const scoped_write_lock lock{sm_mutex}; + + registry1.RegisterImpl(cls, flatten_func, unflatten_func, path_entry_type, registry_namespace); - RegisterImpl(cls, + registry2.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(); + cls.inc_ref(); + flatten_func.inc_ref(); + unflatten_func.inc_ref(); + path_entry_type.inc_ref(); + } } -template -/*static*/ PyTreeTypeRegistry::RegistrationPtr PyTreeTypeRegistry::UnregisterImpl( +PyTreeTypeRegistry::RegistrationPtr 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]] { + const std::string ®istry_namespace, + const bool &is_structsequence_class, + const bool &is_namedtuple_class) { + if (m_builtins_types.find(cls) != m_builtins_types.end()) [[unlikely]] { throw py::value_error("PyTree type " + PyRepr(cls) + " is a built-in type and cannot be unregistered."); } if (registry_namespace.empty()) [[unlikely]] { - const auto it = registry.m_registrations.find(cls); - if (it == registry.m_registrations.end()) [[unlikely]] { + const auto it = m_registrations.find(cls); + if (it == m_registrations.end()) [[unlikely]] { std::ostringstream oss{}; oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { + if (is_structsequence_class) [[unlikely]] { oss << "is a class of `PyStructSequence`, " << "which is not explicitly registered in the global namespace."; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { + } else if (is_namedtuple_class) [[unlikely]] { oss << "is a subclass of `collections.namedtuple`, " << "which is not explicitly registered in the global namespace."; } else [[likely]] { @@ -206,18 +211,17 @@ template throw py::value_error(oss.str()); } RegistrationPtr registration = it->second; - registry.m_registrations.erase(it); + m_registrations.erase(it); return registration; } 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]] { + const auto named_it = m_named_registrations.find(std::make_pair(registry_namespace, cls)); + if (named_it == m_named_registrations.end()) [[unlikely]] { std::ostringstream oss{}; oss << "PyTree type " << PyRepr(cls) << " "; - if (IsStructSequenceClass(cls)) [[unlikely]] { + if (is_structsequence_class) [[unlikely]] { oss << "is a class of `PyStructSequence`, " << "which is not explicitly registered "; - } else if (IsNamedTupleClass(cls)) [[unlikely]] { + } else if (is_namedtuple_class) [[unlikely]] { oss << "is a subclass of `collections.namedtuple`, " << "which is not explicitly registered "; } else [[likely]] { @@ -227,43 +231,70 @@ template throw py::value_error(oss.str()); } RegistrationPtr registration = named_it->second; - registry.m_named_registrations.erase(named_it); + m_named_registrations.erase(named_it); return registration; } } /*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)); - registration1->type.dec_ref(); - registration1->flatten_func.dec_ref(); - registration1->unflatten_func.dec_ref(); - registration1->path_entry_type.dec_ref(); + // 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(); + + { + const scoped_write_lock lock{sm_mutex}; + + const auto registration1 = registry1.UnregisterImpl(cls, + registry_namespace, + is_structsequence_class, + is_namedtuple_class); + const auto registration2 = registry2.UnregisterImpl(cls, + registry_namespace, + is_structsequence_class, + is_namedtuple_class); + 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)); + registration1->type.dec_ref(); + registration1->flatten_func.dec_ref(); + registration1->unflatten_func.dec_ref(); + registration1->path_entry_type.dec_ref(); + } } 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( diff --git a/src/treespec/flatten.cpp b/src/treespec/flatten.cpp index 0894824d..d608ab47 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]] { diff --git a/tests/test_registry.py b/tests/test_registry.py index 3900b3e5..c08b0718 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 @@ -401,6 +406,226 @@ 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_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 +804,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 From 6448d5c2f1de06444e336c4ad222f38980030d03 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 27 Jul 2026 16:57:21 +0800 Subject: [PATCH 2/6] refactor(pytypes): extract WeakKeyCache and make the type caches interpreter-safe The namedtuple classification, PyStructSequence classification and struct sequence field-name caches each hand-rolled the same cache, weakref and locking dance. Extract one `WeakKeyCache` class and route all three through it, then fix the subinterpreter hazards in one place. Entries are keyed by `(interpreter_id, object address)` rather than the address alone. The address can be shared across interpreters, since a type like `int` is immortal and lives at the same address in each, while a computed reference value belongs to the interpreter that produced it. An address-only key would hand that value to another interpreter, which could use it after the owner freed it on finalization. A per-entry weakref evicts an entry when its key is collected, so a later key that reuses the freed 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` and `Py_Initialize` cycle, so a fresh interpreter must not inherit a finalized one's entries. The lookup also releases the read lock and re-acquires the GIL in that order before touching the borrowed value, so the GIL is never re-acquired while the lock is held, which would invert the order against the weakref eviction callback. `hashing.h` gains generic `std::equal_to` and `std::not_equal_to` for `std::pair`, mirroring the `std::hash` specialization already there, so the pair key works without a third copy of the element-wise comparison. --- CHANGELOG.md | 3 + include/optree/hashing.h | 21 +- include/optree/pytypes.h | 276 +++++++++++++---------- tests/concurrent/test_subinterpreters.py | 122 ++++++++++ tests/test_typing.py | 39 ++++ 5 files changed, 345 insertions(+), 116 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62156358..6c49a66c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 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 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). ### Removed diff --git a/include/optree/hashing.h b/include/optree/hashing.h index 862612ee..2d373411 100644 --- a/include/optree/hashing.h +++ b/include/optree/hashing.h @@ -19,7 +19,6 @@ limitations under the License. #include // std::size_t #include // std::hash, std::{not_,}equal_to -#include // std::string #include // std::pair #include @@ -89,6 +88,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/pytypes.h b/include/optree/pytypes.h index a7cac3ea..f2c127a8 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 #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::unordered_set +#include // std::forward, std::pair, std::make_pair, std::move #include @@ -34,6 +37,7 @@ limitations under the License. #include "optree/hashing.h" #include "optree/pymacros.h" +#include "optree/stdutils.h" #include "optree/synchronization.h" namespace py = pybind11; @@ -49,9 +53,6 @@ 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 \ @@ -213,6 +214,149 @@ 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 + [[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)(); + bool inserted = false; + bool register_cleanup = false; + { +#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; + } + register_cleanup = m_registered_interpids.insert(interpreter_id).second; + } + { + // The GIL is held here, so we can safely register the `atexit` callback and increment + // the reference count and create the weakref. + if (register_cleanup) [[unlikely]] { + RegisterInterpreterCleanup(interpreter_id); + } + if (inserted) [[likely]] { + if constexpr (kValueIsPyReference) { + value.inc_ref(); + } + (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(); + } + } + 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_registered_interpids.erase(interpreter_id); + })); + } + + std::unordered_map m_cache{}; + std::unordered_set m_registered_interpids{}; + 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) { // We can only identify namedtuples heuristically, here by the presence of a _fields attribute. @@ -257,40 +401,10 @@ inline bool IsNamedTupleClass(const py::handle &type) { 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) { return IsNamedTupleClass(py::type::handle_of(object)); @@ -368,40 +482,10 @@ inline bool IsStructSequenceClass(const py::handle &type) { 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) { return IsStructSequenceClass(py::type::handle_of(object)); @@ -460,48 +544,10 @@ 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); + }); } inline void TotalOrderSort(py::list &list) { // NOLINT[runtime/references] diff --git a/tests/concurrent/test_subinterpreters.py b/tests/concurrent/test_subinterpreters.py index aec86e9a..ad50ee63 100644 --- a/tests/concurrent/test_subinterpreters.py +++ b/tests/concurrent/test_subinterpreters.py @@ -343,3 +343,125 @@ 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_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/test_typing.py b/tests/test_typing.py index 51cb95cc..781a7841 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -32,10 +32,14 @@ CustomTuple, Py_GIL_DISABLED, Vector2D, + check_script_in_subprocess, disable_systrace, gc_collect, getrefcount, + skipif_android, + skipif_ios, skipif_pypy, + skipif_wasm, ) @@ -660,3 +664,38 @@ 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, + ) From 6079c23cb38cdadfdf6161247c6af5de4932ae1c Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 27 Jul 2026 17:08:15 +0800 Subject: [PATCH 3/6] fix(treespec): map struct sequence fields by offset and validate pickled state Both areas run through `serialization.cpp` and the struct sequence helpers. `tp_members` lists only the named fields, so indexing it by position mislabelled every slot after the first unnamed one: `os.stat_result` slots 7, 8 and 9 were reported as `st_atime`, `st_mtime` and `st_ctime`. Map by the byte offset each member carries instead, and fill the remaining slots with the unnamed marker. `structseq_fields` is fixed the same way, the marker is exported as `PyStructSequence_UnnamedField`, the repr renders an unnamed slot as ``, and `StructSequenceEntry.codify` emits index access for a slot whose name is not an identifier. `PyStructSequenceUnnamedField()` works around the symbol being unexported before 3.11.0a2, where referencing it broke the import. `FromPicklable` trusted its input, so a crafted pickle could build a spec that later read out of bounds or aborted in repr. Validate the kind before the narrowing cast, the arity against the children the traversal provides, node data on childless kinds, dict key count and distinctness, defaultdict shape, deque `maxlen`, namedtuple and struct sequence arity, and the whole traversal rather than only its last node. `ToPicklable` copies the node's mutable containers so the state cannot alias the spec, and `__reduce__` makes protocols 0 and 1 reconstruct through `cls.__new__(cls)` instead of aborting. --- CHANGELOG.md | 5 + docs/source/spelling_wordlist.txt | 1 + include/optree/pymacros.h | 13 + include/optree/pytypes.h | 49 +++- include/optree/registry.h | 8 +- optree/_C.pyi | 3 + optree/accessors.py | 5 +- optree/typing.py | 39 ++- src/optree.cpp | 18 +- src/treespec/serialization.cpp | 301 +++++++++++++++++----- tests/test_treespec.py | 414 +++++++++++++++++++++++++++++- tests/test_typing.py | 53 ++++ 12 files changed, 830 insertions(+), 79 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c49a66c..a094fc1c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 `optree.structseq_fields()` and the struct sequence accessors reporting the trailing hidden field names for unnamed sequence slots (e.g. `os.stat_result` slots 7-9 reported as `st_atime` / `st_mtime` / `st_ctime`), by mapping `tp_members` by offset instead of by position 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 `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 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). ### Removed diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index 3687fbbd..07d4b7b4 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -70,6 +70,7 @@ optree OrderedDict param params +picklable pragma py pypy diff --git a/include/optree/pymacros.h b/include/optree/pymacros.h index e0469047..f326a94f 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) && \ diff --git a/include/optree/pytypes.h b/include/optree/pytypes.h index f2c127a8..57ba40b2 100644 --- a/include/optree/pytypes.h +++ b/include/optree/pytypes.h @@ -25,6 +25,7 @@ limitations under the License. #include // std::unordered_map #include // std::unordered_set #include // std::forward, std::pair, std::make_pair, std::move +#include // std::vector #include @@ -508,23 +509,50 @@ 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]) { + TupleSetItem(fields, i, py::str(PyStructSequenceUnnamedField())); + } } return fields; #endif @@ -646,3 +674,12 @@ inline std::pair DictKeysDifference(const py::list & /*uniqu TotalOrderSort(extra_keys); return std::make_pair(std::move(missing_keys), std::move(extra_keys)); } + +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 2c755c7c..2f5581c6 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; diff --git a/optree/_C.pyi b/optree/_C.pyi index 2643289c..a35ab3af 100644 --- a/optree/_C.pyi +++ b/optree/_C.pyi @@ -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] diff --git a/optree/accessors.py b/optree/accessors.py index 181b3df8..7a8e0729 100644 --- a/optree/accessors.py +++ b/optree/accessors.py @@ -325,7 +325,10 @@ def __repr__(self, /) -> str: def codify(self, /, node: str = '') -> str: """Generate code for accessing the path entry.""" - return f'{node}.{self.field}' + field = self.field + if not field.isidentifier(): # an unnamed PyStructSequence slot -> access by index + return f'{node}[{self.entry}]' + return f'{node}.{field}' class DataclassEntry(GetAttrEntry): diff --git a/optree/typing.py b/optree/typing.py index 77dc4ce4..fd6c6d04 100644 --- a/optree/typing.py +++ b/optree/typing.py @@ -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, ...]: @@ -564,19 +569,33 @@ def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: raise TypeError(f'Expected an instance of PyStructSequence type, got {obj!r}.') 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 = { + 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 = [ - name - for name, member in vars(cls).items() - if isinstance(member, StructSequenceFieldType) - ] - return tuple(fields[: cls.n_sequence_fields]) # type: ignore[attr-defined] + return tuple( + names_by_index.get(index, PyStructSequence_UnnamedField) + for index in range(cls.n_sequence_fields) # type: ignore[attr-defined] + ) + + # pragma: pypy no cover + # CPython's `member_descriptor` does not expose the field offset, so the exact position of an + # unnamed slot is not recoverable in pure Python (the C++ implementation maps by offset). Assume + # unnamed sequence fields trail the named ones (as in `os.stat_result`): keep the named sequence + # fields, then fill the remaining slots with the marker. + named = [ + name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) + ] + n_sequence_fields: int = cls.n_sequence_fields # type: ignore[attr-defined] + n_unnamed_fields: int = cls.n_unnamed_fields # type: ignore[attr-defined] + return tuple(named[: n_sequence_fields - n_unnamed_fields]) + ( + (PyStructSequence_UnnamedField,) * n_unnamed_fields + ) del _tp_cache diff --git a/src/optree.cpp b/src/optree.cpp index f1dd9fd0..a5e51bdd 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{}; @@ -510,7 +511,22 @@ void BuildModule(py::module_ &mod) { // NOLINT[runtime/references] }), "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) diff --git a/src/treespec/serialization.cpp b/src/treespec/serialization.cpp index 04a36052..36df36c9 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,31 @@ py::object PyTreeSpec::ToPicklable() const { const scoped_critical_section2 cs{ node.custom != nullptr ? py::handle{node.custom->type} : py::handle{}, node.node_data}; - 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(), - 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())); + + // 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 = node.node_data.attr("copy")(); + } else if (node.kind == PyTreeKind::DefaultDict) [[unlikely]] { + const auto metadata = py::reinterpret_borrow(node.node_data); + node_data = + py::make_tuple(TupleGetItem(metadata, 0), TupleGetItem(metadata, 1).attr("copy")()); + } + + TupleSetItem( + node_states, + i++, + py::make_tuple(py::int_(static_cast(node.kind)), + py::int_(node.arity), + 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.attr("copy")() : py::none())); } return py::make_tuple(node_states, py::bool_(m_none_is_leaf), py::str(m_namespace)); } @@ -321,9 +354,13 @@ 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 + "."); + }; + 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 +369,154 @@ 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]] { + 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]); + node.node_data = 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 (DistinctCount(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. + if (!(metadata[0].is_none() || + static_cast(PyCallable_Check(metadata[0].ptr())))) [[unlikely]] { + throw malformed("the `default_factory` is not callable"); + } + const auto keys = thread_safe_cast(metadata[1]); + if (ListGetSize(keys) != node.arity) [[unlikely]] { + throw malformed("the number of keys does not match the arity"); + } + if (DistinctCount(keys) != node.arity) [[unlikely]] { + throw malformed("the keys are not distinct"); + } + node.node_data = metadata; + 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; } @@ -390,23 +525,23 @@ py::object PyTreeSpec::ToPicklable() const { 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_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 +550,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."); + } 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)); + for (const py::handle &key : keys) { + const int contains = PyDict_Contains(node.original_keys.ptr(), key.ptr()); + if (contains < 0) [[unlikely]] { + throw py::error_already_set(); + } + if (contains == 0) [[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/tests/test_treespec.py b/tests/test_treespec.py index 643176a0..4a4f5124 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -17,6 +17,7 @@ import contextlib import itertools +import os import pickle import platform import re @@ -24,8 +25,9 @@ import subprocess import sys import tempfile +import time import weakref -from collections import OrderedDict, UserList, defaultdict, deque +from collections import OrderedDict, UserList, defaultdict, deque, namedtuple import pytest @@ -260,6 +262,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 +297,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: @@ -541,6 +606,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 +680,330 @@ 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_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, '')) + + @parametrize( tree=TREES, none_is_leaf=[False, True], diff --git a/tests/test_typing.py b/tests/test_typing.py index 781a7841..d831300b 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,6 +29,7 @@ import optree from helpers import ( PYBIND11_HAS_NATIVE_ENUM, + PYPY, CustomNamedTupleSubclass, CustomTuple, Py_GIL_DISABLED, @@ -551,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, @@ -601,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(): From d63f087674fc1eef9df1b5a32eb93150aedcf573 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 27 Jul 2026 17:09:56 +0800 Subject: [PATCH 4/6] fix(treespec): correct namespace merging, GC reporting and walker recursion Namespace merging. `compose`, `broadcast_to_common_suffix`, `transform` and `treespec_from_collection` let an empty namespace adopt the other side's namespace. If a custom node resolved globally but resolves to a *different* registration under the adopted namespace, the result claimed the namespace while carrying the wrong registration. Add `FindReregisteredCustomType` and reject those merges; a type registered only globally still resolves by fallback and is allowed. The dict constructors also dropped the namespace, losing insertion-ordered key order, and `MakeFromCollection` ignored the return value of `PyErr_WarnEx`, so an escalated warning surfaced as a confusing `SystemError`. GC. `PyTreeIter` never visited or cleared its leaf predicate, so a cycle through that callback leaked. `PyTreeSpec` reports the registration's members only when its node is their sole owner: the registration holds one reference to each member however many nodes point at it, so reporting while shared decrements the same object once per node and underflows the collector's shadow refcount, which is fatal on debug builds. Recursion. `PathsImpl`, `AccessorsImpl` and `BroadcastToCommonSuffixImpl` recurse once per tree level, so a deeply nested spec overflowed the native stack and crashed instead of raising. Thread a depth through them and raise `RecursionError` at `MAX_RECURSION_DEPTH`, which is lowered further for debug builds on Windows, where the frames are large and the throw still has to unwind them. Also fixes `IsPrefix` snapshotting a dict node's subtree from the pristine traversal rather than the working copy (a processed ancestor may have relocated it), `broadcast_to_common_suffix` sorting the argument spec's live key list in place while building an error message, custom node entries being dropped when broadcasting, and `Transform` reading pending counts after `pop_back`. --- CHANGELOG.md | 9 + include/optree/treespec.h | 22 +- src/treespec/constructors.cpp | 50 ++- src/treespec/gc.cpp | 22 ++ src/treespec/richcomparison.cpp | 9 +- src/treespec/treespec.cpp | 145 ++++++- tests/helpers.py | 8 + tests/test_treespec.py | 663 ++++++++++++++++++++++++++++++++ 8 files changed, 900 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a094fc1c..8babf60d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 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 `PyTreeSpec.broadcast_to_common_suffix()` dropping a custom node's entries and falling back to positional ones 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_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 reference cycle passing through a registered custom type not being collectable once the registry no longer holds the registration 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). ### Removed diff --git a/include/optree/treespec.h b/include/optree/treespec.h index 1e132a51..3d8fb2dd 100644 --- a/include/optree/treespec.h +++ b/include/optree/treespec.h @@ -39,7 +39,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.0 +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; @@ -364,7 +372,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 that resolves to a different registration under + // `target_namespace` than it currently holds, 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 FindReregisteredCustomType( + const std::string &target_namespace) const; template [[nodiscard]] py::object WalkImpl( @@ -431,7 +447,7 @@ 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; diff --git a/src/treespec/constructors.cpp b/src/treespec/constructors.cpp index 380f6b84..1664435f 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{}; @@ -114,7 +122,13 @@ template << ", got " << PyRepr(common_registry_namespace) << "."; throw py::value_error(oss.str()); } - } else if (node.kind != PyTreeKind::Custom) [[likely]] { + } else if (node.kind != PyTreeKind::Custom && + ((node.kind != PyTreeKind::Dict && node.kind != PyTreeKind::DefaultDict) || + !is_dict_insertion_ordered_in_current_namespace)) [[likely]] { + // Drop the namespace for namespace-independent nodes. A Dict/DefaultDict whose keys are + // kept in insertion order does depend on the namespace (see the sort at the Dict case), + // so keep it there, mirroring `Flatten` + // (`is_dict_insertion_ordered_in_current_namespace`). registry_namespace = ""; } }; @@ -122,9 +136,11 @@ template switch (node.kind) { case PyTreeKind::Leaf: { node.arity = 0; - PyErr_WarnEx(PyExc_UserWarning, - "PyTreeSpec::MakeFromCollection() is called on a leaf.", - /*stack_level=*/2); + if (PyErr_WarnEx(PyExc_UserWarning, + "PyTreeSpec::MakeFromCollection() is called on a leaf.", + /*stack_level=*/2) < 0) [[unlikely]] { + throw py::error_already_set(); + } break; } @@ -170,8 +186,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 +287,23 @@ 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 reregistered_type = out->FindReregisteredCustomType(registry_namespace)) + [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be composed into a collection: custom PyTree type " + << PyRepr(*reregistered_type) + << " resolves to a different 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/gc.cpp b/src/treespec/gc.cpp index f6c3700d..a030865f 100644 --- a/src/treespec/gc.cpp +++ b/src/treespec/gc.cpp @@ -43,6 +43,17 @@ namespace optree { Py_VISIT(node.node_data.ptr()); Py_VISIT(node.node_entries.ptr()); Py_VISIT(node.original_keys.ptr()); + // Report the registration's members only when this node is their sole owner. + // The registration holds one reference to each member however many nodes point at it, so + // reporting while it is shared would decrement the same object once per node and underflow + // its shadow refcount. While the registry holds the registration it also keeps the members + // alive, so skipping then leaks nothing. + if (node.custom != nullptr && node.custom.use_count() == 1) [[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; } @@ -58,11 +69,13 @@ namespace optree { 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 +87,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 +106,10 @@ namespace optree { } self.m_agenda.clear(); Py_CLEAR(self.m_root.ptr()); + if (self.m_leaf_predicate) [[likely]] { + // Drop the owned leaf predicate reference to break cycles that pass through it. + Py_CLEAR(self.m_leaf_predicate->ptr()); + } 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/treespec.cpp b/src/treespec/treespec.cpp index 44305dd0..0d0f827e 100644 --- a/src/treespec/treespec.cpp +++ b/src/treespec/treespec.cpp @@ -185,13 +185,36 @@ namespace optree { } } +std::optional PyTreeSpec::FindReregisteredCustomType( + 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 reregistered_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, other.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + reregistered_type = FindReregisteredCustomType(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."); + reregistered_type = other.FindReregisteredCustomType(target_namespace); + } + if (reregistered_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*reregistered_type) + << " resolves to a different 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,29 @@ 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 &reregistered_type = + treespec->FindReregisteredCustomType(common_registry_namespace)) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be transformed: custom PyTree type " + << PyRepr(*reregistered_type) + << " resolves to a different 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 +674,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 reregistered_type{}; + if (target_namespace != m_namespace) [[unlikely]] { + EXPECT_EQ(target_namespace, inner.m_namespace, "Namespace mismatch."); + EXPECT_TRUE(m_namespace.empty(), "Namespace mismatch."); + reregistered_type = FindReregisteredCustomType(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."); + reregistered_type = inner.FindReregisteredCustomType(target_namespace); + } + if (reregistered_type) [[unlikely]] { + std::ostringstream oss{}; + oss << "PyTreeSpecs cannot be merged: custom PyTree type " << PyRepr(*reregistered_type) + << " resolves to a different 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 +743,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 +846,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); 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/test_treespec.py b/tests/test_treespec.py index 4a4f5124..942f14e7 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -16,7 +16,9 @@ # pylint: disable=missing-function-docstring,invalid-name import contextlib +import gc import itertools +import math import os import pickle import platform @@ -26,6 +28,7 @@ import sys import tempfile import time +import warnings import weakref from collections import OrderedDict, UserList, defaultdict, deque, namedtuple @@ -35,6 +38,7 @@ import optree from helpers import ( GLOBAL_NAMESPACE, + HAS_DEFERRED_TYPE_REFS, NAMESPACED_TREE, PYPY, STANDARD_DICT_TYPES, @@ -50,6 +54,7 @@ parametrize, recursionlimit, skipif_android, + skipif_deferred_type_refs, skipif_ios, skipif_pypy, skipif_wasm, @@ -409,6 +414,72 @@ 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 +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 + + @disable_systrace def test_treeiter_self_referential(): sentinel = object() @@ -440,6 +511,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 @@ -1156,6 +1245,541 @@ 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='different 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='different 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='different 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(): + # Regression: BroadcastToCommonSuffixImpl rebuilds each non-leaf node with a designated + # initializer that lists `.node_data` then jumps to `.custom`, silently skipping `.node_entries` + # (which is declared between them). C++ value-initializes the omitted member to a null handle, + # so a custom node registered with explicit path entries (a 3-tuple flatten `(children, + # metadata, entries)`) lost them in the broadcasted spec. `entries()`/`paths()`/`accessors()` + # then fell back to the `range(arity)` integer indices, producing not just different but BROKEN + # accessors (`GetAttrEntry(entry=0)` calls `getattr(obj, 0)`, a `TypeError`, instead of + # `getattr(obj, 'a')`). + 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='different 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='different 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='different 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='different 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='different 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], @@ -1832,6 +2456,21 @@ 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) + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -2196,6 +2835,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]) From 1491672ab014e1f66383cdc1c4e6caa6fe95e52d Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 27 Jul 2026 17:11:18 +0800 Subject: [PATCH 5/6] fix(dataclasses,attrs,accessors,ops): correct field selection and broadcasting `children_fields` was derived from the `_FIELDS` attribute, which lists every field, while an integer accessor entry indexes only the tree children: the fields that are both `pytree_node=True` and `init=True`. A class with a metadata or non-init field therefore resolved an integer entry to the wrong attribute. Compute the children from `dataclasses.fields()` / `attrs.fields()` with both conditions applied, in the dataclass and attrs entries alike. `register_node` set its `_FIELDS` guard before `register_pytree_node()` ran, so a failed registration left the class marked as registered and unusable afterwards. Set the guard only on success, and say in the error which API to use to register the class in another namespace. `InitVar` pseudo-fields are excluded from `dataclasses.fields()`, so they are neither children nor metadata and cannot round trip: reject them and point at the generic API. Both modules now document that the generated unflatten rebuilds via `cls(**fields)`, which re-runs `__init__`, so the round trip is exact only when it returns each field unchanged. `tree_broadcast_common` fills an internal tree with a private sentinel and then re-applied the caller's `is_leaf` to it, so a predicate that inspects leaf values could crash on the sentinel or collapse the filled subtree and under-replicate. Do not apply the predicate to the internal tree, and document that a value-based predicate can still yield two trees that re-flatten differently. `treespec_is_prefix` and `treespec_is_suffix` gain the preorder semantics they actually implement: reflexive and transitive, but neither antisymmetric nor total. --- CHANGELOG.md | 4 + docs/source/spelling_wordlist.txt | 2 + optree/accessors.py | 18 +++- optree/dataclasses.py | 42 +++++++++- optree/integrations/attrs.py | 39 +++++++-- optree/ops.py | 91 +++++++++++++++++++- tests/integrations/test_attrs.py | 104 +++++++++++++++++++++-- tests/test_dataclasses.py | 133 +++++++++++++++++++++++++++++- tests/test_ops.py | 93 +++++++++++++++++++++ 9 files changed, 505 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8babf60d..59f650d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 reference cycle passing through a registered custom type not being collectable once the registry no longer holds the registration 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 field rather than only the tree children, which returned the wrong attribute for a class holding a metadata or non-`init` field 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()` marking a class as registered before the registration succeeded, leaving the class permanently unusable after a failure 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). ### Removed diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index 07d4b7b4..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 @@ -72,6 +73,7 @@ param params picklable pragma +preorder py pypy pytree diff --git a/optree/accessors.py b/optree/accessors.py index 7a8e0729..1380e14e 100644 --- a/optree/accessors.py +++ b/optree/accessors.py @@ -348,11 +348,27 @@ 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 child (pytree-node) field names.""" + # pylint: disable-next=import-outside-toplevel,cyclic-import + from optree.dataclasses import _PYTREE_NODE_DEFAULT + + # The children are the `pytree_node=True` fields (a subset of the init fields); an integer + # entry indexes these. Computed from the fields rather than the integration-only + # `__optree_dataclass_fields__` attribute, so it also works for classes registered via the + # generic `register_pytree_node`. + return tuple( + f.name + for f in dataclasses.fields(self.type) + if f.init and f.metadata.get('pytree_node', _PYTREE_NODE_DEFAULT) + ) + @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 diff --git a/optree/dataclasses.py b/optree/dataclasses.py index 298397e9..4513d3b8 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, @@ -609,4 +639,8 @@ 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. + setattr(cls, _FIELDS, (MappingProxyType(children_fields), MappingProxyType(metadata_fields))) return cls diff --git a/optree/integrations/attrs.py b/optree/integrations/attrs.py index 8a43d582..141e0604 100644 --- a/optree/integrations/attrs.py +++ b/optree/integrations/attrs.py @@ -143,11 +143,24 @@ 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 child (pytree-node) field names.""" + # The children are the `pytree_node=True` fields (a subset of the init fields); an integer + # entry indexes these. Computed from the fields rather than read from the `_FIELDS` + # attribute, which is only set by the `optree.integrations.attrs` integration: an attrs + # class registered directly via the generic `register_pytree_node` has no `_FIELDS`. + return tuple( + a.name + for a in attrs.fields(self.type) + if a.init and a.metadata.get('pytree_node', _PYTREE_NODE_DEFAULT) + ) + @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 +245,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 +381,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 +438,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 +472,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, @@ -475,4 +500,8 @@ 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. + setattr(cls, _FIELDS, (MappingProxyType(children_fields), MappingProxyType(metadata_fields))) return cls diff --git a/optree/ops.py b/optree/ops.py index 8db8183c..fb3bdc7d 100644 --- a/optree/ops.py +++ b/optree/ops.py @@ -1727,6 +1727,13 @@ def tree_broadcast_common( ``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]) >>> tree_broadcast_common([1, 2, 3], [4, 5, 6]) @@ -1775,9 +1782,13 @@ def tree_broadcast_common( ) def broadcast_leaves(x: T, subtree: PyTree[T]) -> PyTree[T]: + # `subtree` is a slice of `common_suffix_tree`, whose leaves are the private `sentinel` + # placeholder rather than real input values. Do not run the user's `is_leaf` on it: the + # subtree structure is already fixed by `common_suffix_treespec`, and `is_leaf` may crash or + # wrongly collapse them. Flatten fully to count how many replicas of `x` are needed. subtreespec = tree_structure( subtree, - is_leaf=is_leaf, + is_leaf=None, none_is_leaf=none_is_leaf, namespace=namespace, ) @@ -1821,6 +1832,13 @@ def broadcast_common( ``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. + >>> broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) >>> broadcast_common([1, 2, 3], [4, 5, 6]) @@ -2974,10 +2992,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 +3075,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. diff --git a/tests/integrations/test_attrs.py b/tests/integrations/test_attrs.py index d8b7ac5c..44ea115b 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,52 @@ 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_init_false_class_warns(): @attrs.define(init=False) @@ -356,9 +430,29 @@ 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(): + # An integer entry indexes the tree CHILDREN (the fields that are BOTH `pytree_node=True` and + # `init`), not all init fields and not all `pytree_node` fields. A non-child field interleaved + # between children (a `pytree_node=False` field OR a non-`init` field) must not shift the mapping. + @attrs.define + class Foo: + a: int + b: int = optree.integrations.attrs.field(default=0, pytree_node=False) # not a child + d: int = attrs.field(init=False, default=0) # non-init -> not a tree child + 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_accessor_codify(): diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py index 8a4bae98..ca8c8595 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,94 @@ 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_init_false_class_warns(): @dataclasses.dataclass(init=False) @@ -1131,6 +1234,28 @@ 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(): + # An integer entry indexes the tree CHILDREN (the fields that are BOTH `pytree_node=True` and + # `init`), not all init fields and not all `pytree_node` fields. A non-child field interleaved + # between children (a `pytree_node=False` field OR a non-`init` field) must not shift the mapping. + @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 = dataclasses.field(init=False, default=0) + 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 diff --git a/tests/test_ops.py b/tests/test_ops.py index 04fd8e89..bf0525c7 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -3120,6 +3120,99 @@ 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_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]) From e528bb72d7fd03433b98d00e441931aabcd53ab6 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Wed, 29 Jul 2026 15:27:38 +0800 Subject: [PATCH 6/6] wip --- include/optree/pytypes.h | 51 ++++++-- include/optree/treespec.h | 1 + optree/_C.pyi | 32 +++-- optree/accessors.py | 16 +-- optree/ops.py | 223 +++++++++++++++------------------- optree/typing.py | 54 +++++--- src/treespec/constructors.cpp | 12 +- src/treespec/gc.cpp | 44 +++++-- tests/test_accessors.py | 65 ++++++++++ tests/test_ops.py | 104 ++++++++++++++++ tests/test_treespec.py | 134 ++++++++++++++++++++ tests/test_typing.py | 85 +++++++++++++ 12 files changed, 635 insertions(+), 186 deletions(-) diff --git a/include/optree/pytypes.h b/include/optree/pytypes.h index 57ba40b2..4c2807ea 100644 --- a/include/optree/pytypes.h +++ b/include/optree/pytypes.h @@ -280,11 +280,33 @@ class WeakKeyCache { } ValueType value = std::forward(compute)(); - bool inserted = false; + + // Register the per-interpreter cleanup BEFORE publishing an entry. It runs Python and can + // raise, and the interpreter is marked only once it succeeds: marking first would leave the + // interpreter believing a callback exists and never retry it. bool register_cleanup = false; { #if !defined(Py_GIL_DISABLED) const py::gil_scoped_release_simple gil_release{}; +#endif + const scoped_write_lock lock{m_mutex}; + register_cleanup = m_registered_interpids.insert(interpreter_id).second; + } + if (register_cleanup) [[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_registered_interpids.erase(interpreter_id); + throw; + } + } + + bool inserted = false; + { +#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]] { @@ -293,18 +315,16 @@ class WeakKeyCache { // below). inserted = m_cache.emplace(cache_key, StoredType{value}).second; } - register_cleanup = m_registered_interpids.insert(interpreter_id).second; } - { - // The GIL is held here, so we can safely register the `atexit` callback and increment - // the reference count and create the weakref. - if (register_cleanup) [[unlikely]] { - RegisterInterpreterCleanup(interpreter_id); + 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(); } - if (inserted) [[likely]] { - 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}; @@ -318,6 +338,15 @@ class WeakKeyCache { 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; diff --git a/include/optree/treespec.h b/include/optree/treespec.h index 3d8fb2dd..ee3195c1 100644 --- a/include/optree/treespec.h +++ b/include/optree/treespec.h @@ -269,6 +269,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; diff --git a/optree/_C.pyi b/optree/_C.pyi index a35ab3af..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 ( @@ -102,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: Callable[[Collection[U]], U] | None = None, + 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] = ..., 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: ... @@ -191,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 1380e14e..e8958dc6 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 diff --git a/optree/ops.py b/optree/ops.py index fb3bdc7d..2e8f7952 100644 --- a/optree/ops.py +++ b/optree/ops.py @@ -1320,7 +1320,10 @@ 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 required to infer the inner structure from the first output. Given an explicit + # `inner_treespec`, a leafless outer structure transposes to that structure filled with copies + # of itself, 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 +1339,9 @@ 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) + # With no outer leaves there is nothing to zip, and every inner slot takes an empty copy of the + # outer structure rather than no copy at all. + 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 +1412,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 +1428,9 @@ 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) + # With no outer leaves there is nothing to zip, and every inner slot takes an empty copy of the + # outer structure rather than no copy at all. + 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 +1528,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 +1544,9 @@ 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) + # With no outer leaves there is nothing to zip, and every inner slot takes an empty copy of the + # outer structure rather than no copy at all. + 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 +1715,60 @@ 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. + """ + 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) + if not rests: + return (tree,), 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` is a slice of `common_suffix_tree`, whose leaves are the private `sentinel` + # placeholder rather than real input values. Do not run the user's `is_leaf` on it: the + # subtree structure is already fixed by `common_suffix_treespec`, and `is_leaf` may crash or + # wrongly collapse them. Flatten fully to count how many replicas of `x` are needed. + 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], @@ -1772,41 +1835,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]: - # `subtree` is a slice of `common_suffix_tree`, whose leaves are the private `sentinel` - # placeholder rather than real input values. Do not run the user's `is_leaf` on it: the - # subtree structure is already fixed by `common_suffix_treespec`, and `is_leaf` may crash or - # wrongly collapse them. Flatten fully to count how many replicas of `x` are needed. - 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_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 @@ -1832,13 +1866,6 @@ def broadcast_common( ``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. - >>> broadcast_common(1, [2, 3, 4]) ([1, 1, 1], [2, 3, 4]) >>> broadcast_common([1, 2, 3], [4, 5, 6]) @@ -1876,67 +1903,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], @@ -1992,22 +1972,17 @@ 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( - 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], @@ -2071,19 +2046,15 @@ 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( - 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( @@ -2164,19 +2135,15 @@ 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( - 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 diff --git a/optree/typing.py b/optree/typing.py index fd6c6d04..8cfe338b 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) @@ -568,33 +568,47 @@ 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 # 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 = { + names_by_index: dict[int, str] = { member.index: name # type: ignore[attr-defined] for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) } - return tuple( - names_by_index.get(index, PyStructSequence_UnnamedField) - for index in range(cls.n_sequence_fields) # type: ignore[attr-defined] - ) - - # pragma: pypy no cover - # CPython's `member_descriptor` does not expose the field offset, so the exact position of an - # unnamed slot is not recoverable in pure Python (the C++ implementation maps by offset). Assume - # unnamed sequence fields trail the named ones (as in `os.stat_result`): keep the named sequence - # fields, then fill the remaining slots with the marker. - named = [ - name for name, member in vars(cls).items() if isinstance(member, StructSequenceFieldType) - ] - n_sequence_fields: int = cls.n_sequence_fields # type: ignore[attr-defined] - n_unnamed_fields: int = cls.n_unnamed_fields # type: ignore[attr-defined] - return tuple(named[: n_sequence_fields - n_unnamed_fields]) + ( - (PyStructSequence_UnnamedField,) * n_unnamed_fields + else: # pragma: pypy no cover + # CPython's `member_descriptor` does not expose the field offset, so a field's position is + # not directly readable in pure Python (the C++ implementation maps by offset). `vars()` + # yields the members in field order and only sequence slots can be unnamed, so the leading + # names are the named sequence fields in increasing position order, but their positions are + # still unknown. + named = [ + name + for name, member in vars(cls).items() + if isinstance(member, StructSequenceFieldType) + ][: 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): + 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) ) diff --git a/src/treespec/constructors.cpp b/src/treespec/constructors.cpp index 1664435f..feb14bf5 100644 --- a/src/treespec/constructors.cpp +++ b/src/treespec/constructors.cpp @@ -122,9 +122,11 @@ template << ", got " << PyRepr(common_registry_namespace) << "."; throw py::value_error(oss.str()); } - } else if (node.kind != PyTreeKind::Custom && - ((node.kind != PyTreeKind::Dict && node.kind != PyTreeKind::DefaultDict) || - !is_dict_insertion_ordered_in_current_namespace)) [[likely]] { + } else if (const bool depends_on_namespace = + node.kind == PyTreeKind::Custom || + ((node.kind == PyTreeKind::Dict || node.kind == PyTreeKind::DefaultDict) && + is_dict_insertion_ordered_in_current_namespace); + !depends_on_namespace) [[likely]] { // Drop the namespace for namespace-independent nodes. A Dict/DefaultDict whose keys are // kept in insertion order does depend on the namespace (see the sort at the Dict case), // so keep it there, mirroring `Flatten` @@ -136,6 +138,9 @@ template switch (node.kind) { case PyTreeKind::Leaf: { node.arity = 0; + // 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]] { @@ -146,6 +151,7 @@ template case PyTreeKind::None: { node.arity = 0; + registry_namespace = ""; if constexpr (!NoneIsLeaf) { break; } diff --git a/src/treespec/gc.cpp b/src/treespec/gc.cpp index a030865f..fc5e131f 100644 --- a/src/treespec/gc.cpp +++ b/src/treespec/gc.cpp @@ -15,6 +15,8 @@ limitations under the License. ================================================================================ */ +#include // std::unordered_map + #include "optree/optree.h" inline namespace { @@ -39,20 +41,42 @@ namespace optree { } auto &self = thread_safe_cast(py::handle{self_base}); PYTREESPEC_SANITY_CHECK(self); + std::unordered_map num_holders{}; for (const auto &node : self.m_traversal) { Py_VISIT(node.node_data.ptr()); Py_VISIT(node.node_entries.ptr()); Py_VISIT(node.original_keys.ptr()); - // Report the registration's members only when this node is their sole owner. - // The registration holds one reference to each member however many nodes point at it, so - // reporting while it is shared would decrement the same object once per node and underflow - // its shadow refcount. While the registry holds the registration it also keeps the members - // alive, so skipping then leaks nothing. - if (node.custom != nullptr && node.custom.use_count() == 1) [[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()); + if (node.custom != nullptr) [[unlikely]] { + ++num_holders[node.custom.get()]; // inits to 0 (ssize_t{}) if not present + } + } + if (!num_holders.empty()) [[unlikely]] { + // 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. + // + // 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. + for (const auto &node : self.m_traversal) { + if (node.custom == nullptr) [[likely]] { + continue; + } + const auto it = num_holders.find(node.custom.get()); + if (it == num_holders.end()) [[unlikely]] { + continue; // already handled with an earlier node holding the same registration + } + const ssize_t holders = it->second; + num_holders.erase(it); + if (node.custom.use_count() == 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; diff --git a/tests/test_accessors.py b/tests/test_accessors.py index bd50bc10..79aff679 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -19,6 +19,7 @@ import itertools import re from collections import OrderedDict, UserDict, UserList, defaultdict, deque +from operator import itemgetter from typing import Any, NamedTuple import pytest @@ -228,6 +229,70 @@ def test_pytree_accessor_equal_hash(none_is_leaf): assert hash(accessor1) != hash(accessor2) +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')" + + def test_pytree_entry_init(): for path_entry_type in ( optree.PyTreeEntry, diff --git a/tests/test_ops.py b/tests/test_ops.py index bf0525c7..f5089276 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -2824,6 +2824,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], @@ -3213,6 +3252,71 @@ def is_shape_like(x): ) == ([((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_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_treespec.py b/tests/test_treespec.py index 942f14e7..a13d4a98 100644 --- a/tests/test_treespec.py +++ b/tests/test_treespec.py @@ -445,6 +445,37 @@ class Cyclic: 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. @@ -480,6 +511,77 @@ class Shared: 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() @@ -2036,6 +2138,23 @@ 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 + + @parametrize( tree=TREES, none_is_leaf=[False, True], @@ -2471,6 +2590,21 @@ def test_treespec_from_collection_on_leaf_propagates_escalated_warning(): 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], diff --git a/tests/test_typing.py b/tests/test_typing.py index d831300b..cc96ce26 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -752,3 +752,88 @@ def test_type_caches_register_interpreter_cleanup(): """, 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 failed insert must not be observable: the value is recomputed, not read back from a + # dangling entry. + fields = optree.structseq_fields(time.struct_time) + assert fields[:2] == ('tm_year', 'tm_mon'), fields + + # The interpreter must not be marked as cleaned-up-registered by the failed attempt, or the + # cleanup would never be retried. + before = atexit._ncallbacks() + optree.structseq_fields(time.struct_time) + assert atexit._ncallbacks() >= before + """, + 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, + )