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