update mixins

Author: leej3

ignite/base/mixins.py

@@ -1,11 +1,18 @@
 from collections import OrderedDict
 from collections.abc import Mapping
-from typing import Tuple
+from typing import List, Tuple
 
 
 class Serializable:
-    _state_dict_all_req_keys: Tuple = ()
-    _state_dict_one_of_opt_keys: Tuple = ()
+    _state_dict_all_req_keys: Tuple[str, ...] = ()
+    _state_dict_one_of_opt_keys: Tuple[Tuple[str, ...], ...] = ((),)
+
+    def __init__(self) -> None:
+        self._state_dict_user_keys: List[str] = []
+
+    @property
+    def state_dict_user_keys(self) -> List:
+        return self._state_dict_user_keys
 
     def state_dict(self) -> OrderedDict:
         raise NotImplementedError
@@ -19,6 +26,13 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 raise ValueError(
                     f"Required state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
                 )
-        opts = [k in state_dict for k in self._state_dict_one_of_opt_keys]
-        if len(opts) > 0 and ((not any(opts)) or (all(opts))):
-            raise ValueError(f"state_dict should contain only one of '{self._state_dict_one_of_opt_keys}' keys")
+        for one_of_opt_keys in self._state_dict_one_of_opt_keys:
+            opts = [k in state_dict for k in one_of_opt_keys]
+            if len(opts) > 0 and (not any(opts)) or (all(opts)):
+                raise ValueError(f"state_dict should contain only one of '{one_of_opt_keys}' keys")
+
+        for k in self._state_dict_user_keys:
+            if k not in state_dict:
+                raise ValueError(
+                    f"Required user state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
+                )