Skip to content

Commit d4b888b

Browse files
committed
Fix pickling of exception classes with kw_only attributes
BaseException.__reduce__ returns (cls, self.args, state), passing self.args as positional arguments to cls() on unpickle. When kw_only=True, __init__ only accepts keyword-only arguments, causing unpickling to fail with: TypeError: __init__() takes 1 positional argument but 2 were given This fix adds a custom __reduce__ for exception classes that have kw_only attrs. It uses a _rebuild_exc helper that creates the instance via __new__ (bypassing __init__) and restores attributes directly, then calls BaseException.__init__ to properly set self.args. Fixes GH#734
1 parent 25b74d6 commit d4b888b

File tree

2 files changed

+84
-0
lines changed

2 files changed

+84
-0
lines changed

src/attr/_make.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,31 @@ def __reduce__(self, _none_constructor=type(None), _args=()): # noqa: B008
103103
return _none_constructor, _args
104104

105105

106+
def _rebuild_exc(cls, state):
107+
"""
108+
Rebuild an exception instance without calling ``__init__``.
109+
110+
Used by ``__reduce__`` for exception classes with ``kw_only`` attributes,
111+
where ``BaseException.__reduce__`` would pass positional args that
112+
``kw_only`` rejects.
113+
"""
114+
obj = cls.__new__(cls)
115+
for name, value in state.items():
116+
try:
117+
object.__setattr__(obj, name, value)
118+
except AttributeError:
119+
pass
120+
# Restore BaseException.args which is set by __init__ via
121+
# BaseException.__init__(self, val1, val2, ...).
122+
init_values = tuple(
123+
state[a.name]
124+
for a in cls.__attrs_attrs__
125+
if a.init and a.name in state
126+
)
127+
BaseException.__init__(obj, *init_values)
128+
return obj
129+
130+
106131
def attrib(
107132
default=NOTHING,
108133
validator=None,
@@ -757,6 +782,26 @@ def __init__(
757782
self._cls_dict["__setstate__"],
758783
) = self._make_getstate_setstate()
759784

785+
# Fix pickling for exception classes with kw_only attributes.
786+
# BaseException.__reduce__ returns (cls, self.args, state) which calls
787+
# cls(*args) on unpickle, but kw_only attrs reject positional args.
788+
# Override __reduce__ to use __new__ + state instead of __init__.
789+
if props.is_exception and any(
790+
a.kw_only for a in attrs if a.init
791+
):
792+
_attr_names = self._attr_names
793+
794+
def __reduce__(self, *, _attr_names=_attr_names):
795+
state = {
796+
name: getattr(self, name)
797+
for name in _attr_names
798+
if name != "__weakref__"
799+
and hasattr(self, name)
800+
}
801+
return (_rebuild_exc, (self.__class__, state))
802+
803+
self._cls_dict["__reduce__"] = __reduce__
804+
760805
# tuples of script, globs, hook
761806
self._script_snippets: list[
762807
tuple[str, dict, Callable[[dict, dict], Any]]

tests/test_functional.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -624,6 +624,45 @@ class FooError(Exception):
624624

625625
FooError(1)
626626

627+
@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1))
628+
def test_auto_exc_kw_only_pickle(self, slots, frozen, protocol):
629+
"""
630+
Exception classes with kw_only attributes can be pickled.
631+
632+
BaseException.__reduce__ passes positional args, which kw_only
633+
rejects. Our custom __reduce__ handles this correctly.
634+
Regression test for GH#734.
635+
"""
636+
637+
@attr.s(auto_exc=True, slots=slots, frozen=frozen, kw_only=True)
638+
class KwOnlyError(Exception):
639+
msg = attr.ib()
640+
x = attr.ib()
641+
642+
e = KwOnlyError(msg="hello", x=42)
643+
e2 = pickle.loads(pickle.dumps(e, protocol))
644+
645+
assert e2.msg == "hello"
646+
assert e2.x == 42
647+
assert e2.args == e.args
648+
649+
def test_auto_exc_mixed_kw_only_pickle(self, slots):
650+
"""
651+
Exception classes with mixed positional and kw_only attributes
652+
can be pickled.
653+
"""
654+
655+
@attr.s(auto_exc=True, slots=slots, kw_only=False)
656+
class MixedError(Exception):
657+
pos = attr.ib()
658+
kw = attr.ib(kw_only=True)
659+
660+
e = MixedError(10, kw=20)
661+
e2 = pickle.loads(pickle.dumps(e))
662+
663+
assert e2.pos == 10
664+
assert e2.kw == 20
665+
627666
def test_eq_only(self, slots, frozen):
628667
"""
629668
Classes with order=False cannot be ordered.

0 commit comments

Comments
 (0)