Skip to content

Commit c9935a1

Browse files
committed
Add reduce special method to transform wrappers
Summary: Add the ___reduce__ method to transform wrappers (model/generator). This allows transform wrapped objects to be pickled (with original pickle module, not just dill) Test Plan: tests should pass
1 parent c4446fc commit c9935a1

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

aepsych/transforms/parameters.py

+8
Original file line numberDiff line numberDiff line change
@@ -477,6 +477,10 @@ def eval(self):
477477
f"{self._base_obj.__class__.__name__} has no attribute 'eval'"
478478
)
479479

480+
def __reduce__(self):
481+
# Helps pickle work (not dill)
482+
return (ParameterTransformedGenerator, (self._base_obj, self.transforms))
483+
480484
@classmethod
481485
def get_config_options(
482486
cls,
@@ -725,6 +729,10 @@ def eval(self):
725729
f"{self._base_obj.__class__.__name__} has no attribute 'eval'"
726730
)
727731

732+
def __reduce__(self):
733+
# Helps pickle work (not dill)
734+
return (ParameterTransformedModel, (self._base_obj, self.transforms))
735+
728736
@classmethod
729737
def get_config_options(
730738
cls,

0 commit comments

Comments
 (0)