Fixes custom pickle protocol to handle None values as well (#1469)

* Fixes custom pickle protocol to handle `None` values as well

* Updates `CHANGELOG.md`

* Update test_pickle.py

* Update tests/test_primitives/test_container/test_base_container/test_pickle.py

Co-authored-by: Nikita Sobolev <mail@sobolevn.me>

Author: thepabloaguilar

CHANGELOG.md

@@ -11,6 +11,7 @@ See [0Ver](https://0ver.org/).
 ### Bugfixes
 
 - Fixes a problem with `do-notation` and type aliases
+- Fixes custom pickle protocol to handle `None` values gracefully
 
 
 ## 0.19.0 aka The Do Notation

returns/primitives/container.py

@@ -1,13 +1,28 @@
 from abc import ABCMeta
 from typing import Any, TypeVar
 
+from typing_extensions import TypedDict
+
 from returns.interfaces.equable import Equable
 from returns.primitives.hkt import Kind1
 from returns.primitives.types import Immutable
 
 _EqualType = TypeVar('_EqualType', bound=Equable)
 
 
+class _PickleState(TypedDict):
+    """Dict to represent the `BaseContainer` state to be pickled."""
+
+    # TODO: Remove `__slots__` from here when `slotscheck` allow ignore classes
+    # by using comments. We don't need the slots here since this class is just
+    # a representation of a dictionary and should not be instantiated by any
+    # means.
+    # See: https://github.com/ariebovenberg/slotscheck/issues/71
+    __slots__ = ('container_value',)  # type: ignore
+
+    container_value: Any
+
+
 class BaseContainer(Immutable, metaclass=ABCMeta):
     """Utility class to provide all needed magic methods to the context."""
 
@@ -37,13 +52,15 @@ def __hash__(self) -> int:
         """Used to use this value as a key."""
         return hash(self._inner_value)
 
-    def __getstate__(self) -> Any:
+    def __getstate__(self) -> _PickleState:
         """That's how this object will be pickled."""
-        return self._inner_value
+        return {'container_value': self._inner_value}  # type: ignore
 
-    def __setstate__(self, state: Any) -> None:
+    def __setstate__(self, state: _PickleState) -> None:
         """Loading state from pickled data."""
-        object.__setattr__(self, '_inner_value', state)  # noqa: WPS609
+        object.__setattr__(  # noqa: WPS609
+            self, '_inner_value', state['container_value'],
+        )
 
 
 def container_equality(

tests/test_io/test_io_container/test_io_pickle.py

@@ -3,11 +3,11 @@
 
 def test_io_pickle():
     """Tests how pickle protocol works for containers."""
-    assert IO(1).__getstate__() == 1  # noqa: WPS609
+    assert IO(1).__getstate__() == {'container_value': 1}  # noqa: WPS609
 
 
 def test_io_pickle_restore():
     """Ensures that object can be restored."""
     container = IO(2)
-    container.__setstate__(1)  # noqa: WPS609
+    container.__setstate__({'container_value': 1})  # type: ignore  # noqa: WPS609, E501
     assert container == IO(1)

tests/test_primitives/test_container/test_base_container/test_pickle.py

@@ -0,0 +1,36 @@
+import pickle  # noqa: S403
+from typing import Any
+
+from hypothesis import example, given
+from hypothesis import strategies as st
+
+from returns.primitives.container import BaseContainer
+
+
+class _CustomClass(object):
+    def __init__(self, inner_value: Any) -> None:
+        self.inner_value = inner_value
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            type(other) == type(self) and  # noqa: WPS516
+            self.inner_value == other.inner_value
+        )
+
+
+@given(
+    st.one_of(
+        st.integers(),
+        st.floats(allow_nan=False),
+        st.text(),
+        st.booleans(),
+        st.lists(st.text()),
+        st.dictionaries(st.text(), st.integers()),
+        st.builds(_CustomClass, st.text()),
+    ),
+)
+@example(None)
+def test_pickle(container_value: Any):
+    """Ensures custom pickle protocol works as expected."""
+    container = BaseContainer(container_value)
+    assert pickle.loads(pickle.dumps(container)) == container  # noqa: S301