From b2d42a738816128888a36f479d934ccd0dd1bf8d Mon Sep 17 00:00:00 2001 From: Sskift <163287349+Sskift@users.noreply.github.com> Date: Tue, 3 Nov 2026 02:02:47 +0800 Subject: [PATCH 1/2] Pass only supported context to custom encoders --- tests/test_encoder_dispatch.py | 158 +++++++++++++++++++++++++++++++++ tomlkit/items.py | 34 +++++-- 2 files changed, 185 insertions(+), 7 deletions(-) create mode 100644 tests/test_encoder_dispatch.py diff --git a/tests/test_encoder_dispatch.py b/tests/test_encoder_dispatch.py new file mode 100644 index 00000000..f10ac17d --- /dev/null +++ b/tests/test_encoder_dispatch.py @@ -0,0 +1,158 @@ +import functools + +from collections.abc import Iterator +from typing import Any + +import pytest + +from tomlkit import api +from tomlkit.exceptions import ConvertError +from tomlkit.items import CUSTOM_ENCODERS +from tomlkit.items import Item + + +@pytest.fixture(autouse=True) +def restore_encoders() -> Iterator[None]: + original = CUSTOM_ENCODERS[:] + CUSTOM_ENCODERS.clear() + yield + CUSTOM_ENCODERS[:] = original + + +@pytest.mark.parametrize("parameter", ["_parent", "_sort_keys"]) +@pytest.mark.parametrize("required", [False, True]) +def test_individual_keyword_context(parameter: str, required: bool) -> None: + captured: dict[str, Any] = {} + namespace: dict[str, Any] = {"api": api, "captured": captured} + default = "" if required else "=None" + exec( + f"def encode(value, *, {parameter}{default}):\n" + f' captured["value"] = {parameter}\n' + f' return api.string("converted")\n', + namespace, + ) + parent = api.table() + api.register_encoder(namespace["encode"]) + assert ( + api.item(object(), _parent=parent, _sort_keys=True).as_string() == '"converted"' + ) + assert captured["value"] is (parent if parameter == "_parent" else True) + + +@pytest.mark.parametrize("name", ["_parent", "_sort_keys"]) +@pytest.mark.parametrize("syntax", ["", ", /"]) +def test_primary_value_name_not_reused_as_context(name: str, syntax: str) -> None: + namespace: dict[str, Any] = {"api": api} + exec( + f'def encode({name}{syntax}):\n return api.string("converted")\n', namespace + ) + api.register_encoder(namespace["encode"]) + assert ( + api.item(object(), _parent=api.table(), _sort_keys=True).as_string() + == '"converted"' + ) + + +def test_bound_method_with_context_subset() -> None: + class Encoder: + def __init__(self) -> None: + self.captured: bool | None = None + + def encode(self, value: Any, *, _sort_keys: bool = False) -> Item: + self.captured = _sort_keys + return api.string("converted") + + encoder = Encoder() + api.register_encoder(encoder.encode) + assert api.item(object(), _sort_keys=True).as_string() == '"converted"' + assert encoder.captured is True + + +def test_partial_encoder_with_context_subset() -> None: + captured: dict[str, Any] = {} + + def encode(prefix: str, value: Any, *, _parent: Item | None = None) -> Item: + captured["parent"] = _parent + return api.string(prefix) + + parent = api.table() + api.register_encoder(functools.partial(encode, "converted")) + assert api.item(object(), _parent=parent).as_string() == '"converted"' + assert captured["parent"] is parent + + +def test_kwargs_and_legacy_chain() -> None: + calls: list[Any] = [] + + def legacy(value: Any) -> Item: + calls.append("legacy") + raise ConvertError("not supported") + + def encode(value: Any, **kwargs: Any) -> Item: + calls.append(kwargs) + return api.string("converted") + + api.register_encoder(legacy) + api.register_encoder(encode) + parent = api.table() + assert ( + api.item(object(), _parent=parent, _sort_keys=True).as_string() == '"converted"' + ) + assert calls == ["legacy", {"_parent": parent, "_sort_keys": True}] + + +def test_internal_typeerror_propagates_once() -> None: + calls: list[Any] = [] + error = TypeError("encoder sentinel") + + def encode(value: Any, *, _parent: Item | None = None) -> Item: + calls.append(value) + raise error + + api.register_encoder(encode) + with pytest.raises(TypeError) as caught: + api.item(object()) + assert caught.value is error + assert len(calls) == 1 + + +def test_invalid_encoder_return_is_not_accepted() -> None: + def invalid(value: Any) -> Any: + return None + + api.register_encoder(invalid) + with pytest.raises(ConvertError, match="expected to return"): + api.item(object()) + + +@pytest.mark.parametrize("syntax", ["_parent, /", "*_parent"]) +def test_primary_nonkeyword_name_still_allows_context_in_kwargs(syntax: str) -> None: + captured: dict[str, Any] = {} + namespace: dict[str, Any] = {"api": api, "captured": captured} + exec( + f"def encode({syntax}, **kwargs):\n" + f" captured.update(kwargs)\n" + f' return api.string("converted")\n', + namespace, + ) + parent = api.table() + api.register_encoder(namespace["encode"]) + assert ( + api.item(object(), _parent=parent, _sort_keys=True).as_string() == '"converted"' + ) + assert captured == {"_parent": parent, "_sort_keys": True} + + +def test_bound_primary_keyword_name_does_not_receive_context_twice() -> None: + captured: dict[str, Any] = {} + + def encode(_parent: Any, **kwargs: Any) -> Item: + captured.update(kwargs) + return api.string("converted") + + api.register_encoder(encode) + assert ( + api.item(object(), _parent=api.table(), _sort_keys=True).as_string() + == '"converted"' + ) + assert captured == {"_sort_keys": True} diff --git a/tomlkit/items.py b/tomlkit/items.py index e84a40e4..72f37e92 100644 --- a/tomlkit/items.py +++ b/tomlkit/items.py @@ -223,14 +223,34 @@ def _sort_key(i: tuple[Any, Any]) -> Any: try: # Check if encoder accepts keyword arguments for backward compatibility sig = inspect.signature(encoder) - if "_parent" in sig.parameters or any( + has_kwargs = any( p.kind == p.VAR_KEYWORD for p in sig.parameters.values() - ): - # New style encoder that can accept additional parameters - rv = encoder(value, _parent=_parent, _sort_keys=_sort_keys) # type: ignore[call-arg] - else: - # Old style encoder that only accepts value - rv = encoder(value) + ) + first_parameter = next(iter(sig.parameters.values()), None) + bound_name = ( + first_parameter.name + if first_parameter is not None + and first_parameter.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD + else None + ) + context = {"_parent": _parent, "_sort_keys": _sort_keys} + kwargs = { + name: argument + for name, argument in context.items() + if name != bound_name + and ( + has_kwargs + or ( + name in sig.parameters + and sig.parameters[name].kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + ) + ) + } + rv = encoder(value, **kwargs) except ConvertError: pass else: From 1c42d44fd20151001a90b34935eb9cbfe9bd9e25 Mon Sep 17 00:00:00 2001 From: Sskift <163287349+Sskift@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:32:21 +0800 Subject: [PATCH 2/2] test: defer encoder annotations on Python 3.9 --- tests/test_encoder_dispatch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_encoder_dispatch.py b/tests/test_encoder_dispatch.py index f10ac17d..e906d865 100644 --- a/tests/test_encoder_dispatch.py +++ b/tests/test_encoder_dispatch.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import functools from collections.abc import Iterator