Open
Description
Describe the bug
The OAuth2Session class is not pickleable.
The problem is caused by the lambda function within the __init__ method.
In consequence, the class can't be used in a multiprocessing/threading context.
How to reproduce
Try to pickle the class.
And try to use the class in a multiprocessing/threading context
Expected behavior
The class should support pickling without any issue.
Additional context
Here's the exception:
concurrent.futures.process._RemoteTraceback:
"""
Traceback (most recent call last):
File "/usr/local/lib/python3.10/multiprocessing/queues.py", line 244, in _feed
obj = _ForkingPickler.dumps(obj)
File "/usr/local/lib/python3.10/multiprocessing/reduction.py", line 51, in dumps
cls(buf, protocol).dump(obj)
AttributeError: Can't pickle local object 'OAuth2Session.__init__.<locals>.<lambda>'
"""
Here's a wrapper that fixes the issue.
I tried 2 versions, and they all work (at least for my use case) and with python 3.10.
class OAuth2SessionPickleableV1(OAuth2Session):
"""A wrapper around OAuth2Session to make it pickleable."""
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
# This is the attribute that was not pickleable
# `self.auth = lambda r: r` was causing the issue
# since `lambda` is not pickleable
self.auth = self.auth_method
def auth_method(self, r: Any) -> Any:
return r
def __getstate__(self) -> Dict[str, Any]:
"""Customize pickling behavior"""
state = self.__dict__.copy()
return state
def __setstate__(self, state: Dict[str, Any]) -> None:
"""Customize pickling behavior"""
self.__dict__.update(state)
class OAuth2SessionPickleableV2(OAuth2Session):
"""A wrapper around OAuth2Session to make it pickleable."""
def __getstate__(self) -> Dict[str, Any]:
"""Customize pickling behavior"""
state = self.__dict__.copy()
state["auth"] = None # Remove the non-pickleable 'auth' attribute
return state
def __setstate__(self, state: Dict[str, Any]) -> None:
"""Customize pickling behavior"""
self.__dict__.update(state)
self.auth = lambda r: r # Reassign the 'auth' attribute after unpickling
Metadata
Metadata
Assignees
Labels
No labels