Make DeferredImportModule / ModuleUnavailable objects picklable

Pyomo · May 23, 2022 · bedc0ef · bedc0ef
1 parent bb01dc0
commit bedc0ef
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/pyomo/common/dependencies.py b/pyomo/common/dependencies.py
@@ -74,6 +74,12 @@ def __getattr__(self, attr):
                                  % (type(self).__name__, attr))
         raise DeferredImportError(self._moduleunavailable_message())
 
+    def __getstate__(self):
+        return (self.__name__, self._moduleunavailable_info_)
+
+    def __setstate__(self, state):
+        self.__name__, self._moduleunavailable_info_ = state
+
     def _moduleunavailable_message(self, msg=None):
         _err, _ver, _imp, _package = self._moduleunavailable_info_
         if msg is None:
@@ -151,6 +157,12 @@ def __getattr__(self, attr):
                 _mod = getattr(_mod, _sub)
         return getattr(_mod, attr)
 
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state):
+        for k, v in state.items():
+            super().__setattr__(k, v)
 
 class _DeferredImportIndicatorBase(object):
     def __and__(self, other):