support type arguments for FlatDict&OrderedDict

Signed-off-by: Zhiyuan Chen <[email protected]>

chanfig/flat_dict.py

@@ -9,13 +9,12 @@
 from json import loads as json_loads
 from os import PathLike
 from os.path import splitext
-from typing import IO, Any, Callable, Iterable, Optional, Union
+from typing import IO, Any, Callable, Iterable, Optional, Union, TypeVar
 
-from yaml import SafeLoader
 from yaml import dump as yaml_dump
 from yaml import load as yaml_load
 
-from .utils import FileError, JsonEncoder, YamlDumper
+from .utils import FileError, JsonEncoder, YamlDumper, YamlLoader
 from .variable import Variable
 
 PathStr = Union[PathLike, str, bytes]
@@ -25,8 +24,11 @@
 JSON = ("json",)
 PYTHON = ("py",)
 
+K = TypeVar("K")
+V = TypeVar("V")
 
-class FlatDict(OrderedDict):
+
+class FlatDict(OrderedDict[K, V]):
 
     # pylint: disable=R0904
 
@@ -711,7 +713,7 @@ def from_yamls(cls, string: str, *args, **kwargs) -> FlatDict:
         """
 
         if "Loader" not in kwargs:
-            kwargs["Loader"] = SafeLoader
+            kwargs["Loader"] = YamlLoader
         return cls(**yaml_load(string, *args, **kwargs))
 
     def dump(self, file: File, method: Optional[str] = None, *args, **kwargs) -> None:  # pylint: disable=W1113

chanfig/nested_dict.py

@@ -3,13 +3,16 @@
 from collections.abc import Mapping
 from functools import wraps
 from os import PathLike
-from typing import Any, Callable, Iterable, Optional, Union
+from typing import Any, Callable, Iterable, Optional, Union, TypeVar
 
 from .flat_dict import FlatDict, PathStr
 from .variable import Variable
 
+K = TypeVar("K")
+V = TypeVar("V")
 
-class NestedDict(FlatDict):
+
+class NestedDict(FlatDict[K, V]):
     r"""
     NestedDict is basically an FlatDict object that create a nested structure with delimiter.
 

chanfig/utils.py

@@ -1,7 +1,7 @@
 from json import JSONEncoder
 from typing import Any
 
-from yaml import SafeDumper
+from yaml import SafeDumper, SafeLoader
 
 from .variable import Variable
 
@@ -28,5 +28,13 @@ def increase_indent(self, flow: bool = False, indentless: bool = False):  # pyli
         return super().increase_indent(flow, indentless)
 
 
-class FileError(ValueError):  # pylint: disable=C0115
-    pass
+class YamlLoader(SafeLoader):  # pylint: disable=R0901
+    """
+    YAML Loader for Config.
+    """
+
+
+class FileError(ValueError):
+    """
+    Error for file operations.
+    """