OAuth2Session is not pickleable

[f-tremblay]

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

[jasoriya]

@JonathanHuot I've run into the same issue and its solvable by the suggested fix. Would you be wiling to accept a PR for the same?

[auvipy]

you can come with a PR but that need proper unit test