Skip to content

Commit 5b2eab6

Browse files
committed
Added support for pickling
1 parent 8bc0861 commit 5b2eab6

File tree

2 files changed

+23
-0
lines changed

2 files changed

+23
-0
lines changed

tap/tap.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -668,3 +668,16 @@ def __deepcopy__(self, memo: Dict[int, Any] = None) -> TapType:
668668
copied.__dict__[k] = deepcopy(v, memo)
669669

670670
return copied
671+
672+
def __getstate__(self) -> Dict[str, Any]:
673+
"""Gets the state of the object for pickling."""
674+
return self.as_dict()
675+
676+
def __setstate__(self, d: Dict[str, Any]) -> None:
677+
"""
678+
Initializes the object with the provided dictionary of arguments for unpickling.
679+
680+
:param d: A dictionary of arguments.
681+
"""
682+
self.__init__()
683+
self.from_dict(d)

tests/test_integration.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from copy import deepcopy
22
import os
33
from pathlib import Path
4+
import pickle
45
import sys
56
from tempfile import TemporaryDirectory
67
from typing import Any, Iterable, List, Optional, Set, Tuple
@@ -1373,6 +1374,15 @@ def configure(self) -> None:
13731374

13741375
self.assertEqual(args.a, 1)
13751376

1377+
1378+
class TestPickle(TestCase):
1379+
def test_pickle(self):
1380+
args = IntegrationDefaultTap().parse_args([])
1381+
pickle_str = pickle.dumps(args)
1382+
loaded_args: IntegrationDefaultTap = pickle.loads(pickle_str)
1383+
self.assertEqual(loaded_args.as_dict(), args.as_dict())
1384+
1385+
13761386
"""
13771387
- crash if default type not supported
13781388
- user specifying process_args

0 commit comments

Comments
 (0)