Support pickle.loads (#499) (#500)

mmcv/utils/config.py

@@ -372,6 +372,15 @@ class Config:
     def __iter__(self):
         return iter(self._cfg_dict)
 
+    def __getstate__(self):
+        return (self._cfg_dict, self._filename, self._text)
+
+    def __setstate__(self, state):
+        _cfg_dict, _filename, _text = state
+        super(Config, self).__setattr__('_cfg_dict', _cfg_dict)
+        super(Config, self).__setattr__('_filename', _filename)
+        super(Config, self).__setattr__('_text', _text)
+
     def dump(self, file=None):
         cfg_dict = super(Config, self).__getattribute__('_cfg_dict').to_dict()
         if self.filename.endswith('.py'):

tests/test_config.py

@@ -7,7 +7,7 @@ import tempfile
 import pytest
 import yaml
 
-from mmcv import Config, DictAction
+from mmcv import Config, DictAction, dump, load
 
 
 def test_construct():
@@ -354,3 +354,15 @@ def test_syntax_error():
             f'file {temp_cfg_path}'):
         Config.fromfile(temp_cfg_path)
     temp_cfg_file.close()
+
+
+def test_pickle_support():
+    cfg_file = osp.join(osp.dirname(__file__), 'data/config/n.py')
+    cfg = Config.fromfile(cfg_file)
+
+    with tempfile.TemporaryDirectory() as temp_config_dir:
+        pkl_cfg_filename = osp.join(temp_config_dir, '_pickle.pkl')
+        dump(cfg, pkl_cfg_filename)
+        pkl_cfg = load(pkl_cfg_filename)
+
+    assert pkl_cfg._cfg_dict == cfg._cfg_dict