type check

.github/workflows/python.yaml

@@ -28,6 +28,9 @@ jobs:
       - 
         name: Lint
         run: poetry run ruff check .
+      -
+        name: Type check
+        run: poetry run mypy mlconfig
       - 
         name: Test
         run: poetry run pytest -v -s --cov=mlconfig --cov-report=xml tests

mlconfig/conf.py

@@ -1,14 +1,18 @@
 import copy
 import functools
+from typing import Any
 from typing import Optional
+from typing import Union
 
+from omegaconf import DictConfig
+from omegaconf import ListConfig
 from omegaconf import OmegaConf
 
 _registry = {}
 _key = "name"
 
 
-def load(f=None, obj=None) -> OmegaConf:
+def load(f: Optional[str] = None, obj: Optional[Any] = None) -> Union[ListConfig, DictConfig]:
     """Load configuration file or structured object.
     If both are provided, then the configuration from file
     will override the configuration from structured object.
@@ -37,7 +41,7 @@ def load(f=None, obj=None) -> OmegaConf:
     return OmegaConf.merge(*configs)
 
 
-def register(func_or_cls=None, name: str = None):
+def register(func_or_cls=None, name: Optional[str] = None):
     r"""Register function or class
 
     Args:

mlconfig/torch/__init__.py

@@ -3,7 +3,7 @@
 import mlconfig
 
 try:
-    from torch import optim
+    from torch import optim  # type: ignore
 except ImportError:
     print("Failed to import torch.")
 

pyproject.toml

@@ -14,6 +14,7 @@ pytest = "^7.4.3"
 pytest-cov = "^4.1.0"
 ruff = "^0.2.0"
 toml = "^0.10.2"
+mypy = "^1.10.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]